diff --git a/inspect_sft_file_prefix.py b/inspect_sft_file_prefix.py new file mode 100644 index 00000000000..7fc304a77c4 --- /dev/null +++ b/inspect_sft_file_prefix.py @@ -0,0 +1,332 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +#!/usr/bin/env python3 +"""Inspect pretokenized SFT samples from Megatron-LM .bin/.idx file pairs. + +Usage: + python inspect_sft_file_prefix.py + +Example: + python inspect_sft_file_prefix.py /path/to/dataset-materialized_text_document +""" + +import argparse +import struct +import numpy +from functools import lru_cache +from typing import Optional, Tuple + +# ────────────────────────────────────────────────────────────────────────────── +# Hardcoded config for nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8 chat template +# ────────────────────────────────────────────────────────────────────────────── +TOKENIZER_NAME = "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8" + +# These are pre-computed from: +# tokenizer.encode("<|im_start|>system\n", add_special_tokens=False) etc. +ROLE_START_TOKENS = { + "system": [10, 25708, 1010], # <|im_start|>system\n + "user": [10, 3263, 1010], # <|im_start|>user\n + "assistant": [10, 1503, 19464, 1010], # <|im_start|>assistant\n +} +END_TOKENS = [11, 1010] # <|im_end|>\n +THINK_START_ID = [12] # +THINK_END_ID = [13] # +TOOL_CALL_START = [14, 1010] # \n +TOOL_CALL_END = [15, 1010] # \n +TOOL_RESPONSE_START = [16, 1010] # \n +TOOL_RESPONSE_END = [17, 1010] # \n + +# ────────────────────────────────────────────────────────────────────────────── +# Index / bin readers (from Megatron-LM) +# ────────────────────────────────────────────────────────────────────────────── +_INDEX_HEADER = b"MMIDIDX\x00\x00" + + +class _MMapBinReader: + def __init__(self, bin_path: str) -> None: + self._bin_file_reader = open(bin_path, mode="rb") + self._bin_buffer_mmap = numpy.memmap(self._bin_file_reader, mode="r", order="C") + self._bin_buffer = memoryview(self._bin_buffer_mmap.data) + + def read(self, dtype, count: int, offset: int) -> numpy.ndarray: + return numpy.frombuffer(self._bin_buffer, dtype=dtype, count=count, offset=offset) + + def __del__(self) -> None: + if self._bin_buffer_mmap is not None: + self._bin_buffer_mmap._mmap.close() + if self._bin_file_reader is not None: + self._bin_file_reader.close() + del self._bin_buffer_mmap + del self._bin_file_reader + + +class _IndexReader: + def __init__(self, idx_path: str) -> None: + with open(idx_path, "rb") as stream: + header = stream.read(9) + assert header == _INDEX_HEADER, f"bad header, cannot read: {idx_path}" + + version = struct.unpack(" int: + return self.sequence_count + + @lru_cache(maxsize=8) + def __getitem__(self, idx: int) -> Tuple[numpy.int64, numpy.int32]: + return (self.sequence_pointers[idx], self.sequence_lengths[idx]) + + def __del__(self) -> None: + if hasattr(self, "bin_buffer_mmap"): + self.bin_buffer_mmap._mmap.close() + del self.bin_buffer_mmap + + +# ────────────────────────────────────────────────────────────────────────────── +# Segment extraction logic +# ────────────────────────────────────────────────────────────────────────────── +def find_subsequence(sequence, subsequence, start=0): + sub_len = len(subsequence) + for i in range(start, len(sequence) - sub_len + 1): + if sequence[i : i + sub_len] == subsequence: + return i + return -1 + + +NL_TOKEN = 1010 # \n token id + +def split_tool_calls(tokens, offset): + """Split a token sequence into assistant text and tool_call sub-segments. + + Whitespace-only assistant fragments (e.g. a lone \\n between and + ) are dropped so we don't produce meaningless segments. + """ + tc_start_len = len(TOOL_CALL_START) + tc_end_len = len(TOOL_CALL_END) + result = [] + pos = 0 + while pos < len(tokens): + # Find next \n + tc_start = find_subsequence(tokens, TOOL_CALL_START, pos) + + if tc_start == -1: + # No more tool calls, rest is regular assistant content + if pos < len(tokens): + result.append({"role": "assistant", "tokens": tokens[pos:], "start": offset + pos, "end": offset + len(tokens)}) + break + + # Assistant content before tool_call (skip if whitespace-only) + if tc_start > pos: + frag = tokens[pos:tc_start] + if not all(t == NL_TOKEN for t in frag): + result.append({"role": "assistant", "tokens": frag, "start": offset + pos, "end": offset + tc_start}) + + # Find matching \n + content_start = tc_start + tc_start_len + tc_end = find_subsequence(tokens, TOOL_CALL_END, content_start) + + if tc_end == -1: + # No closing tag, treat rest as tool_call + result.append({"role": "tool_call", "tokens": tokens[content_start:], "start": offset + content_start, "end": offset + len(tokens)}) + break + + # Tool call content (excluding markers) + result.append({"role": "tool_call", "tokens": tokens[content_start:tc_end], "start": offset + content_start, "end": offset + tc_end}) + pos = tc_end + tc_end_len + + # Also drop trailing whitespace-only assistant fragments + if result and result[-1]["role"] == "assistant" and all(t == NL_TOKEN for t in result[-1]["tokens"]): + result.pop() + + return result + + +def extract_segments(tokenized_conversation, role_start_tokens, end_tokens, think_start_id, think_end_id): + markers = [] + for role, start_tokens in role_start_tokens.items(): + pos = 0 + while True: + idx = find_subsequence(tokenized_conversation, start_tokens, pos) + if idx == -1: + break + markers.append((idx, role, len(start_tokens))) + pos = idx + len(start_tokens) + markers.sort(key=lambda x: x[0]) + + segments = [] + for start_pos, role, marker_len in markers: + content_start = start_pos + marker_len + end_pos = find_subsequence(tokenized_conversation, end_tokens, content_start) + if end_pos == -1: + content_end = len(tokenized_conversation) + else: + content_end = end_pos + content_tokens = tokenized_conversation[content_start:content_end] + + # Check if this user turn is actually a tool response + if role == "user" and len(content_tokens) >= len(TOOL_RESPONSE_START) and content_tokens[:len(TOOL_RESPONSE_START)] == TOOL_RESPONSE_START: + segments.append({"role": "tool_response", "tokens": content_tokens, "start": content_start, "end": content_end}) + continue + + if role == "assistant": + think_start_idx = find_subsequence(content_tokens, think_start_id) + if think_start_idx != -1: + think_end_idx = find_subsequence(content_tokens, think_end_id, think_start_idx + len(think_start_id)) + else: + think_end_idx = -1 + + if think_start_idx != -1 and think_end_idx != -1: + reasoning_tokens = content_tokens[think_start_idx + len(think_start_id) : think_end_idx] + response_tokens = content_tokens[think_end_idx + len(think_end_id) :] + if reasoning_tokens: + abs_start = content_start + think_start_idx + len(think_start_id) + abs_end = content_start + think_end_idx + segments.append({"role": "reasoning", "tokens": reasoning_tokens, "start": abs_start, "end": abs_end}) + # Split the response part by tool calls + if response_tokens: + abs_start = content_start + think_end_idx + len(think_end_id) + segments.extend(split_tool_calls(response_tokens, abs_start)) + continue + + # No think tags — split entire content by tool calls + if find_subsequence(content_tokens, TOOL_CALL_START) != -1: + segments.extend(split_tool_calls(content_tokens, content_start)) + continue + + segments.append({"role": role, "tokens": content_tokens, "start": content_start, "end": content_end}) + + return segments + + +# ────────────────────────────────────────────────────────────────────────────── +# Main +# ────────────────────────────────────────────────────────────────────────────── +def main(): + parser = argparse.ArgumentParser(description="Inspect pretokenized SFT samples.") + parser.add_argument("file_prefix", help="Path prefix for .bin/.idx files (without extension)") + args = parser.parse_args() + + print(f"Loading tokenizer: {TOKENIZER_NAME}") + from transformers import AutoTokenizer + tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_NAME, trust_remote_code=True) + + # Verify hardcoded token ids match this tokenizer + assert tokenizer.encode("<|im_start|>system\n", add_special_tokens=False) == ROLE_START_TOKENS["system"] + assert tokenizer.encode("<|im_start|>user\n", add_special_tokens=False) == ROLE_START_TOKENS["user"] + assert tokenizer.encode("<|im_start|>assistant\n", add_special_tokens=False) == ROLE_START_TOKENS["assistant"] + assert tokenizer.encode("<|im_end|>\n", add_special_tokens=False) == END_TOKENS + assert tokenizer.encode("", add_special_tokens=False) == THINK_START_ID + assert tokenizer.encode("", add_special_tokens=False) == THINK_END_ID + assert tokenizer.encode("\n", add_special_tokens=False) == TOOL_CALL_START + assert tokenizer.encode("\n", add_special_tokens=False) == TOOL_CALL_END + assert tokenizer.encode("\n", add_special_tokens=False) == TOOL_RESPONSE_START + assert tokenizer.encode("\n", add_special_tokens=False) == TOOL_RESPONSE_END + + print(f"Loading index: {args.file_prefix}.idx") + index = _IndexReader(args.file_prefix + ".idx") + print(f"Loading bin: {args.file_prefix}.bin") + reader = _MMapBinReader(args.file_prefix + ".bin") + print(f"Total sequences: {len(index)}\n") + + while True: + try: + raw = input(f"Enter sample index [0-{len(index) - 1}] (q to quit): ").strip() + except (EOFError, KeyboardInterrupt): + print() + break + + if raw.lower() == "q": + break + + # Parse optional -d suffix for raw detokenized output + detokenize_only = raw.endswith("-d") + if detokenize_only: + raw = raw[:-2].strip() + + try: + sample_idx = int(raw) + except ValueError: + print(f"Invalid input: {raw!r}") + continue + + if sample_idx < 0 or sample_idx >= len(index): + print(f"Out of range. Must be 0-{len(index) - 1}") + continue + + pointer, length = index[sample_idx] + sample = reader.read(numpy.int32, int(length), int(pointer)) + + if detokenize_only: + print(f"\n{'=' * 80}") + print(f"Sample {sample_idx} | {len(sample)} total tokens (raw detokenized)") + print(f"{'=' * 80}\n") + print(tokenizer.decode(sample.tolist())) + print(f"\n{'=' * 80}\n") + continue + + segments = extract_segments( + sample.tolist(), ROLE_START_TOKENS, END_TOKENS, THINK_START_ID, THINK_END_ID + ) + + print(f"\n{'=' * 80}") + print(f"Sample {sample_idx} | {len(sample)} total tokens") + print(f"{'=' * 80}") + + assistant_tokens = 0 + reasoning_tokens = 0 + tool_call_tokens = 0 + + for seg in segments: + decoded = tokenizer.decode(seg["tokens"]) + role_label = seg["role"].upper() + n_tokens = len(seg["tokens"]) + print(f"\n[{role_label:>13}] ({n_tokens:>5} tokens | {seg['start']}:{seg['end']})") + print(decoded) + + if seg["role"] == "assistant": + assistant_tokens += n_tokens + elif seg["role"] == "reasoning": + reasoning_tokens += n_tokens + elif seg["role"] == "tool_call": + tool_call_tokens += n_tokens + + has_tools = tool_call_tokens > 0 + has_reasoning = reasoning_tokens > 0 + + print(f"\n{'-' * 80}") + print(f"Training tokens (assistant only): {assistant_tokens}") + if has_tools: + print(f"Training tokens (assistant + tool calls): {assistant_tokens + tool_call_tokens}") + if has_reasoning: + print(f"Training tokens (assistant + reasoning): {assistant_tokens + reasoning_tokens}") + if has_tools or has_reasoning: + print(f"Training tokens (assistant + tool calls + reasoning): {assistant_tokens + tool_call_tokens + reasoning_tokens}") + print(f"{'=' * 80}\n") + + +if __name__ == "__main__": + main() diff --git a/megatron/core/datasets/sft.MD b/megatron/core/datasets/sft.MD new file mode 100644 index 00000000000..53b883ce5f8 --- /dev/null +++ b/megatron/core/datasets/sft.MD @@ -0,0 +1,291 @@ +# SFT Dataset + +The `SFTDataset` class (`sft_dataset.py`) provides supervised fine-tuning over +pretokenized chat conversations stored in Megatron's `IndexedDataset` format +(`.bin` / `.idx` pairs). + +--- + +## 1. `__getitem__` outputs + +| Key | Shape | Dtype | Description | +|-----|-------|-------|-------------| +| `tokens` | `(S-1,)` | `int64` | Input token ids (all but the **last** token of the packed sample) | +| `labels` | `(S-1,)` | `int64` | Target token ids (all but the **first** token — i.e. `tokens` shifted right by one) | +| `loss_mask` | `(S-1,)` | `float32` | Per-token loss mask (`0.0`/`1.0`). Which tokens are trained on depends on config (see §3) | +| `cu_seqlens` | `(D+1,)` | `int32` | Cumulative sequence lengths. `D` is the number of documents packed into this sample. First element is always `0`, last element equals `S` (total tokens). Used for Flash Attention varlen kernels and to reset position ids per document | + +`S` is the sum of the token lengths of every document packed into the sample +(always `<= sequence_length`). + +--- + +## 2. Index-building logic + +Three numpy arrays are built (or loaded from cache) at construction time: + +### 2.1 Document index — `document_index` + +A **1-D** `int32` array that maps flat positions to document ids in the +underlying `IndexedDataset`. + +``` +document_index = [docA, docB, docC, docD, docE, ...] +``` + +Built by flattening the output of `pack_samples`. + +### 2.2 Sample index — `sample_index` + +A **1-D** `int32` array of cumulative bin sizes. Its length is +`num_packed_samples + 1`. The documents belonging to packed sample `i` are: + +``` +document_index[ sample_index[i] : sample_index[i+1] ] +``` + +**Example** — 6 documents packed into 3 bins `[[0,2], [1,4,5], [3]]`: + +``` +document_index = [0, 2, 1, 4, 5, 3] +sample_index = [0, 2, 5, 6] + │ │ │ └─ sample 2: document_index[5:6] → [3] + │ │ └──── sample 1: document_index[2:5] → [1, 4, 5] + │ └─────── sample 0: document_index[0:2] → [0, 2] + └────────── always starts at 0 +``` + +### 2.3 Shuffle index — `shuffle_index` + +A **1-D** `int32` array that is a random permutation of +`range(num_packed_samples)`, extended across multiple epochs when `num_samples` +exceeds the number of packed samples. + +```python +for epoch in range(num_epochs): + shuffle_index.extend(torch.randperm(num_packed)) +``` + +### 2.4 Sample packing — `pack_samples` (MFFD) + +Variable-length documents are packed into fixed-capacity bins using the +**Modified First-Fit Decreasing** algorithm. Documents longer than +`sequence_length` are dropped with a warning. + +| Phase | Action | +|-------|--------| +| 0 | Classify items: **large** `(C/2, C]`, **medium** `(C/3, C/2]`, **small** `(C/6, C/3]`, **tiny** `(0, C/6]` | +| 1 | Open one bin per large item | +| 2 | Forward pass: add one medium item to each large bin if it fits | +| 3 | Backward pass: fill large-only bins with two small items | +| 4 | Greedy forward: fit remaining items into existing bins | +| 5 | FFD on leftovers: create new bins for anything still unplaced | + +### 2.5 Caching + +Indices are saved as `.npy` files under +`/cache/SFTDataset_indices/`. On subsequent runs, they are +memory-mapped (`mmap_mode="r"`) for fast startup. + +### 2.6 Query path (`__getitem__` → indices → tokens) + +``` +idx + │ + ▼ +shuffle_index[idx] → packed_sample_id + │ + ▼ +sample_index[packed_id] → (doc_index_beg, doc_index_end) +sample_index[packed_id + 1] + │ + ▼ +for i in range(beg, end): + document_index[i] → doc_id + dataset.get(doc_id) → token array for that document + │ + ▼ +concatenate all token arrays → tokens (length S) +concatenate all loss masks → loss_mask (length S) +per-document lengths → lengths (used to build cu_seqlens) +``` + +--- + +## 3. Loss masking — which tokens do we train on? + +### 3.1 Conversation structure + +Every pretokenized sample is a complete conversation. There are two families: + +**Chat-only** (no tool use): + +``` +<|im_start|>system\n ... <|im_end|>\n +<|im_start|>user\n ... <|im_end|>\n +<|im_start|>assistant\n ... response <|im_end|>\n + (possibly more user / assistant turns) +``` + +**Tool-use** (agentic trajectories): + +``` +<|im_start|>system\n ... (tool definitions) ... <|im_end|>\n +<|im_start|>user\n ... (task) ... <|im_end|>\n +<|im_start|>assistant\n ... text \n ... \n <|im_end|>\n +<|im_start|>user\n \n ... \n <|im_end|>\n + (repeat tool-call / tool-response turns) +<|im_start|>assistant\n ... final answer <|im_end|>\n +``` + +Key structural observations: + +- Every assistant turn starts with ``. The reasoning trace may be empty + (``) but the tags are always present. +- Tool responses are wrapped inside **user turns** + (`<|im_start|>user\n ...\n \n <|im_end|>\n`). + They are never a separate role. +- Tool calls live inside **assistant turns**, enclosed by + `\n ... \n`. +- In tool-use samples, the system prompt contains `` + tokens as part of the format description. These are **not** actual tool calls. + +### 3.2 Segment roles + +`extract_segments` parses each tokenized conversation and classifies every span +of content tokens into one of these roles: + +| Role | Where it appears | Content | +|------|-----------------|---------| +| `system` | System turn content | System prompt / tool definitions | +| `user` | User turn content | User messages | +| `tool_response` | User turn starting with `` | Environment output | +| `reasoning` | Between `` and `` inside an assistant turn | Chain-of-thought trace | +| `assistant` | After `` inside an assistant turn (excluding tool call regions) | Model response text | +| `tool_call` | Between `\n` and `\n` inside an assistant turn | Tool invocations | + +### 3.3 What gets `loss_mask = 1` + +The loss mask determines which token positions contribute to the training loss. +Three config flags control the behavior (`train_on_tool_calls` and +`train_on_thinking_traces` require `train_on_assistant_responses_only = True`): + +| `train_on_assistant_responses_only` | `train_on_tool_calls` | `train_on_thinking_traces` | `loss_mask = 1` on | +|---|---|---|---| +| `False` | — | — | **All** tokens | +| `True` | `False` | `False` | `assistant` only | +| `True` | `True` | `False` | `assistant` + `tool_call` | +| `True` | `False` | `True` | `assistant` + `reasoning` | +| `True` | `True` | `True` | `assistant` + `tool_call` + `reasoning` | + +`tool_response` segments (environment output) are **never** trained on — they +are always masked regardless of configuration. + +### 3.4 Worked example — chat-only + +A 2-turn conversation with reasoning traces. 1 186 tokens total. + +``` + POSITION TOKEN(S) ROLE loss_mask + ──────── ────────────────────────── ─────────── ───────── + [0:3] <|im_start|>system\n (template) 0 + [3:3] (empty system content) system 0 + [3:5] <|im_end|>\n (template) 0 + [5:8] <|im_start|>user\n (template) 0 + [8:40] "Ugh, fill this: ..." user 0 + [40:42] <|im_end|>\n (template) 0 + [42:46] <|im_start|>assistant\n (template) 0 + [46] (template) 0 + [47] (template) 0 + [47:173] "The worst thing ..." assistant ← 1 + [173:175] <|im_end|>\n (template) 0 + [175:178] <|im_start|>user\n (template) 0 + [178:241] "Okay, my Alexa just ..." user 0 + [241:243] <|im_end|>\n (template) 0 + [243:247] <|im_start|>assistant\n (template) 0 + [247] (template) 0 + [248:965] (717 tokens of reasoning) reasoning ← depends on config + [965] (template) 0 + [966:1184] "The real horror ..." assistant ← 1 + [1184:1186] <|im_end|>\n (template) 0 +``` + +| Setting | Training tokens | +|---------|-----------------| +| assistant only | **343** | +| assistant + reasoning | **1 060** | + +### 3.5 Worked example — tool-use + +An agentic coding trajectory. 17 955 tokens total, 39 assistant turns, 38 tool +responses, ~39 tool calls. + +``` + POSITION TOKEN(S) ROLE loss_mask + ────────── ─────────────────────────────── ───────────── ───────── + [0:3] <|im_start|>system\n (template) 0 + [3:1107] system prompt + tool defs system 0 + (includes + as format examples — NOT real) + [1107:1109] <|im_end|>\n (template) 0 + [1109:1112] <|im_start|>user\n (template) 0 + [1112:2865] task description user 0 + [2865:2867] <|im_end|>\n (template) 0 + + ┌─── assistant turn (with tool call) ──────────────────────────────── + │ [2867:2871] <|im_start|>assistant\n (template) 0 + │ [2871] (template) 0 + │ [2872] (template) 0 + │ [2873:3631] "I'll help you implement ..." assistant ← 1 + │ [3631:3633] \n (template) 0 + │ [3633:3665] tool_call ← depends on config + │ view ... + │ [3665:3667] \n (template) 0 + │ [3667:3669] <|im_end|>\n (template) 0 + └──────────────────────────────────────────────────────────────────── + + ┌─── tool response (inside a user turn) ───────────────────────────── + │ [3669:3672] <|im_start|>user\n (template) 0 + │ [3672:3674] \n (template) 0 + │ [3674:4608] OBSERVATION: ... tool_response 0 + │ [4608:4610] \n (template) 0 + │ [4610:4612] <|im_end|>\n (template) 0 + └──────────────────────────────────────────────────────────────────── + + ... (pattern repeats for each tool-call / tool-response pair) ... + + ┌─── final assistant turn (text, no tool call) ────────────────────── + │ [17855:17859] <|im_start|>assistant\n (template) 0 + │ [17859] (template) 0 + │ [17860] (template) 0 + │ [17861:17874] "Let me remove the ..." assistant ← 1 + │ [17874:17876] \n (template) 0 + │ [17876:17910] rm ... tool_call ← depends on config + │ [17910:17912] \n (template) 0 + │ [17912:17914] <|im_end|>\n (template) 0 + └──────────────────────────────────────────────────────────────────── +``` + +| Setting | Training tokens | +|---------|-----------------| +| assistant only | **2 790** | +| assistant + tool calls | **5 732** | + +### 3.6 Template tokens and the loss mask + +The following structural tokens always have `loss_mask = 0`: + +| Token(s) | Id(s) | Purpose | +|----------|-------|---------| +| `<\|im_start\|>` | `10` | Turn boundary | +| `system\n` / `user\n` / `assistant\n` | varies | Role identifier | +| `<\|im_end\|>\n` | `11, 1010` | Turn terminator | +| `` | `12` | Reasoning open | +| `` | `13` | Reasoning close | +| `\n` | `14, 1010` | Tool call open | +| `\n` | `15, 1010` | Tool call close | +| `\n` | `16, 1010` | Tool response open | +| `\n` | `17, 1010` | Tool response close | + +These tokens delineate message boundaries and roles. They are part of the chat +template, not model-generated content, so they are always masked out. diff --git a/megatron/core/datasets/sft_dataset.py b/megatron/core/datasets/sft_dataset.py new file mode 100644 index 00000000000..cfa1d7bd6ca --- /dev/null +++ b/megatron/core/datasets/sft_dataset.py @@ -0,0 +1,826 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +import logging +import os +import time +from bisect import bisect +from dataclasses import dataclass, field +from typing import Dict, List, Optional, Tuple + +import numpy +import torch + +from megatron.core.datasets.gpt_dataset import GPTDatasetConfig +from megatron.core.datasets.indexed_dataset import IndexedDataset +from megatron.core.datasets.megatron_dataset import MegatronDataset +from megatron.core.datasets.object_storage_utils import ObjectStorageConfig, is_object_storage_path +from megatron.core.datasets.utils import Split +from megatron.core.utils import log_single_rank + +logger = logging.getLogger(__name__) + +IGNORE_INDEX = -100 + + +@dataclass +class ChatTemplateConfig: + """Configuration for chat template delimiter strings used to parse tokenized conversations. + + Each field holds the string form of a delimiter token sequence. During SFT dataset + construction these strings are tokenized once and then used to locate role boundaries, + thinking traces, tool calls, and tool responses inside a tokenized conversation so that + a per-token loss mask can be built. + """ + + system_start_str: str + user_start_str: str + assistant_start_str: str + end_str: str + think_start_str: str + think_end_str: str + tool_call_start_str: str + tool_call_end_str: str + tool_response_start_str: str + tool_response_end_str: str + + +@dataclass +class Nemotron3ChatTemplateConfig(ChatTemplateConfig): + """Default ChatTemplateConfig pre-filled with Nemotron-3 delimiters.""" + + system_start_str: str = "<|im_start|>system\n" + user_start_str: str = "<|im_start|>user\n" + assistant_start_str: str = "<|im_start|>assistant\n" + end_str: str = "<|im_end|>\n" + think_start_str: str = "" + think_end_str: str = "" + tool_call_start_str: str = "\n" + tool_call_end_str: str = "\n" + tool_response_start_str: str = "\n" + tool_response_end_str: str = "\n" + + +@dataclass +class SFTDatasetConfig(GPTDatasetConfig): + """Configuration for SFT (Supervised Fine-Tuning) datasets. + + Extends GPTDatasetConfig with options that control which parts of a + multi-turn conversation contribute to the training loss. When + ``train_on_assistant_responses_only`` is True, the chat template delimiter + strings are tokenized in ``__post_init__`` and stored as token-id lists + so that ``extract_segments`` can build per-token loss masks at training time. + """ + + train_on_assistant_responses_only: bool = True + """If True, only train on completions, otherwise train on full conversations""" + + train_on_thinking_traces: bool = True + """If False, mask thinking traces from the completions""" + + train_on_tool_calls: bool = True + """If True, include tool call content in the loss computation""" + + chat_template_config: ChatTemplateConfig = field(default_factory=Nemotron3ChatTemplateConfig) + + role_start_tokens: Dict[str, List[int]] = field(init=False, default=None) + end_tokens: List[int] = field(init=False, default=None) + think_start_tokens: List[int] = field(init=False, default=None) + think_end_tokens: List[int] = field(init=False, default=None) + tool_call_start_tokens: List[int] = field(init=False, default=None) + tool_call_end_tokens: List[int] = field(init=False, default=None) + tool_response_start_tokens: List[int] = field(init=False, default=None) + tool_response_end_tokens: List[int] = field(init=False, default=None) + + def __post_init__(self) -> None: + """Do asserts and set fields post init""" + super().__post_init__() + + assert ( + not self.train_on_thinking_traces or self.train_on_assistant_responses_only + ), "train_on_assistant_responses_only must be True when train_on_thinking_traces is True" + + assert ( + not self.train_on_tool_calls or self.train_on_assistant_responses_only + ), "train_on_assistant_responses_only must be True when train_on_tool_calls is True" + + assert ( + not self.train_on_assistant_responses_only or self.chat_template_config is not None + ), "chat_template_config must be provided when train_on_assistant_responses_only is True" + + if self.train_on_assistant_responses_only: + self.role_start_tokens = { + "system": self.tokenizer.tokenize( + self.chat_template_config.system_start_str, add_special_tokens=False + ), + "user": self.tokenizer.tokenize( + self.chat_template_config.user_start_str, add_special_tokens=False + ), + "assistant": self.tokenizer.tokenize( + self.chat_template_config.assistant_start_str, add_special_tokens=False + ), + } + self.end_tokens = self.tokenizer.tokenize( + self.chat_template_config.end_str, add_special_tokens=False + ) + self.think_start_tokens = self.tokenizer.tokenize( + self.chat_template_config.think_start_str, add_special_tokens=False + ) + self.think_end_tokens = self.tokenizer.tokenize( + self.chat_template_config.think_end_str, add_special_tokens=False + ) + self.tool_call_start_tokens = self.tokenizer.tokenize( + self.chat_template_config.tool_call_start_str, add_special_tokens=False + ) + self.tool_call_end_tokens = self.tokenizer.tokenize( + self.chat_template_config.tool_call_end_str, add_special_tokens=False + ) + self.tool_response_start_tokens = self.tokenizer.tokenize( + self.chat_template_config.tool_response_start_str, add_special_tokens=False + ) + self.tool_response_end_tokens = self.tokenizer.tokenize( + self.chat_template_config.tool_response_end_str, add_special_tokens=False + ) + + +class SFTDataset(MegatronDataset): + """A dataset for Supervised Fine-Tuning on multi-turn conversations. + + Reads pre-tokenized conversations from an IndexedDataset, packs multiple + conversations into fixed-length samples using Modified First-Fit Decreasing + bin-packing, and builds per-token loss masks that optionally restrict the + training signal to assistant responses, thinking traces, and/or tool calls. + + Each ``__getitem__`` returns a dict with keys ``tokens``, ``labels``, + ``loss_mask``, and ``cu_seqlens`` (cumulative sequence lengths for + variable-length packing). + """ + + def __init__( + self, + indexed_dataset: IndexedDataset, + dataset_path: Optional[str], + indexed_indices: numpy.ndarray, + num_samples: Optional[int], + index_split: Split, + config: SFTDatasetConfig, + ) -> None: + super().__init__( + indexed_dataset, dataset_path, indexed_indices, num_samples, index_split, config + ) + + (self.document_index, self.sample_index, self.shuffle_index) = ( + self._build_document_sample_shuffle_indices() + ) + + @staticmethod + def numel_low_level_dataset(low_level_dataset: IndexedDataset) -> int: + """Abstract method implementation + + For GPT, the underlying IndexedDataset should be split by sequence, as opposed to, say, + BERT, which should be split by document + + Args: + low_level_dataset (IndexedDataset): The underlying IndexedDataset + + Returns: + int: The number of unique elements in the underlying IndexedDataset + """ + return low_level_dataset.sequence_lengths.shape[0] + + @staticmethod + def build_low_level_dataset(dataset_path: str, config: SFTDatasetConfig) -> IndexedDataset: + """Abstract method implementation + + Args: + dataset_path (str): The real path prefix to the IndexedDataset .bin and .idx files + + config (SFTDatasetConfig): The config + + Returns: + IndexedDataset: The underlying IndexedDataset + """ + if is_object_storage_path(dataset_path): + assert config.object_storage_cache_path is not None + return IndexedDataset( + dataset_path, + multimodal=False, + mmap=config.mmap_bin_files, + object_storage_config=ObjectStorageConfig( + path_to_idx_cache=config.object_storage_cache_path + ), + ) + return IndexedDataset(dataset_path, multimodal=False, mmap=config.mmap_bin_files) + + def __len__(self) -> int: + """Abstract method implementation + + Returns: + int: The length of the dataset + """ + + return self.shuffle_index.shape[0] + + def __getitem__(self, idx: Optional[int]) -> Dict[str, torch.Tensor]: + """Abstract method implementation + + Args: + idx (Optional[int]): The index into the dataset + + Returns: + Dict[str, torch.Tensor]: The sample information wrapped in a dictionary + """ + tokens, loss_mask, lengths = self._query_document_sample_shuffle_indices(idx) + + tokens = torch.from_numpy(tokens).long() + loss_mask = torch.from_numpy(loss_mask).float() + cu_seqlens = torch.cat( + (torch.zeros(1, dtype=torch.int32), torch.cumsum(torch.from_numpy(lengths), dim=0)) + ).to( + torch.int32 + ) # NOTE(asolergi-nv): torch.cumsum promotes int32 to int64 + + return { + 'tokens': tokens[:-1].contiguous(), + 'labels': tokens[1:].contiguous(), + 'loss_mask': loss_mask[:-1].contiguous(), + 'cu_seqlens': cu_seqlens, + } + + def _query_document_sample_shuffle_indices( + self, idx: int + ) -> Tuple[numpy.ndarray, numpy.ndarray, numpy.ndarray]: + """Get the text (token ids), loss mask, and cu_seqlens for a given index + + Args: + idx (int): The index into the dataset + + Returns: + Tuple[numpy.ndarray, numpy.ndarray, numpy.ndarray]: The text ids, + loss mask, and cu_seqlens + """ + # Do the shuffle mapping + idx = self.shuffle_index[idx] + + # Get the beginning and end documents and offsets + doc_index_beg = self.sample_index[idx] + doc_index_end = self.sample_index[idx + 1] + + document_ids = [] + sample_parts = [] + loss_masks = [] + + # TODO(asolergi-nv): For PP add flag to JUST read cu_seqlens + + for i in range(doc_index_beg, doc_index_end): + sample = self.dataset.get(self.indices[self.document_index[i]]) + if self.config.train_on_assistant_responses_only: + segments = extract_segments( + sample.tolist(), + self.config.role_start_tokens, + self.config.end_tokens, + self.config.think_start_tokens, + self.config.think_end_tokens, + self.config.tool_call_start_tokens, + self.config.tool_call_end_tokens, + self.config.tool_response_start_tokens, + ) + loss_mask = numpy.zeros(sample.shape[0], dtype=numpy.int64) + for seg in segments: + if seg["role"] == "assistant": + loss_mask[seg["start"] : seg["end"]] = 1 + elif seg["role"] == "reasoning" and self.config.train_on_thinking_traces: + loss_mask[seg["start"] : seg["end"]] = 1 + elif seg["role"] == "tool_call" and self.config.train_on_tool_calls: + loss_mask[seg["start"] : seg["end"]] = 1 + else: + loss_mask = numpy.ones(sample.shape[0], dtype=numpy.int64) + + # Add the sample part & loss mask + sample_parts.append(sample) + loss_masks.append(loss_mask) + + return ( + numpy.concatenate(sample_parts, dtype=numpy.int64), + numpy.concatenate(loss_masks, dtype=numpy.int64), + numpy.array([len(sample_part) for sample_part in sample_parts], dtype=numpy.int32), + ) + + def _build_document_sample_shuffle_indices( + self, + ) -> Tuple[numpy.ndarray, numpy.ndarray, numpy.ndarray]: + """Build the document index, the sample index, and the shuffle index + + The document index: + -- 1-D + -- An ordered array of document ids + + The sample index: + -- 2-D + -- The document indices and offsets which mark the start of every sample + + The shuffle index: + -- 1-D + -- A random permutation of index range of the sample index + + Returns: + Tuple[numpy.ndarray, numpy.ndarray, numpy.ndarray]: The document index, the sample + index, and the shuffle index + """ + path_to_cache = self.config.path_to_cache + if path_to_cache is None and not self.config.mock: + path_to_cache = os.path.join( + self.dataset.path_prefix, "cache", f"{type(self).__name__}_indices" + ) + + if path_to_cache: + base = f"{self.unique_description_hash}-{type(self).__name__}-{self.index_split.name}" + get_path_to = lambda affix: os.path.join(path_to_cache, f"{base}-{affix}") + path_to_description = get_path_to("description.txt") + path_to_document_index = get_path_to("document_index.npy") + path_to_sample_index = get_path_to("sample_index.npy") + path_to_shuffle_index = get_path_to("shuffle_index.npy") + cache_hit = ( + True + if self.config.fast_cache_load + else all( + map( + os.path.isfile, + [ + path_to_description, + path_to_document_index, + path_to_sample_index, + path_to_shuffle_index, + ], + ) + ) + ) + else: + cache_hit = False + + if not path_to_cache or ( + not cache_hit + and (not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0) + ): + log_single_rank( + logger, + logging.DEBUG, + f"Build and save the {type(self).__name__} {self.index_split.name} indices", + ) + t_beg = time.time() + + sequence_length = self.config.sequence_length + sample_lengths = self.dataset.sequence_lengths[self.indices] + num_samples = 1 if self.num_samples is None else self.num_samples + + packed_samples = pack_samples(sample_lengths, sequence_length) + + document_index = numpy.concatenate( + [numpy.array(pack, dtype=numpy.int32) for pack in packed_samples] + ) + + sample_index = numpy.empty(len(packed_samples) + 1, dtype=numpy.int32) + sample_index[0] = 0 + numpy.cumsum( + numpy.array([len(pack) for pack in packed_samples], dtype=numpy.int32), + out=sample_index[1:], + ) + + num_packed = len(packed_samples) + num_epochs = -(-num_samples // num_packed) + + numpy_random_state = numpy.random.RandomState(self.config.random_seed) + shuffle_index = numpy.empty(num_epochs * num_packed, dtype=numpy.int32) + for epoch in range(num_epochs): + epoch_perm = numpy.arange(num_packed, dtype=numpy.int32) + numpy_random_state.shuffle(epoch_perm) + shuffle_index[epoch * num_packed : (epoch + 1) * num_packed] = epoch_perm + shuffle_index = shuffle_index[:num_samples] + + if path_to_cache: + os.makedirs(path_to_cache, exist_ok=True) + # Write the description + with open(path_to_description, "wt") as writer: + writer.write(self.unique_description) + numpy.save(path_to_document_index, document_index, allow_pickle=True) + numpy.save(path_to_sample_index, sample_index, allow_pickle=True) + numpy.save(path_to_shuffle_index, shuffle_index, allow_pickle=True) + else: + log_single_rank( + logger, + logging.WARNING, + f"Unable to save {type(self).__name__} indexes because path_to_cache is None", + ) + + t_end = time.time() + log_single_rank(logger, logging.DEBUG, f"\t> time elapsed: {t_end - t_beg:4f} seconds") + + log_single_rank( + logger, logging.DEBUG, f"> total number of samples: {sample_index.shape[0] - 1}" + ) + log_single_rank(logger, logging.DEBUG, f"> total number of epochs: {num_epochs}") + + return document_index, sample_index, shuffle_index + + log_single_rank( + logger, logging.DEBUG, f"Load the {type(self).__name__} {self.index_split.name} indices" + ) + + log_single_rank( + logger, + logging.DEBUG, + f"\tLoad the document index from {os.path.basename(path_to_document_index)}", + ) + t_beg = time.time() + document_index = numpy.load(path_to_document_index, allow_pickle=True, mmap_mode="r") + t_end = time.time() + log_single_rank(logger, logging.DEBUG, f"\t> time elapsed: {t_end - t_beg:4f} seconds") + + log_single_rank( + logger, + logging.DEBUG, + f"\tLoad the sample index from {os.path.basename(path_to_sample_index)}", + ) + t_beg = time.time() + sample_index = numpy.load(path_to_sample_index, allow_pickle=True, mmap_mode="r") + t_end = time.time() + log_single_rank(logger, logging.DEBUG, f"\t> time elapsed: {t_end - t_beg:4f} seconds") + + log_single_rank( + logger, + logging.DEBUG, + f"\tLoad the shuffle index from {os.path.basename(path_to_shuffle_index)}", + ) + t_beg = time.time() + shuffle_index = numpy.load(path_to_shuffle_index, allow_pickle=True, mmap_mode="r") + t_end = time.time() + log_single_rank(logger, logging.DEBUG, f"\t> time elapsed: {t_end - t_beg:4f} seconds") + + log_single_rank( + logger, logging.DEBUG, f"> total number of samples: {sample_index.shape[0] - 1}" + ) + + return document_index, sample_index, shuffle_index + + +def _classify_items( + items: List[Tuple[int, int]], bin_capacity: int +) -> Tuple[ + List[Tuple[int, int]], List[Tuple[int, int]], List[Tuple[int, int]], List[Tuple[int, int]] +]: + """Split items into large / medium / small / tiny classes. + + Follows the classification used by Johnson & Garey: + large : (C/2, C] + medium : (C/3, C/2] + small : (C/6, C/3] + tiny : (0 , C/6] + + Args: + items: List of (index, size) tuples + + Returns: + Tuple of four lists (large, medium, small, tiny) without additional sorting. + """ + large, medium, small, tiny = [], [], [], [] + for idx, size in items: + if size > bin_capacity / 2: + large.append((idx, size)) + elif size > bin_capacity / 3: + medium.append((idx, size)) + elif size > bin_capacity / 6: + small.append((idx, size)) + else: + tiny.append((idx, size)) + return large, medium, small, tiny + + +def pack_samples(sequence_lengths: List[int], bin_capacity: int) -> List[List[int]]: + """Pack sequences using the Modified First-Fit Decreasing algorithm. + + Args: + sequence_lengths: A list of sequence lengths to pack. + + Returns: + A list of bins, where each bin is a list of indices into the original + sequence_lengths list. + """ + # Validate inputs + if bin_capacity <= 0: + raise ValueError("bin_capacity must be positive") + if any(l <= 0 for l in sequence_lengths): + raise ValueError("sequence lengths must be positive") + + # Drop documents that exceed capacity and warn + long_mask = [l > bin_capacity for l in sequence_lengths] + if any(long_mask): + n_dropped = sum(long_mask) + log_single_rank( + logger, + logging.WARNING, + f"Dropping {n_dropped} document(s) with sequence length > bin_capacity " + f"(bin_capacity={bin_capacity}).", + ) + items: List[Tuple[int, int]] = [ + (i, l) for i, l in enumerate(sequence_lengths) if l <= bin_capacity + ] + + # Phase-0: classify + large, medium, small, tiny = _classify_items(items, bin_capacity) + + # Sort according to the rules of MFFD + large.sort(key=lambda x: x[1], reverse=True) # descending size + medium.sort(key=lambda x: x[1], reverse=True) + small.sort(key=lambda x: x[1]) # ascending size + tiny.sort(key=lambda x: x[1]) + + # Phase-1: start one bin per large item + bins: List[List[Tuple[int, int]]] = [[item] for item in large] + + # Phase-2: try to add one medium item to each large bin (forward pass) + for b in bins: + remaining = bin_capacity - sum(size for _, size in b) + for i, (idx, size) in enumerate(medium): + if size <= remaining: + b.append(medium.pop(i)) + break + + # Phase-3: backward pass – fill with two small items where possible + for b in reversed(bins): + has_medium = any(bin_capacity / 3 < size <= bin_capacity / 2 for _, size in b) + if has_medium or len(small) < 2: + continue + remaining = bin_capacity - sum(size for _, size in b) + if small[0][1] + small[1][1] > remaining: + continue + first_small = small.pop(0) + # pick the *largest* small that fits with first_small (so iterate from end) + second_idx = None + for j in range(len(small) - 1, -1, -1): + if small[j][1] <= remaining - first_small[1]: + second_idx = j + break + if second_idx is not None: + second_small = small.pop(second_idx) + b.extend([first_small, second_small]) + + # Phase-4: forward greedy fit of remaining items + remaining_items = sorted(medium + small + tiny, key=lambda x: x[1], reverse=True) + for b in bins: + while remaining_items: + rem = bin_capacity - sum(size for _, size in b) + # if even the smallest remaining doesn't fit we break + if rem < remaining_items[-1][1]: + break + + # pick the first (largest) that fits + chosen_idx = None + for i, (_, size) in enumerate(remaining_items): + if size <= rem: + chosen_idx = i + break + if chosen_idx is None: + break + b.append(remaining_items.pop(chosen_idx)) + + # Phase-5: FFD on leftovers + leftovers = remaining_items # renamed for clarity + + # New O(n * logn) implementation + ffd_bins: List[List[Tuple[int, int]]] = [[]] + ffd_bin_sizes: List[int] = [0] + for idx, size in sorted(leftovers, key=lambda x: x[1], reverse=True): + # We only need to check the first bin since we guarantee the order + # of ffd_bin_sizes to be sorted from smallest to largest. + if size <= (bin_capacity - ffd_bin_sizes[0]): + new_bin = ffd_bins.pop(0) + new_bin_size = ffd_bin_sizes.pop(0) + else: + new_bin = [] + new_bin_size = 0 + + new_bin.append((idx, size)) + new_bin_size += size + + new_idx = bisect(ffd_bin_sizes, new_bin_size) + ffd_bins.insert(new_idx, new_bin) + ffd_bin_sizes.insert(new_idx, new_bin_size) + + bins.extend(ffd_bins) + + # Convert to list of index lists (discard sizes) + return [[idx for idx, _ in b] for b in bins if b] + + +def find_subsequence(sequence, subsequence, start=0): + """Return the index of the first occurrence of *subsequence* in *sequence*, or -1. + + Args: + sequence (list): The sequence to search in. + subsequence (list): The contiguous subsequence to find. + start (int): Position in *sequence* at which to begin the search. + + Returns: + int: Index of the first match, or -1 if not found. + """ + sub_len = len(subsequence) + for i in range(start, len(sequence) - sub_len + 1): + if sequence[i : i + sub_len] == subsequence: + return i + return -1 + + +def _split_tool_calls(tokens, offset, tool_call_start_tokens, tool_call_end_tokens): + """Split a token sequence into assistant text and tool_call sub-segments. + + Whitespace-only (\n) assistant fragments between and + are dropped so we don't produce meaningless segments. + """ + NL_TOKEN = tool_call_start_tokens[-1] # \n token is the last token in the start marker + tc_start_len = len(tool_call_start_tokens) + tc_end_len = len(tool_call_end_tokens) + result = [] + pos = 0 + while pos < len(tokens): + tc_start = find_subsequence(tokens, tool_call_start_tokens, pos) + + if tc_start == -1: + if pos < len(tokens): + result.append( + { + "role": "assistant", + "tokens": tokens[pos:], + "start": offset + pos, + "end": offset + len(tokens), + } + ) + break + + # Assistant content before tool_call (skip if whitespace-only) + if tc_start > pos: + frag = tokens[pos:tc_start] + if not all(t == NL_TOKEN for t in frag): + result.append( + { + "role": "assistant", + "tokens": frag, + "start": offset + pos, + "end": offset + tc_start, + } + ) + + content_start = tc_start + tc_start_len + tc_end = find_subsequence(tokens, tool_call_end_tokens, content_start) + + if tc_end == -1: + result.append( + { + "role": "tool_call", + "tokens": tokens[content_start:], + "start": offset + content_start, + "end": offset + len(tokens), + } + ) + break + + result.append( + { + "role": "tool_call", + "tokens": tokens[content_start:tc_end], + "start": offset + content_start, + "end": offset + tc_end, + } + ) + pos = tc_end + tc_end_len + + # Drop trailing whitespace-only assistant fragments + if ( + result + and result[-1]["role"] == "assistant" + and all(t == NL_TOKEN for t in result[-1]["tokens"]) + ): + result.pop() + + return result + + +def extract_segments( + tokenized_conversation, + role_start_tokens, + end_tokens, + think_start_tokens, + think_end_tokens, + tool_call_start_tokens, + tool_call_end_tokens, + tool_response_start_tokens, +): + """Parse a tokenized conversation into labeled segments for loss masking. + + Scans *tokenized_conversation* for role-start and end delimiter token + sequences, then classifies each segment as one of: ``system``, ``user``, + ``assistant``, ``reasoning``, ``tool_call``, or ``tool_response``. + Assistant segments are further split on ````/```` and + ````/```` boundaries. + + Args: + tokenized_conversation (List[int]): Full token-id sequence of one conversation. + role_start_tokens (Dict[str, List[int]]): Mapping from role name to its + start-delimiter token ids. + end_tokens (List[int]): Token ids for the end-of-turn delimiter. + think_start_tokens (List[int]): Token ids for ````. + think_end_tokens (List[int]): Token ids for ````. + tool_call_start_tokens (List[int]): Token ids for ````. + tool_call_end_tokens (List[int]): Token ids for ````. + tool_response_start_tokens (List[int]): Token ids for ````. + + Returns: + List[dict]: Each dict has keys ``role`` (str), ``tokens`` (List[int]), + ``start`` (int), and ``end`` (int) giving absolute indices into + *tokenized_conversation*. + """ + markers = [] + for role, start_tokens in role_start_tokens.items(): + pos = 0 + while True: + idx = find_subsequence(tokenized_conversation, start_tokens, pos) + if idx == -1: + break + markers.append((idx, role, len(start_tokens))) + pos = idx + len(start_tokens) + markers.sort(key=lambda x: x[0]) + + segments = [] + for start_pos, role, marker_len in markers: + content_start = start_pos + marker_len + end_pos = find_subsequence(tokenized_conversation, end_tokens, content_start) + if end_pos == -1: + content_end = len(tokenized_conversation) + else: + content_end = end_pos + content_tokens = tokenized_conversation[content_start:content_end] + + # Check if this user turn is actually a tool response + tr_start_len = len(tool_response_start_tokens) + if ( + role == "user" + and len(content_tokens) >= tr_start_len + and content_tokens[:tr_start_len] == tool_response_start_tokens + ): + segments.append( + { + "role": "tool_response", + "tokens": content_tokens, + "start": content_start, + "end": content_end, + } + ) + continue + + if role == "assistant": + think_start_idx = find_subsequence(content_tokens, think_start_tokens) + if think_start_idx != -1: + think_end_idx = find_subsequence( + content_tokens, think_end_tokens, think_start_idx + len(think_start_tokens) + ) + else: + think_end_idx = -1 + + if think_start_idx != -1 and think_end_idx != -1: + reasoning_tokens = content_tokens[ + think_start_idx + len(think_start_tokens) : think_end_idx + ] + response_tokens = content_tokens[think_end_idx + len(think_end_tokens) :] + if reasoning_tokens: + abs_start = content_start + think_start_idx + len(think_start_tokens) + abs_end = content_start + think_end_idx + segments.append( + { + "role": "reasoning", + "tokens": reasoning_tokens, + "start": abs_start, + "end": abs_end, + } + ) + # Split the response part by tool calls + if response_tokens: + abs_start = content_start + think_end_idx + len(think_end_tokens) + segments.extend( + _split_tool_calls( + response_tokens, abs_start, tool_call_start_tokens, tool_call_end_tokens + ) + ) + continue + + # No think tags — split entire content by tool calls + if find_subsequence(content_tokens, tool_call_start_tokens) != -1: + segments.extend( + _split_tool_calls( + content_tokens, content_start, tool_call_start_tokens, tool_call_end_tokens + ) + ) + continue + + segments.append( + {"role": role, "tokens": content_tokens, "start": content_start, "end": content_end} + ) + + return segments diff --git a/megatron/core/tokenizers/text/libraries/abstract_tokenizer.py b/megatron/core/tokenizers/text/libraries/abstract_tokenizer.py index 8f5b7d3b4f5..9b8b38990f9 100644 --- a/megatron/core/tokenizers/text/libraries/abstract_tokenizer.py +++ b/megatron/core/tokenizers/text/libraries/abstract_tokenizer.py @@ -62,7 +62,7 @@ def ids_to_tokens(self, ids: List[int]) -> List[str]: pass @abstractmethod - def text_to_ids(self, text: str) -> List[int]: + def text_to_ids(self, text: str, add_special_tokens: bool = True) -> List[int]: """ Converts text to ids. diff --git a/megatron/core/tokenizers/text/libraries/bytelevel_tokenizer.py b/megatron/core/tokenizers/text/libraries/bytelevel_tokenizer.py index 909f6cd4518..6e6944b28ab 100644 --- a/megatron/core/tokenizers/text/libraries/bytelevel_tokenizer.py +++ b/megatron/core/tokenizers/text/libraries/bytelevel_tokenizer.py @@ -70,7 +70,7 @@ def tokens_to_text(self, tokens): """ return self.ids_to_text(tokens) - def text_to_ids(self, text): + def text_to_ids(self, text, add_special_tokens=True): """ Convert a text to a list of IDs. """ diff --git a/megatron/core/tokenizers/text/libraries/huggingface_tokenizer.py b/megatron/core/tokenizers/text/libraries/huggingface_tokenizer.py index bed8d9c5ad3..d4137599583 100644 --- a/megatron/core/tokenizers/text/libraries/huggingface_tokenizer.py +++ b/megatron/core/tokenizers/text/libraries/huggingface_tokenizer.py @@ -259,7 +259,7 @@ def ids_to_tokens(self, ids: List[int]) -> List[str]: tokens = self.tokenizer.convert_ids_to_tokens(ids) return tokens - def text_to_ids(self, text: str) -> List[int]: + def text_to_ids(self, text: str, add_special_tokens: bool = True) -> List[int]: """Converts text to tokens ids.""" if self.include_special_tokens: return self.tokenizer(text).input_ids diff --git a/megatron/core/tokenizers/text/libraries/null_tokenizer.py b/megatron/core/tokenizers/text/libraries/null_tokenizer.py index 96a0d3afd57..14745a07c26 100644 --- a/megatron/core/tokenizers/text/libraries/null_tokenizer.py +++ b/megatron/core/tokenizers/text/libraries/null_tokenizer.py @@ -19,9 +19,9 @@ def __init__(self, vocab_size, eod_id=None, pad_id=-1, **kwargs): self._eod_id = int(eod_id) if eod_id is not None else self._vocab_size - 1 self._pad_id = int(pad_id) - def text_to_ids(self, text): + def text_to_ids(self, text, add_special_tokens=True): """Converts text to ids.""" - return [int(x) for x in text.split(' ')] + return [ord(x) % self._vocab_size for x in text] def ids_to_text(self, ids): """Converts ids to text.""" diff --git a/megatron/core/tokenizers/text/libraries/sentencepiece_tokenizer.py b/megatron/core/tokenizers/text/libraries/sentencepiece_tokenizer.py index feaf1c4e9a1..866f8af047e 100644 --- a/megatron/core/tokenizers/text/libraries/sentencepiece_tokenizer.py +++ b/megatron/core/tokenizers/text/libraries/sentencepiece_tokenizer.py @@ -138,7 +138,7 @@ def text_to_tokens(self, text: str) -> List[str]: tokens = list(filter(lambda x: x != self.extra_space_token, tokens)) return tokens - def text_to_ids(self, text, sample_alpha=None) -> List[int]: + def text_to_ids(self, text, sample_alpha=None, add_special_tokens=True) -> List[int]: """Converts text to tokens ids.""" if isinstance(text, str): return self._text_to_ids(text, sample_alpha) diff --git a/megatron/core/tokenizers/text/libraries/sft_tokenizer.py b/megatron/core/tokenizers/text/libraries/sft_tokenizer.py index bc29ddfab68..ebdce8d29bf 100644 --- a/megatron/core/tokenizers/text/libraries/sft_tokenizer.py +++ b/megatron/core/tokenizers/text/libraries/sft_tokenizer.py @@ -175,7 +175,10 @@ def tokenize_conversation( if turn["role"].lower() == "assistant" and len(turn["content"]) == 0: raise ValueError(f"empty assistant turn in conversation: {conversation}.") if turn["role"].lower() == "assistant": - assert conversation[turn_idx - 1]["role"].lower() in ("user", "tool") + assert conversation[turn_idx - 1]["role"].lower() in ( + "user", + "tool", + ), "Assistant turn must be preceded by a user or tool turn" turn_tokens = self._extract_token_ids( self._tokenizer.apply_chat_template( @@ -208,7 +211,7 @@ def tokenize_conversation( return tokens, target - def text_to_ids(self, text: Union[str, List[Dict]]): + def text_to_ids(self, text: Union[str, List[Dict]], add_special_tokens: bool = True): """Tokenize conversation or string input.""" if isinstance(text, list): # This code path is used by the inference code currently. @@ -216,7 +219,7 @@ def text_to_ids(self, text: Union[str, List[Dict]]): text, return_target=False, add_generation_prompt=True ).tolist() - return self._tokenizer.encode(text) + return self._tokenizer.encode(text, add_special_tokens=add_special_tokens) def tokens_to_ids(self, tokens: List[str]): """Convert tokens to IDs.""" diff --git a/megatron/core/tokenizers/text/libraries/tiktoken_tokenizer.py b/megatron/core/tokenizers/text/libraries/tiktoken_tokenizer.py index eb98ab98c05..9cdedd68839 100644 --- a/megatron/core/tokenizers/text/libraries/tiktoken_tokenizer.py +++ b/megatron/core/tokenizers/text/libraries/tiktoken_tokenizer.py @@ -206,7 +206,7 @@ def ids_to_tokens(self, token_ids: List[int]) -> List[str]: return tokens - def text_to_ids(self, text: str) -> List[int]: + def text_to_ids(self, text: str, add_special_tokens: bool = True) -> List[int]: """Converts text to list of ids.""" tokens = self.tokenizer.encode(text, allowed_special="all") return tokens diff --git a/megatron/core/tokenizers/text/text_tokenizer.py b/megatron/core/tokenizers/text/text_tokenizer.py index f7927776138..5746029be22 100644 --- a/megatron/core/tokenizers/text/text_tokenizer.py +++ b/megatron/core/tokenizers/text/text_tokenizer.py @@ -63,7 +63,7 @@ def _restore_model(self, **kwargs) -> MegatronTokenizerTextAbstract: else: return library_class(self.path, **kwargs) - def tokenize(self, text: str) -> List[int]: + def tokenize(self, text: str, add_special_tokens: bool = True) -> List[int]: """ Text tokenization. @@ -74,7 +74,7 @@ def tokenize(self, text: str) -> List[int]: list: list of ids. """ - return self._tokenizer.text_to_ids(text) + return self._tokenizer.text_to_ids(text, add_special_tokens=add_special_tokens) def detokenize(self, ids: List[int], skip_special_tokens: Optional[bool] = None) -> str: """ diff --git a/megatron/core/tokenizers/utils/build_tokenizer.py b/megatron/core/tokenizers/utils/build_tokenizer.py index 5f87c0ea34a..83ed8b1831d 100644 --- a/megatron/core/tokenizers/utils/build_tokenizer.py +++ b/megatron/core/tokenizers/utils/build_tokenizer.py @@ -22,6 +22,8 @@ logger = logging.getLogger(__name__) +NULL_TOKENIZERS = {'NullTokenizer': 'null-text', 'NullMultimodalTokenizer': 'null-multimodal'} + def build_tokenizer(args, **kwargs): """Initialize tokenizer.""" @@ -84,10 +86,8 @@ def build_tokenizer(args, **kwargs): tokenizer_library = 'sft' tokenizer_path = args.tokenizer_model kwargs['prompt_format'] = args.sft_tokenizer_prompt_format - elif args.tokenizer_type in ['NullTokenizer', 'NullMultimodalTokenizer']: - tokenizer_library = ( - 'null-text' if args.tokenizer_type == 'NullTokenizer' else 'null-multimodal' - ) + elif args.tokenizer_type in NULL_TOKENIZERS.keys(): + tokenizer_library = NULL_TOKENIZERS[args.tokenizer_type] metadata = {'library': tokenizer_library} if args.vocab_size: kwargs['vocab_size'] = args.vocab_size diff --git a/megatron/core/utils.py b/megatron/core/utils.py index 2cc5d635f48..d6ae6098df8 100644 --- a/megatron/core/utils.py +++ b/megatron/core/utils.py @@ -2031,6 +2031,295 @@ def is_submodule(module, parent_module, strict=True): return False +######################## +###### sft utils ####### +######################## + + +def pad_thd_sequences_for_cp( + tensors_with_pad_values: List[Tuple[torch.Tensor, Union[int, float]]], + cu_seqlens: torch.Tensor, + divisibility_factor: int, +) -> Tuple[List[torch.Tensor], torch.Tensor]: + """Round every sub-sequence of one or more THD-packed tensors up to a length + that is a multiple of ``divisibility_factor``. + + This rounding is required for context-parallel sequence sharding, which + splits each segment evenly across ``2 * cp_size`` chunks (and an additional + factor of ``tp_size`` when sequence parallelism is enabled). + + All input tensors share the same segment layout described by ``cu_seqlens`` + and are padded together so the returned ``cu_seqlens_padded`` is the single + authoritative segment layout for every output tensor. + + Args: + tensors_with_pad_values: List of ``(tensor, pad_value)`` pairs. Each + tensor must be 1-D ``(N,)`` or 2-D ``(1, N)`` (a leading batch dim + of 1 is squeezed). Every tensor must share the same length ``N``, + which has to equal ``cu_seqlens[-1]``. ``pad_value`` is cast to the + tensor's dtype (e.g. ``padding_token_id`` for ``input_ids``, ``-100`` + for ``labels``, ``0`` for ``loss_mask``). + cu_seqlens: 1-D cumulative segment lengths starting at 0 with + ``cu_seqlens[-1] == N``. + divisibility_factor: Round each segment length up to the next multiple + of this value. + + Returns: + - ``padded_tensors``: List of padded 1-D tensors in the same order as + the input. Each has shape ``(N_padded,)`` where ``N_padded`` is the + sum of the rounded-up segment lengths. + - ``cu_seqlens_padded``: 1-D tensor of cumulative padded segment + lengths. Same length and dtype as ``cu_seqlens``. + """ + # Squeeze the optional leading batch dim once per tensor. + flat_tensors = [ + (t.squeeze(0) if t.dim() == 2 else t, pad) for t, pad in tensors_with_pad_values + ] + + # Per-segment original and padded lengths (preserve cu_seqlens dtype throughout). + seqlens = cu_seqlens[1:] - cu_seqlens[:-1] + padded_lengths = ( + (seqlens + divisibility_factor - 1) // divisibility_factor + ) * divisibility_factor + pad_amounts = (padded_lengths - seqlens).tolist() + + starts = cu_seqlens[:-1].tolist() + ends = cu_seqlens[1:].tolist() + + padded_tensors: List[torch.Tensor] = [] + for tensor, pad_value in flat_tensors: + pieces: List[torch.Tensor] = [] + for start, end, pad in zip(starts, ends, pad_amounts): + pieces.append(tensor[start:end]) + if pad > 0: + pieces.append(tensor.new_full((pad,), pad_value)) + padded_tensors.append(torch.cat(pieces)) + + # `torch.cumsum` promotes int32 to int64 by default; write directly into a + # buffer of the right dtype so cu_seqlens_padded matches cu_seqlens. + cu_seqlens_padded = torch.empty_like(cu_seqlens) + cu_seqlens_padded[0] = 0 + torch.cumsum(padded_lengths, dim=0, out=cu_seqlens_padded[1:]) + + return padded_tensors, cu_seqlens_padded + + +def pad_or_truncate_thd_tensors( + input_ids: torch.Tensor, + labels: torch.Tensor, + loss_mask: torch.Tensor, + cu_seqlens: torch.Tensor, + cu_seqlens_padded: torch.Tensor, + sequence_length: int, + padding_token_id: int, + padding_label_id: int, +): + """Pad or truncate THD-format (Token-Head-Dim packed) tensors to a fixed sequence length. + + For SFT with packed sequences, the total number of tokens across all sub-sequences + in a sample may not match the target sequence length. This function either: + - Truncates: clips input_ids, labels, loss_mask, and cu_seqlens/cu_seqlens_padded + so that the total token count equals ``sequence_length``. Sub-sequences that + extend beyond the boundary are dropped, and a final cumulative length entry is + appended. + - Pads: extends input_ids, labels, and loss_mask with padding values to reach + ``sequence_length``, and updates the last entry of cu_seqlens (and + cu_seqlens_padded if present) accordingly. + + All input tensors are expected to be 1-D (no batch dimension). + + Args: + input_ids (torch.Tensor): 1-D token IDs for the packed sample. + labels (torch.Tensor): 1-D label IDs for the packed sample. + loss_mask (torch.Tensor): 1-D loss mask for the packed sample (1 = train, 0 = mask). + cu_seqlens (torch.Tensor): 1-D cumulative sequence lengths (int32), starting + at 0 and ending at the total token count. + cu_seqlens_padded (torch.Tensor | None): 1-D cumulative sequence lengths after + CP padding (from ``pad_thd_sequences_for_cp``), or None when CP is disabled. + sequence_length (int): Target total token count (i.e., ``max_seq_len``). + padding_token_id (int): Token ID used for padding input_ids. + padding_label_id (int): Label ID used for padding labels (typically -100 or similar + ignore index). + + Returns: + Tuple of (input_ids, labels, loss_mask, cu_seqlens, cu_seqlens_padded), all + adjusted to ``sequence_length``. + """ + if input_ids.shape[0] > sequence_length: # Truncate + input_ids = input_ids[:sequence_length] + labels = labels[:sequence_length] + loss_mask = loss_mask[:sequence_length] + if cu_seqlens_padded is not None: + # NOTE(asolergi-nv): When CP padding is active, cu_seqlens_padded + # determines the actual token layout. Because padded entries are + # always >= original entries, cu_seqlens_padded hits + # sequence_length first. Truncate BOTH at the same segment + # boundary so they always have the same number of entries. + idx = (cu_seqlens_padded < sequence_length).nonzero(as_tuple=True)[0] + num_keep = idx[-1] + 1 + cu_seqlens = torch.cat( + [cu_seqlens[:num_keep], cu_seqlens.new_tensor([sequence_length])] + ) + cu_seqlens_padded = torch.cat( + [cu_seqlens_padded[:num_keep], cu_seqlens_padded.new_tensor([sequence_length])] + ) + else: + # NOTE(asolergi-nv): Truncate cu_seqlens + # Find the largest index such that cu_seqlens[index] < sequence_length + idx = (cu_seqlens < sequence_length).nonzero(as_tuple=True)[0] + cu_seqlens = torch.cat( + [cu_seqlens[: idx[-1] + 1], cu_seqlens.new_tensor([sequence_length])] + ) + else: # Pad + input_ids = torch.cat( + [ + input_ids, + torch.full( + (sequence_length - input_ids.shape[0],), + padding_token_id, + dtype=input_ids.dtype, + device=input_ids.device, + ), + ] + ) + labels = torch.cat( + [ + labels, + torch.full( + (sequence_length - labels.shape[0],), + padding_label_id, + dtype=labels.dtype, + device=labels.device, + ), + ] + ) + loss_mask = torch.cat( + [ + loss_mask, + torch.zeros( + sequence_length - loss_mask.shape[0], + dtype=loss_mask.dtype, + device=loss_mask.device, + ), + ] + ) + cu_seqlens = torch.cat([cu_seqlens[:-1], cu_seqlens.new_tensor([sequence_length])]) + # NOTE(asolergi-nv): Pad cu_seqlens_padded if CP + if cu_seqlens_padded is not None: + cu_seqlens_padded = torch.cat( + [cu_seqlens_padded[:-1], cu_seqlens_padded.new_tensor([sequence_length])] + ) + + return input_ids, labels, loss_mask, cu_seqlens, cu_seqlens_padded + + +def preprocess_sft_batch( + batch: Dict[str, Any], + tp_rank: int, + cp_size: int, + tp_size: int, + sp: bool, + padding_token_id: int, + padding_label_id: int, + max_seq_len: int, +): + """Preprocess an SFT batch on TP rank 0 before broadcasting to other TP ranks. + + Performs the following preprocessing steps (only on ``tp_rank == 0``): + 1. Pads individual sub-sequences for CP divisibility via + ``pad_thd_sequences_for_cp`` (when ``cp_size > 1``), producing + ``cu_seqlens_padded``. + 2. Pads or truncates the packed sample to exactly ``max_seq_len`` tokens. + 3. Creates a loss mask that zeros out padding tokens and prompt tokens. + 4. Computes per-segment position IDs from cumulative sequence lengths. + 5. Computes ``max_seqlen`` (the length of the longest sub-sequence). + + After this function, all sequence tensors have shape ``[1, max_seq_len]`` + (with a batch dimension). ``cu_seqlens`` and ``cu_seqlens_padded`` carry the + same leading batch dim (shape ``[1, n]``) so they round-trip through + ``get_batch_on_this_tp_rank``'s length-prefixed broadcast. ``max_seqlen`` + remains 1-D with a single element. + + Args: + batch (Dict[str, Any]): Raw batch from the dataloader containing at minimum + 'tokens', 'labels', and 'cu_seqlens' (each with a leading batch dim of 1). + tp_rank (int): Tensor-parallel rank. Preprocessing is only done on rank 0. + cp_size (int): Context-parallel world size. + tp_size (int): Tensor-parallel world size. + sp (bool): Whether sequence parallelism is enabled. + padding_token_id (int): Token ID used for input padding. + padding_label_id (int): Label ID used for prompt/padding masking (e.g., -100). + max_seq_len (int): Target sequence length for the batch. + + Returns: + Dict[str, Any]: Preprocessed batch dict with keys 'tokens', 'labels', + 'loss_mask', 'position_ids', 'cu_seqlens', 'cu_seqlens_padded', and + 'max_seqlen'. + """ + if tp_rank == 0: + tokens, labels, loss_mask, cu_seqlens = ( + batch["tokens"].squeeze(0), + batch["labels"].squeeze(0), + batch["loss_mask"].squeeze(0), + batch["cu_seqlens"].squeeze(0), + ) # NOTE(asolergi-nv): PyTorch DataLoader `default_collate` adds batch dimension, + # so we need to remove it since TE expects cu_seqlens to be 1D + + if cp_size > 1: + divisibility_factor = cp_size * 2 + if tp_size > 1 and sp: + divisibility_factor *= tp_size + + (tokens, labels, loss_mask), cu_seqlens_padded = pad_thd_sequences_for_cp( + [(tokens, padding_token_id), (labels, padding_label_id), (loss_mask, 0)], + cu_seqlens, + divisibility_factor, + ) + else: + cu_seqlens_padded = None + + tokens, labels, loss_mask, cu_seqlens, cu_seqlens_padded = pad_or_truncate_thd_tensors( + tokens, + labels, + loss_mask, + cu_seqlens, + cu_seqlens_padded, + max_seq_len, + padding_token_id, + padding_label_id, + ) + + # Position ids. + seg_lengths = ( + cu_seqlens[1:] - cu_seqlens[:-1] + if cu_seqlens_padded is None + else cu_seqlens_padded[1:] - cu_seqlens_padded[:-1] + ) + position_ids = torch.cat([torch.arange(length) for length in seg_lengths]) + + seq_lens = ( + cu_seqlens[1:] - cu_seqlens[:-1] + if cu_seqlens_padded is None + else cu_seqlens_padded[1:] - cu_seqlens_padded[:-1] + ) + max_seqlen = torch.tensor([seq_lens.max().item()], dtype=torch.int32) + + batch = { + 'tokens': tokens.unsqueeze(0), # NOTE(asolergi-nv): Add back batch dimension + 'labels': labels.unsqueeze(0), # NOTE(asolergi-nv): Add back batch dimension + 'loss_mask': loss_mask.unsqueeze(0), # NOTE(asolergi-nv): Add batch dimension + 'position_ids': position_ids.unsqueeze(0), # NOTE(asolergi-nv): Add batch dimension + # cu_seqlens / cu_seqlens_padded carry the dataloader's batch dim (1, n) + # through get_batch_on_this_tp_rank's length-prefixed broadcast. + 'cu_seqlens': cu_seqlens.unsqueeze(0), + 'cu_seqlens_padded': ( + cu_seqlens_padded.unsqueeze(0) if cu_seqlens_padded is not None else None + ), + 'max_seqlen': max_seqlen, + } + return batch + + ######################## ### tensor parallel #### ######################## diff --git a/megatron/elastification/pretrain_hybrid_flex.py b/megatron/elastification/pretrain_hybrid_flex.py index 13eeca1f7c8..25f67bf2aa9 100644 --- a/megatron/elastification/pretrain_hybrid_flex.py +++ b/megatron/elastification/pretrain_hybrid_flex.py @@ -10,6 +10,7 @@ from megatron.core import mpu, parallel_state from megatron.core.datasets.blended_megatron_dataset_builder import BlendedMegatronDatasetBuilder from megatron.core.datasets.gpt_dataset import GPTDataset, GPTDatasetConfig, MockGPTDataset +from megatron.core.datasets.sft_dataset import SFTDataset from megatron.core.enums import ModelType from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_with_transformer_engine_spec from megatron.core.models.hybrid.hybrid_model import HybridModel @@ -50,7 +51,6 @@ ) from megatron.training.argument_utils import pretrain_cfg_container_from_args from megatron.training.arguments import core_transformer_config_from_args, parse_and_validate_args -from megatron.training.datasets.sft_dataset import SFTDataset from megatron.training.utils import get_blend_and_blend_per_split, is_first_or_last_pipeline_stage # modelopt distillation diff --git a/megatron/training/datasets/sft_dataset.py b/megatron/training/datasets/sft_dataset.py deleted file mode 100644 index 3f93927387d..00000000000 --- a/megatron/training/datasets/sft_dataset.py +++ /dev/null @@ -1,201 +0,0 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. - -import atexit, json -from collections import Counter -from typing import Any, Dict, Optional - -import numpy as np -import torch - -from megatron.core.datasets.gpt_dataset import GPTDatasetConfig -from megatron.core.datasets.megatron_dataset import LowLevelDataset, MegatronDataset -from megatron.core.datasets.utils import Split - -IGNORE_INDEX = -100 - - -class SFTLowLevelDataset: - """The low-level dataset loading jsonl data for SFT - - Args: - dataset_path (str): The path to jsonl data - Each line of the jsonl must have key "messages" (List[Dict]), - which is a sequence of system/user/assistant messages. - Must be in the following format: - [ - {"role": "system", "content": "something"}, - {"role": "user", "content": "something1"}, - {"role": "assistant", "content": "something2"}, - ] - A jsonl line can contain multiple conversations packed together into on list. Each - conversation starts with the system role, and conversations can have multiple turns - of the user and assistant roles. - """ - - def __init__(self, dataset_path: str) -> None: - try: - from datasets import load_dataset - except ImportError: - raise ImportError( - "SFTDataset currently requires datasets library to be installed" - ) - self.dataset = load_dataset("json", data_files=dataset_path, split="all") - - def __len__(self) -> int: - return len(self.dataset) - - def __getitem__(self, idx: int) -> list: - return self.dataset[idx]["messages"] - - -class SFTDataset(MegatronDataset): - """The dataset used during SFT""" - - def __init__( - self, - dataset: LowLevelDataset, - dataset_path: Optional[str], - indices: np.ndarray, - num_samples: Optional[int], - index_split: Split, - config: GPTDatasetConfig, - ) -> None: - super().__init__(dataset, dataset_path, indices, num_samples, index_split, config) - - @staticmethod - def numel_low_level_dataset(low_level_dataset: LowLevelDataset) -> int: - return len(low_level_dataset) - - @staticmethod - def build_low_level_dataset(dataset_path: str, config: GPTDatasetConfig) -> LowLevelDataset: - return SFTLowLevelDataset(dataset_path) - - def __len__(self) -> int: - return self.num_samples - - def _split_conversations(self, merged_conversations): - split_conversations = [] - current = [] - for msg in merged_conversations: - # Whenever we see a new system message, start a new conversation - if msg["role"] == "system": - if current: # If previously accumulating a conversation, then store it - split_conversations.append(current) - current = [msg] # Then start the new conversation - else: - current.append(msg) # Continue accumulating the current conversation - if current: # Store any remaining conversation - split_conversations.append(current) - return split_conversations - - def __getitem__(self, idx: int) -> Dict[str, Any]: - - tokenizer = self.config.tokenizer - pack_length = self.config.sequence_length - - merged_conversations = self.dataset[int(self.indices[idx % len(self.indices)])] - split_conversations = self._split_conversations(merged_conversations) - - def extend_with_padding(tokens, targets, positions, pad_len): - tokens.extend([pad] * pad_len) - targets.extend([pad] * pad_len) - positions.extend(range(positions[-1]+1, positions[-1]+1+pad_len)) - - pack_tokens = [] - pack_targets = [] - pack_positions = [] - cu_seqlens = [0] - eod = tokenizer.eod - pad = tokenizer.pad - # TODO(duncan): Track number of convs dropped and/or truncated and amount of end-padding - for conversation in split_conversations: - - tokens, targets = tokenizer.tokenize_conversation( - conversation, return_target=True, add_generation_prompt=False - ) - - tokens_list = tokens.tolist() - targets_list = targets.tolist() - - - pack_tokens.extend(tokens_list) - pack_targets.extend(targets_list) - - assert not self.config.reset_position_ids - pack_positions.extend(range(len(tokens_list))) - - if self.config.context_parallel_size > 1: - pad_granularity = self.config.context_parallel_size * 2 - mod_token_count = len(pack_tokens) % pad_granularity - if mod_token_count != 0: - pad_len = pad_granularity - mod_token_count - extend_with_padding(pack_tokens, pack_targets, pack_positions, pad_len) - - # TODO(duncan): Consider also padding to multiple of number of tokens here. This might - # be needed for efficiency (and potentially set via command-line argument). - - cu_seqlens.append(len(pack_tokens)) - - # Handle any necessary truncation - if len(pack_tokens) >= pack_length + 1: # +1 here to account for later alignment - # Truncate on the right - max_body = pack_length - pack_tokens = pack_tokens[:max_body] - pack_targets = pack_targets[:max_body] - pack_tokens.append(pad) - pack_targets.append(pad) - pack_positions = pack_positions[:pack_length+1] - # Note len({pack_tokens, pack_targets, pack_positions}) should be pack_length + 1 - cu_seqlens[-1] = len(pack_tokens) - 1 - break - - # Handle any necessary padding - if len(pack_tokens) < pack_length + 1: # +1 here to account for later alignment - pad_len = pack_length + 1 - len(pack_tokens) - extend_with_padding(pack_tokens, pack_targets, pack_positions, pad_len) - # Note len({pack_tokens, pack_targets, pack_positions}) should be pack_length + 1 - cu_seqlens[-1] = len(pack_tokens) - 1 - - assert len(pack_tokens) == pack_length + 1 - assert len(pack_targets) == pack_length + 1 - assert len(pack_positions) == pack_length + 1 - - # Align and convert to tensors - input_ids = torch.tensor(pack_tokens[:-1], dtype=torch.int64) - labels = torch.tensor(pack_targets[1:], dtype=torch.int64) - position_ids = torch.tensor(pack_positions[:-1], dtype=torch.int64) - - # Loss mask. - loss_mask = torch.ones(pack_length, dtype=torch.float32) - loss_mask[labels == pad] = 0.0 # Mask paddings - loss_mask[labels == IGNORE_INDEX] = 0.0 # mask prompts - - # TODO(duncan): Optionally create an attention mask - assert not self.config.create_attention_mask and not self.config.reset_attention_mask - # attention_mask = None - - assert len(cu_seqlens) >= 2 - cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.int32) - # Calculating max_seqlen here, rather than incrementally above, because of possible - # effects of truncation and padding - adjacent_diffs = cu_seqlens[1:] - cu_seqlens[:-1] - max_seqlen = adjacent_diffs.max() # max_seqlen is a 0-D tensor - - # Pad cu_seqlens to a fixed length so that default_collate can - # stack samples with different numbers of documents. Trailing - # entries are filled with pack_length; the merge helper strips - # them later. - padded_cu_seqlens = torch.full( - (pack_length + 1,), pack_length, dtype=torch.int32, - ) - padded_cu_seqlens[:cu_seqlens.numel()] = cu_seqlens - - return { - 'tokens': input_ids, - 'labels': labels, - # 'attention_mask': attention_mask, # PyTorch collate cannot handle NoneType - 'loss_mask': loss_mask, - 'position_ids': position_ids, - 'cu_seqlens': padded_cu_seqlens, - 'max_seqlen': max_seqlen, - } diff --git a/pretrain_gpt.py b/pretrain_gpt.py index 47c1935eb90..1413eb796f4 100644 --- a/pretrain_gpt.py +++ b/pretrain_gpt.py @@ -27,9 +27,10 @@ from megatron.core import mpu from megatron.core.datasets.blended_megatron_dataset_builder import BlendedMegatronDatasetBuilder from megatron.core.datasets.gpt_dataset import GPTDataset, GPTDatasetConfig, MockGPTDataset +from megatron.core.datasets.sft_dataset import IGNORE_INDEX, SFTDataset, SFTDatasetConfig from megatron.core.enums import ModelType -from megatron.core.package_info import __version__ as mcore_version from megatron.core.models.gpt import GPTModel +from megatron.core.package_info import __version__ as mcore_version from megatron.core.packed_seq_params import PackedSeqParams from megatron.core.parallel_state import ( get_context_parallel_group, @@ -37,7 +38,9 @@ ) from megatron.core.rerun_state_machine import get_rerun_state_machine from megatron.core.tokenizers.utils.build_tokenizer import build_tokenizer -from megatron.core.transformer.multi_token_prediction import get_mtp_ranks +from megatron.core.transformer.multi_token_prediction import ( + get_mtp_ranks, +) from megatron.core.transformer.multi_token_prediction import ( mtp_on_this_rank as mtp_on_this_rank_func, ) @@ -49,10 +52,12 @@ get_batch_on_this_tp_rank, get_te_version, get_torch_version, + preprocess_sft_batch, ) from megatron.training import ( get_args, get_timers, + get_tokenizer, inprocess_restart, pretrain, print_rank_0, @@ -61,7 +66,6 @@ from megatron.training.argument_utils import gpt_config_from_args, pretrain_cfg_container_from_args from megatron.training.arguments import core_transformer_config_from_args, parse_and_validate_args from megatron.training.datasets.fim_dataset import GPTFIMDataset, GPTFIMDatasetConfig -from megatron.training.datasets.sft_dataset import SFTDataset from megatron.training.training import update_seqlen_stats_from_cu_seqlens from megatron.training.utils import get_blend_and_blend_per_split, is_first_or_last_pipeline_stage from model_provider import model_provider @@ -101,7 +105,10 @@ def get_batch(data_iterator, vp_stage: Optional[int] = None): config = core_transformer_config_from_args(args) cp_size = args.context_parallel_size + tp_size = args.tensor_model_parallel_size tp_rank = mpu.get_tensor_model_parallel_rank() + sp = args.sequence_parallel + max_seq_len = args.seq_length is_sft = args.sft has_cu_seqlens = is_sft or args.dataloader_inter_document_masking create_attention_mask_in_dataloader = args.create_attention_mask_in_dataloader @@ -123,6 +130,17 @@ def get_batch(data_iterator, vp_stage: Optional[int] = None): batch = {} if tp_rank == 0: batch = next(data_iterator) + if is_sft: + batch = preprocess_sft_batch( + batch, + tp_rank=tp_rank, + cp_size=cp_size, + tp_size=tp_size, + sp=sp, + padding_token_id=get_tokenizer().pad, + padding_label_id=IGNORE_INDEX, + max_seq_len=max_seq_len, + ) for key in BATCH_KEYS: batch[key] = ( batch[key].cuda(non_blocking=True) @@ -438,6 +456,9 @@ def core_gpt_dataset_config_from_args(args: Any) -> GPTDatasetConfig: ) return GPTFIMDatasetConfig(**data_args) + if args.sft: + return SFTDatasetConfig(**data_args) + return GPTDatasetConfig(**data_args) diff --git a/pretrain_hybrid.py b/pretrain_hybrid.py index 39bc7f30b57..e9889e95ca9 100644 --- a/pretrain_hybrid.py +++ b/pretrain_hybrid.py @@ -26,9 +26,10 @@ from megatron.core import mpu from megatron.core.datasets.blended_megatron_dataset_builder import BlendedMegatronDatasetBuilder from megatron.core.datasets.gpt_dataset import GPTDataset, GPTDatasetConfig, MockGPTDataset +from megatron.core.datasets.sft_dataset import IGNORE_INDEX, SFTDataset, SFTDatasetConfig from megatron.core.enums import ModelType -from megatron.core.package_info import __version__ as mcore_version from megatron.core.models.hybrid.hybrid_model import HybridModel +from megatron.core.package_info import __version__ as mcore_version from megatron.core.packed_seq_params import PackedSeqParams from megatron.core.parallel_state import ( get_context_parallel_group, @@ -47,10 +48,12 @@ get_batch_on_this_tp_rank, get_te_version, get_torch_version, + preprocess_sft_batch, ) from megatron.training import ( get_args, get_timers, + get_tokenizer, inprocess_restart, pretrain, print_rank_0, @@ -61,7 +64,6 @@ pretrain_cfg_container_from_args, ) from megatron.training.arguments import core_transformer_config_from_args, parse_and_validate_args -from megatron.training.datasets.sft_dataset import SFTDataset from megatron.training.training import update_seqlen_stats_from_cu_seqlens from megatron.training.utils import get_blend_and_blend_per_split, is_first_or_last_pipeline_stage from model_provider import model_provider @@ -101,7 +103,10 @@ def get_batch(data_iterator, vp_stage=None): config = core_transformer_config_from_args(args) cp_size = args.context_parallel_size + tp_size = args.tensor_model_parallel_size tp_rank = mpu.get_tensor_model_parallel_rank() + sp = args.sequence_parallel + max_seq_len = args.seq_length is_sft = args.sft has_cu_seqlens = is_sft or args.dataloader_inter_document_masking create_attention_mask_in_dataloader = args.create_attention_mask_in_dataloader @@ -123,6 +128,17 @@ def get_batch(data_iterator, vp_stage=None): batch = {} if tp_rank == 0: batch = next(data_iterator) + if is_sft: + batch = preprocess_sft_batch( + batch, + tp_rank=tp_rank, + cp_size=cp_size, + tp_size=tp_size, + sp=sp, + padding_token_id=get_tokenizer().pad, + padding_label_id=IGNORE_INDEX, + max_seq_len=max_seq_len, + ) for key in BATCH_KEYS: batch[key] = ( batch[key].cuda(non_blocking=True) @@ -375,34 +391,39 @@ def core_gpt_dataset_config_from_args(args: Any) -> GPTDatasetConfig: with open(args.per_dataset_sequences_path, "r") as f: sequences_per_dataset = json.load(f) - return GPTDatasetConfig( - random_seed=args.seed, - sequence_length=args.seq_length, - blend=blend, - blend_per_split=blend_per_split, - split=args.split, - multiple_validation_sets=args.multiple_validation_sets, - full_validation=args.full_validation, - num_dataset_builder_threads=args.num_dataset_builder_threads, - path_to_cache=args.data_cache_path, - mmap_bin_files=args.mmap_bin_files, - tokenizer=tokenizer, - reset_position_ids=args.reset_position_ids, - reset_attention_mask=args.reset_attention_mask, - eod_mask_loss=args.eod_mask_loss, - create_attention_mask=args.create_attention_mask_in_dataloader, - object_storage_cache_path=args.object_storage_cache_path, - mid_level_dataset_surplus=args.mid_level_dataset_surplus, - allow_ambiguous_pad_tokens=args.allow_ambiguous_pad_tokens, - fast_cache_load=args.dataloader_fast_cache_load, - sequences_per_dataset=sequences_per_dataset, - defer_npy_index_mmap=args.dataloader_defer_npy_index_mmap, - context_parallel_size=args.context_parallel_size, - data_parallel_size=args.data_parallel_size, - sequence_parallel_size=args.tensor_model_parallel_size * args.sequence_parallel, - hybrid_context_parallel=args.hybrid_context_parallel, - inter_document_masking=args.dataloader_inter_document_masking, - ) + data_args = { + "random_seed": args.seed, + "sequence_length": args.seq_length, + "blend": blend, + "blend_per_split": blend_per_split, + "split": args.split, + "multiple_validation_sets": args.multiple_validation_sets, + "full_validation": args.full_validation, + "num_dataset_builder_threads": args.num_dataset_builder_threads, + "path_to_cache": args.data_cache_path, + "mmap_bin_files": args.mmap_bin_files, + "tokenizer": tokenizer, + "reset_position_ids": args.reset_position_ids, + "reset_attention_mask": args.reset_attention_mask, + "eod_mask_loss": args.eod_mask_loss, + "create_attention_mask": args.create_attention_mask_in_dataloader, + "object_storage_cache_path": args.object_storage_cache_path, + "mid_level_dataset_surplus": args.mid_level_dataset_surplus, + "allow_ambiguous_pad_tokens": args.allow_ambiguous_pad_tokens, + "fast_cache_load": args.dataloader_fast_cache_load, + "sequences_per_dataset": sequences_per_dataset, + "defer_npy_index_mmap": args.dataloader_defer_npy_index_mmap, + "context_parallel_size": args.context_parallel_size, + "data_parallel_size": args.data_parallel_size, + "sequence_parallel_size": args.tensor_model_parallel_size * args.sequence_parallel, + "hybrid_context_parallel": args.hybrid_context_parallel, + "inter_document_masking": args.dataloader_inter_document_masking, + } + + if args.sft: + return SFTDatasetConfig(**data_args) + + return GPTDatasetConfig(**data_args) def train_valid_test_datasets_provider(train_val_test_num_samples, vp_stage=None): diff --git a/tests/unit_tests/data/test_cp_utils.py b/tests/unit_tests/data/test_cp_utils.py new file mode 100644 index 00000000000..f204c6ebae4 --- /dev/null +++ b/tests/unit_tests/data/test_cp_utils.py @@ -0,0 +1,232 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Unit tests for context-parallel data utilities in megatron.core.utils.""" + +import pytest +import torch + +from megatron.core.utils import pad_thd_sequences_for_cp + + +def _reference_pad(tensors_with_pad_values, cu_seqlens, divisibility_factor): + """Slow but obviously-correct reference implementation.""" + out_tensors = [[] for _ in tensors_with_pad_values] + cu_padded = [0] + for start, end in zip(cu_seqlens[:-1].tolist(), cu_seqlens[1:].tolist()): + seg_len = end - start + padded_len = ( + (seg_len + divisibility_factor - 1) // divisibility_factor + ) * divisibility_factor + pad = padded_len - seg_len + for i, (tensor, pad_value) in enumerate(tensors_with_pad_values): + out_tensors[i].append(tensor[start:end]) + if pad > 0: + out_tensors[i].append(torch.full((pad,), pad_value, dtype=tensor.dtype)) + cu_padded.append(cu_padded[-1] + padded_len) + return ( + [torch.cat(pieces) for pieces in out_tensors], + torch.tensor(cu_padded, dtype=cu_seqlens.dtype), + ) + + +class TestPadThdSequencesForCp: + """Unit tests for ``pad_thd_sequences_for_cp``.""" + + def test_shorter_than_divisibility_factor(self): + """All segments shorter than the divisibility factor (matches TE unit test).""" + input_ids = torch.tensor([1, 1, 1, 2, 2, 3, 3, 3, 3]) + labels = torch.tensor([-100, -100, -100, -100, -100, -100, -100, 13, -100]) + cu_seqlens = torch.tensor([0, 3, 5, 9], dtype=torch.int32) + + (ids_p, lab_p), cu_p = pad_thd_sequences_for_cp( + [(input_ids, 777), (labels, -200)], cu_seqlens, divisibility_factor=8 + ) + (ref_ids, ref_lab), ref_cu = _reference_pad( + [(input_ids, 777), (labels, -200)], cu_seqlens, 8 + ) + assert torch.equal(ids_p, ref_ids) + assert torch.equal(lab_p, ref_lab) + assert torch.equal(cu_p, ref_cu) + assert cu_p[-1].item() == 24 + + def test_mixed_sequence_lengths(self): + """Segments mixing lengths shorter and longer than the divisibility factor.""" + input_ids = torch.tensor( + [1, 1, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4] + ) + labels = torch.arange(input_ids.numel(), dtype=torch.int64) + cu_seqlens = torch.tensor([0, 2, 9, 13, 23], dtype=torch.int32) + + (ids_p, lab_p), cu_p = pad_thd_sequences_for_cp( + [(input_ids, 999), (labels, -300)], cu_seqlens, divisibility_factor=6 + ) + (ref_ids, ref_lab), ref_cu = _reference_pad( + [(input_ids, 999), (labels, -300)], cu_seqlens, 6 + ) + assert torch.equal(ids_p, ref_ids) + assert torch.equal(lab_p, ref_lab) + assert torch.equal(cu_p, ref_cu) + # Per-segment padded lengths: 6, 12, 6, 12. + assert cu_p.tolist() == [0, 6, 18, 24, 36] + + def test_longer_than_divisibility_factor(self): + """Segments longer than the divisibility factor are rounded up to the next multiple.""" + # Seg 1: 7 -> 8 (pad 1); Seg 2: 11 -> 12 (pad 1); Seg 3: 5 -> 8 (pad 3). + input_ids = torch.cat( + [ + torch.ones(7, dtype=torch.int64), + torch.full((11,), 2, dtype=torch.int64), + torch.full((5,), 3, dtype=torch.int64), + ] + ) + labels = torch.arange(input_ids.numel(), dtype=torch.int64) + 100 + cu_seqlens = torch.tensor([0, 7, 18, 23], dtype=torch.int32) + + (ids_p, lab_p), cu_p = pad_thd_sequences_for_cp( + [(input_ids, 888), (labels, -400)], cu_seqlens, divisibility_factor=4 + ) + (ref_ids, ref_lab), ref_cu = _reference_pad( + [(input_ids, 888), (labels, -400)], cu_seqlens, 4 + ) + assert torch.equal(ids_p, ref_ids) + assert torch.equal(lab_p, ref_lab) + assert torch.equal(cu_p, ref_cu) + assert cu_p.tolist() == [0, 8, 20, 28] + + def test_already_divisible_is_noop(self): + """When every segment already satisfies divisibility, output equals input.""" + input_ids = torch.arange(16, dtype=torch.int64) + labels = torch.arange(16, dtype=torch.int64) + 100 + cu_seqlens = torch.tensor([0, 8, 16], dtype=torch.int32) + + (ids_p, lab_p), cu_p = pad_thd_sequences_for_cp( + [(input_ids, 0), (labels, -100)], cu_seqlens, divisibility_factor=4 + ) + assert torch.equal(ids_p, input_ids) + assert torch.equal(lab_p, labels) + assert torch.equal(cu_p, cu_seqlens) + + def test_accepts_2d_input(self): + """A leading batch dim of 1 (from DataLoader collation) is squeezed away.""" + input_ids = torch.tensor([[1, 2, 3, 4, 5]]) + labels = torch.tensor([[10, 20, 30, 40, 50]]) + cu_seqlens = torch.tensor([0, 2, 5], dtype=torch.int32) + + (ids_p, lab_p), cu_p = pad_thd_sequences_for_cp( + [(input_ids, 0), (labels, -100)], cu_seqlens, divisibility_factor=4 + ) + # Seg 1: [1, 2] + 2 pads = 4; Seg 2: [3, 4, 5] + 1 pad = 4. + assert ids_p.dim() == 1 + assert ids_p.tolist() == [1, 2, 0, 0, 3, 4, 5, 0] + assert lab_p.tolist() == [10, 20, -100, -100, 30, 40, 50, -100] + assert cu_p.tolist() == [0, 4, 8] + + def test_preserves_cu_seqlens_int32_dtype(self): + """``torch.cumsum`` promotes int32 -> int64 by default; verify we preserve int32.""" + input_ids = torch.tensor([1, 2, 3]) + cu_seqlens = torch.tensor([0, 3], dtype=torch.int32) + + _, cu_p = pad_thd_sequences_for_cp([(input_ids, 0)], cu_seqlens, divisibility_factor=8) + assert cu_p.dtype == torch.int32 + + def test_preserves_cu_seqlens_int64_dtype(self): + """Dtype preservation also holds when cu_seqlens is int64.""" + input_ids = torch.tensor([1, 2, 3]) + cu_seqlens = torch.tensor([0, 3], dtype=torch.int64) + + _, cu_p = pad_thd_sequences_for_cp([(input_ids, 0)], cu_seqlens, divisibility_factor=8) + assert cu_p.dtype == torch.int64 + + def test_preserves_tensor_dtypes(self): + """Per-tensor dtypes propagate to the padded outputs.""" + input_ids = torch.tensor([1, 2, 3], dtype=torch.int32) + labels = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float32) + loss_mask = torch.tensor([1.0, 1.0, 1.0], dtype=torch.float32) + cu_seqlens = torch.tensor([0, 3], dtype=torch.int32) + + (ids_p, lab_p, lm_p), _ = pad_thd_sequences_for_cp( + [(input_ids, 0), (labels, 0), (loss_mask, 0)], cu_seqlens, divisibility_factor=4 + ) + assert ids_p.dtype == torch.int32 + assert lab_p.dtype == torch.float32 + assert lm_p.dtype == torch.float32 + + def test_pads_multiple_tensors_in_one_call(self): + """Passing N tensors yields N padded outputs sharing one cu_seqlens_padded.""" + input_ids = torch.tensor([10, 11, 20, 21, 22]) + labels = torch.tensor([100, 101, 200, 201, 202]) + loss_mask = torch.tensor([1, 1, 1, 1, 1], dtype=torch.float32) + cu_seqlens = torch.tensor([0, 2, 5], dtype=torch.int32) + + (ids_p, lab_p, lm_p), cu_p = pad_thd_sequences_for_cp( + [(input_ids, 0), (labels, -100), (loss_mask, 0.0)], cu_seqlens, divisibility_factor=4 + ) + # Seg 1: 2 -> 4 (2 pads); Seg 2: 3 -> 4 (1 pad). + assert ids_p.tolist() == [10, 11, 0, 0, 20, 21, 22, 0] + assert lab_p.tolist() == [100, 101, -100, -100, 200, 201, 202, -100] + assert lm_p.tolist() == [1.0, 1.0, 0.0, 0.0, 1.0, 1.0, 1.0, 0.0] + assert cu_p.tolist() == [0, 4, 8] + + def test_single_tensor(self): + """A single (tensor, pad_value) pair is also valid input.""" + input_ids = torch.tensor([7, 7, 7]) + cu_seqlens = torch.tensor([0, 3], dtype=torch.int32) + + (ids_p,), cu_p = pad_thd_sequences_for_cp( + [(input_ids, 99)], cu_seqlens, divisibility_factor=4 + ) + assert ids_p.tolist() == [7, 7, 7, 99] + assert cu_p.tolist() == [0, 4] + + def test_divisibility_factor_one_is_noop(self): + """divisibility_factor=1 means every length is already valid -> no padding.""" + input_ids = torch.tensor([1, 2, 3, 4, 5]) + labels = torch.tensor([6, 7, 8, 9, 10]) + cu_seqlens = torch.tensor([0, 2, 5], dtype=torch.int32) + + (ids_p, lab_p), cu_p = pad_thd_sequences_for_cp( + [(input_ids, 0), (labels, -100)], cu_seqlens, divisibility_factor=1 + ) + assert torch.equal(ids_p, input_ids) + assert torch.equal(lab_p, labels) + assert torch.equal(cu_p, cu_seqlens) + + @pytest.mark.parametrize("divisibility_factor", [2, 4, 8, 16]) + def test_post_condition_segments_divisible(self, divisibility_factor): + """The padded cu_seqlens differences are all multiples of divisibility_factor.""" + torch.manual_seed(0) + seg_lens = [int(x) for x in torch.randint(1, 50, (8,))] + cu_seqlens = torch.tensor( + [0] + torch.tensor(seg_lens).cumsum(0).tolist(), dtype=torch.int32 + ) + total = cu_seqlens[-1].item() + input_ids = torch.arange(total, dtype=torch.int64) + labels = torch.arange(total, dtype=torch.int64) + + _, cu_p = pad_thd_sequences_for_cp( + [(input_ids, 0), (labels, -100)], cu_seqlens, divisibility_factor=divisibility_factor + ) + diffs = (cu_p[1:] - cu_p[:-1]).tolist() + assert all(d % divisibility_factor == 0 for d in diffs), diffs + + @pytest.mark.parametrize("divisibility_factor", [2, 4, 8]) + def test_against_reference(self, divisibility_factor): + """Random shapes against the slow reference implementation.""" + torch.manual_seed(42) + seg_lens = [int(x) for x in torch.randint(1, 30, (6,))] + cu_seqlens = torch.tensor( + [0] + torch.tensor(seg_lens).cumsum(0).tolist(), dtype=torch.int32 + ) + total = cu_seqlens[-1].item() + input_ids = torch.randint(0, 1000, (total,), dtype=torch.int64) + labels = torch.randint(-100, 1000, (total,), dtype=torch.int64) + loss_mask = torch.ones(total, dtype=torch.float32) + spec = [(input_ids, 42), (labels, -100), (loss_mask, 0.0)] + + padded_tensors, cu_p = pad_thd_sequences_for_cp( + spec, cu_seqlens, divisibility_factor=divisibility_factor + ) + ref_tensors, ref_cu = _reference_pad(spec, cu_seqlens, divisibility_factor) + for got, want in zip(padded_tensors, ref_tensors): + assert torch.equal(got, want) + assert torch.equal(cu_p, ref_cu) diff --git a/tests/unit_tests/data/test_get_batch.py b/tests/unit_tests/data/test_get_batch.py index 3d96caee968..bb8ad35176a 100644 --- a/tests/unit_tests/data/test_get_batch.py +++ b/tests/unit_tests/data/test_get_batch.py @@ -71,11 +71,14 @@ def initialize_test_environment( def create_sft_data_iterator(max_seq_length: int = 1024): - """Create a mock SFT data iterator matching the old SFTDataset output after DataLoader collation. - - The old SFTDataset (megatron/training/datasets/sft_dataset.py) returns per-sample dicts with - keys: tokens, labels, loss_mask, position_ids, cu_seqlens, max_seqlen — all padded to - seq_length. After PyTorch DataLoader default_collate, tensors get a leading batch dim of 1. + """Create a mock SFT data iterator matching the new SFTDataset output after DataLoader collation. + + The new SFTDataset (megatron/core/datasets/sft_dataset.py) returns per-sample dicts with + keys: tokens, labels, loss_mask, cu_seqlens — un-padded (length == sum of segment lengths, + cu_seqlens last entry == total real tokens). After PyTorch DataLoader default_collate, the + tensors get a leading batch dim of 1. ``preprocess_sft_batch`` (invoked inside get_batch) + is then responsible for padding to ``max_seq_length``, computing position_ids, and + appending the padding segment to cu_seqlens. """ min_len = max(1, int(0.1 * max_seq_length)) max_len = max(2, int(0.4 * max_seq_length)) @@ -97,58 +100,25 @@ def create_sft_data_iterator(max_seq_length: int = 1024): # Generate packed token sequence (num_real_tokens + 1 for labels shift) text = torch.randint(0, 10000, (num_real_tokens + 1,), dtype=torch.int64) - # Pad to max_seq_length (mimics old SFTDataset padding) - pad_len = max_seq_length - num_real_tokens - pad_token = 0 - - tokens = torch.cat([text[:-1], torch.full((pad_len,), pad_token, dtype=torch.int64)]) - labels = torch.cat([text[1:], torch.full((pad_len,), pad_token, dtype=torch.int64)]) + tokens = text[:-1].contiguous() + labels = text[1:].contiguous() + loss_mask = torch.ones(num_real_tokens, dtype=torch.float32) - # Position IDs: per-segment positions, then padding positions - position_ids = torch.cat([torch.arange(l, dtype=torch.int64) for l in lengths]) - position_ids = torch.cat( - [ - position_ids, - torch.arange( - position_ids[-1].item() + 1, - position_ids[-1].item() + 1 + pad_len, - dtype=torch.int64, - ), - ] - ) - - # Loss mask: 1 for real tokens, 0 for padding - loss_mask = torch.cat( - [ - torch.ones(num_real_tokens, dtype=torch.float32), - torch.zeros(pad_len, dtype=torch.float32), - ] - ) - - # cu_seqlens: cumulative lengths ending at max_seq_length (last entry = seq_length after padding) + # cu_seqlens ends at total real tokens — preprocess_sft_batch appends the padding + # segment when it pads/truncates to max_seq_length. cu_seqlens = torch.cat( ( torch.zeros(1, dtype=torch.int32), torch.cumsum(torch.tensor(lengths, dtype=torch.int64), dim=0).to(torch.int32), ) ) - cu_seqlens[-1] = max_seq_length # last entry is padded to seq_length - - # max_seqlen: max segment length - seg_lengths = cu_seqlens[1:] - cu_seqlens[:-1] - max_seqlen = torch.tensor([seg_lengths.max().item()], dtype=torch.int32) # Add batch dimension to all per-sample tensors to mimic DataLoader default_collate. - # The dataset emits cu_seqlens as 1-D (S+1,) and max_seqlen as 0-D; default_collate - # stacks them with a leading batch dim of 1. get_batch_on_this_tp_rank's sender is - # responsible for squeezing the batch dim of cu_seqlens before broadcast. batch = { "tokens": tokens.unsqueeze(0), "labels": labels.unsqueeze(0), "loss_mask": loss_mask.unsqueeze(0), - "position_ids": position_ids.unsqueeze(0), "cu_seqlens": cu_seqlens.unsqueeze(0), - "max_seqlen": max_seqlen, } return iter([batch]), num_real_tokens @@ -207,7 +177,15 @@ def test_sft_batch(tp_size, pp_size, cp_size, seq_length): assert attention_mask is None assert hybrid_cp_group is None assert local_cp_size is None - assert cu_seqlens_padded is None + if cp_size > 1: + assert cu_seqlens_padded is not None + assert cu_seqlens_padded.dim() == 2 + assert cu_seqlens_padded.shape[0] == 1 + assert cu_seqlens_padded.dtype == torch.int32 + assert cu_seqlens_padded[0, 0].item() == 0 + assert cu_seqlens_padded[0, -1].item() == seq_length + else: + assert cu_seqlens_padded is None assert tokens.shape == ( 1, @@ -255,7 +233,15 @@ def test_sft_batch(tp_size, pp_size, cp_size, seq_length): assert attention_mask is None assert hybrid_cp_group is None assert local_cp_size is None - assert cu_seqlens_padded is None + if cp_size > 1: + assert cu_seqlens_padded is not None + assert cu_seqlens_padded.dim() == 2 + assert cu_seqlens_padded.shape[0] == 1 + assert cu_seqlens_padded.dtype == torch.int32 + assert cu_seqlens_padded[0, 0].item() == 0 + assert cu_seqlens_padded[0, -1].item() == seq_length + else: + assert cu_seqlens_padded is None assert tokens.shape == ( 1, @@ -291,7 +277,15 @@ def test_sft_batch(tp_size, pp_size, cp_size, seq_length): assert attention_mask is None assert hybrid_cp_group is None assert local_cp_size is None - assert cu_seqlens_padded is None + if cp_size > 1: + assert cu_seqlens_padded is not None + assert cu_seqlens_padded.dim() == 2 + assert cu_seqlens_padded.shape[0] == 1 + assert cu_seqlens_padded.dtype == torch.int32 + assert cu_seqlens_padded[0, 0].item() == 0 + assert cu_seqlens_padded[0, -1].item() == seq_length + else: + assert cu_seqlens_padded is None assert labels.shape == ( 1, @@ -330,7 +324,15 @@ def test_sft_batch(tp_size, pp_size, cp_size, seq_length): assert cu_seqlens is not None assert max_seqlen is not None - assert cu_seqlens_padded is None + if cp_size > 1: + assert cu_seqlens_padded is not None + assert cu_seqlens_padded.dim() == 2 + assert cu_seqlens_padded.shape[0] == 1 + assert cu_seqlens_padded.dtype == torch.int32 + assert cu_seqlens_padded[0, 0].item() == 0 + assert cu_seqlens_padded[0, -1].item() == seq_length + else: + assert cu_seqlens_padded is None assert cu_seqlens.dim() == 2 assert cu_seqlens.shape[0] == 1 @@ -543,7 +545,7 @@ def test_inter_document_masking_batch(tp_size, pp_size, cp_size, seq_length): data_iterator = None if mpu.get_tensor_model_parallel_rank() == 0: - data_iterator, _ = create_sft_data_iterator(seq_length) + data_iterator = create_inter_document_masking_data_iterator(seq_length, micro_batch_size=1) ( attention_mask, @@ -770,6 +772,7 @@ def test_pretrain_batch( assert loss_mask is not None assert position_ids is not None assert cu_seqlens is None + assert cu_seqlens_padded is None assert max_seqlen is None assert hybrid_cp_group is None @@ -1061,7 +1064,15 @@ def test_hybrid_cp_batch(tp_size, cp_size, seq_length, create_attention_mask): assert cu_seqlens_padded[0, -1].item() == total_seq_len assert hybrid_cp_group is not None else: - assert cu_seqlens_padded is None + if cp_size > 1: + assert cu_seqlens_padded is not None + assert cu_seqlens_padded.dim() == 2 + assert cu_seqlens_padded.shape[0] == 1 + assert cu_seqlens_padded.dtype == torch.int32 + assert cu_seqlens_padded[0, 0].item() == 0 + assert cu_seqlens_padded[0, -1].item() == seq_length + else: + assert cu_seqlens_padded is None assert hybrid_cp_group is None Utils.destroy_model_parallel() diff --git a/tests/unit_tests/data/test_sft_dataset.py b/tests/unit_tests/data/test_sft_dataset.py new file mode 100644 index 00000000000..65179dcf554 --- /dev/null +++ b/tests/unit_tests/data/test_sft_dataset.py @@ -0,0 +1,399 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +## +# Compile megatron.core.datasets.helpers_cpp dependencies before BlendedDataset import +## + +import os +import random +from argparse import Namespace +from dataclasses import dataclass + +import pytest +import torch + +from megatron.core.datasets.blended_megatron_dataset_builder import BlendedMegatronDatasetBuilder +from megatron.core.datasets.indexed_dataset import DType, IndexedDatasetBuilder +from megatron.core.datasets.sft_dataset import ( + ChatTemplateConfig, + SFTDataset, + SFTDatasetConfig, + extract_segments, +) +from megatron.core.datasets.utils import compile_helpers +from megatron.core.tokenizers.utils.build_tokenizer import build_tokenizer +from megatron.training.utils import get_blend_and_blend_per_split +from tests.unit_tests.dist_checkpointing import TempNamedDir +from tests.unit_tests.test_utilities import Utils + + +@dataclass +class TestChatTemplateConfig(ChatTemplateConfig): + """Chat template config compatible with NullSFTTokenizer's character-level encoding. + + Marker strings use ASCII characters (token IDs < 128). + Think markers must be single characters (tokenize(str)[0] returns a single ID). + Tool call markers end with '\\n' as required by _split_tool_calls. + """ + + system_start_str: str = "[S]" + user_start_str: str = "[U]" + assistant_start_str: str = "[A]" + end_str: str = "[E]" + think_start_str: str = "{" + think_end_str: str = "}" + tool_call_start_str: str = "~\n" + tool_call_end_str: str = "`\n" + tool_response_start_str: str = "^\n" + tool_response_end_str: str = "|\n" + + +# Safe range for content tokens: [200, 1000) avoids collisions with ASCII marker tokens (< 128) +CONTENT_TOKEN_MIN = 200 +CONTENT_TOKEN_MAX = 1000 + + +def _random_content_tokens(min_len=5, max_len=20): + """Generate random content token IDs in a safe range that won't collide with markers.""" + length = random.randint(min_len, max_len) + return [random.randint(CONTENT_TOKEN_MIN, CONTENT_TOKEN_MAX - 1) for _ in range(length)] + + +def build_tokenized_conversation( + tokenizer, chat_template_config, conversation_type, with_system=True +): + """Build a pre-tokenized conversation as a flat list of token IDs. + + Args: + tokenizer: Tokenizer with tokenize() method. + chat_template_config: TestChatTemplateConfig instance. + conversation_type: One of 'simple', 'with_thinking', 'with_tool_calls', + 'with_thinking_and_tool_calls'. + with_system: Whether to include a system message. + + Returns: + List[int]: The tokenized conversation. + """ + cfg = chat_template_config + + sys_start = tokenizer.tokenize(cfg.system_start_str, add_special_tokens=False) + usr_start = tokenizer.tokenize(cfg.user_start_str, add_special_tokens=False) + ast_start = tokenizer.tokenize(cfg.assistant_start_str, add_special_tokens=False) + end = tokenizer.tokenize(cfg.end_str, add_special_tokens=False) + think_start = tokenizer.tokenize(cfg.think_start_str, add_special_tokens=False) + think_end = tokenizer.tokenize(cfg.think_end_str, add_special_tokens=False) + tc_start = tokenizer.tokenize(cfg.tool_call_start_str, add_special_tokens=False) + tc_end = tokenizer.tokenize(cfg.tool_call_end_str, add_special_tokens=False) + tr_start = tokenizer.tokenize(cfg.tool_response_start_str, add_special_tokens=False) + tr_end = tokenizer.tokenize(cfg.tool_response_end_str, add_special_tokens=False) + + tokens = [] + + # Optional system message + if with_system: + tokens += sys_start + _random_content_tokens() + end + + if conversation_type == "simple": + # user + assistant + tokens += usr_start + _random_content_tokens() + end + tokens += ast_start + _random_content_tokens() + end + + elif conversation_type == "with_thinking": + # user + assistant(think + response) + tokens += usr_start + _random_content_tokens() + end + tokens += ( + ast_start + + think_start + + _random_content_tokens() + + think_end + + _random_content_tokens() + + end + ) + + elif conversation_type == "with_tool_calls": + # user + assistant(tool_call) + user(tool_response) + assistant(response) + tokens += usr_start + _random_content_tokens() + end + tokens += ast_start + tc_start + _random_content_tokens() + tc_end + end + tokens += usr_start + tr_start + _random_content_tokens() + tr_end + end + tokens += ast_start + _random_content_tokens() + end + + elif conversation_type == "with_thinking_and_tool_calls": + # user + assistant(think + tool_call) + user(tool_response) + assistant(response) + tokens += usr_start + _random_content_tokens() + end + tokens += ( + ast_start + + think_start + + _random_content_tokens() + + think_end + + tc_start + + _random_content_tokens() + + tc_end + + end + ) + tokens += usr_start + tr_start + _random_content_tokens() + tr_end + end + tokens += ast_start + _random_content_tokens() + end + + return tokens + + +CONVERSATION_TYPES = ["simple", "with_thinking", "with_tool_calls", "with_thinking_and_tool_calls"] + + +def create_file_prefixes( + tokenizer, chat_template_config, number_of_files, max_conversations_per_file, dataset_dir +): + """Create indexed dataset files with pre-tokenized conversations of all types.""" + os.makedirs(dataset_dir, exist_ok=True) + + file_prefixes = [] + for i in range(number_of_files): + file_prefix_path = os.path.join(dataset_dir, f"file_{i}") + builder = IndexedDatasetBuilder( + file_prefix_path + ".bin", dtype=DType.optimal_dtype(tokenizer.vocab_size) + ) + for _ in range(random.randint(10, max_conversations_per_file)): + conv_type = random.choice(CONVERSATION_TYPES) + with_system = random.choice([True, False]) + tokens = build_tokenized_conversation( + tokenizer, chat_template_config, conv_type, with_system + ) + builder.add_document(tokens, [len(tokens)]) + builder.finalize(file_prefix_path + ".idx") + file_prefixes.append(file_prefix_path) + + return file_prefixes + + +def verify_loss_mask(sample, config): + """Verify the loss mask of a packed sample against the expected pattern. + + For each document in the packed sample (split by cu_seqlens), runs + extract_segments and checks that loss_mask values match the flag settings. + """ + tokens = sample['tokens'] + labels = sample['labels'] + loss_mask = sample['loss_mask'] + cu_seqlens = sample['cu_seqlens'] + + # The __getitem__ returns tokens[:-1], labels=tokens[1:], loss_mask[:-1] + # To reconstruct the full token sequence per document, we use cu_seqlens. + # cu_seqlens gives cumulative lengths of the original (unshifted) documents. + # The total packed length is cu_seqlens[-1], and output length is cu_seqlens[-1] - 1. + + num_docs = len(cu_seqlens) - 1 + # Reconstruct the full token sequence (before the [:-1] / [1:] split) + full_tokens = torch.cat([tokens[:1], labels]) # length = cu_seqlens[-1] + # Reconstruct the full loss mask (before the [:-1] slice) + # loss_mask has length cu_seqlens[-1] - 1, original had length cu_seqlens[-1] + # We don't know the last element, but we don't need it for per-doc verification + + offset = 0 + for doc_idx in range(num_docs): + doc_len = (cu_seqlens[doc_idx + 1] - cu_seqlens[doc_idx]).item() + doc_tokens = full_tokens[offset : offset + doc_len].tolist() + + # loss_mask covers positions [0, cu_seqlens[-1] - 1) of the full sequence + # For this document, loss_mask positions are [offset, offset + doc_len) + # but the last position of the entire packed sequence has no loss_mask entry + doc_loss_mask_end = min(offset + doc_len, len(loss_mask)) + doc_loss_mask = loss_mask[offset:doc_loss_mask_end] + + if not config.train_on_assistant_responses_only: + # All tokens should be trained on + assert ( + doc_loss_mask == 1.0 + ).all(), f"Doc {doc_idx}: expected all loss_mask=1 when train_on_assistant_responses_only=False" + else: + segments = extract_segments( + doc_tokens, + config.role_start_tokens, + config.end_tokens, + config.think_start_tokens, + config.think_end_tokens, + config.tool_call_start_tokens, + config.tool_call_end_tokens, + config.tool_response_start_tokens, + ) + + for seg in segments: + # Segment positions are relative to doc_tokens, but loss_mask is + # indexed from the start of the packed sequence. Adjust with offset. + seg_start = seg["start"] + seg_end = min(seg["end"], doc_loss_mask_end - offset) + if seg_start >= seg_end: + continue + + seg_mask = doc_loss_mask[seg_start:seg_end] + role = seg["role"] + + if role == "assistant": + expected = 1.0 + elif role == "reasoning": + expected = 1.0 if config.train_on_thinking_traces else 0.0 + elif role == "tool_call": + expected = 1.0 if config.train_on_tool_calls else 0.0 + else: + # system, user, tool_response — always masked + expected = 0.0 + + assert (seg_mask == expected).all(), ( + f"Doc {doc_idx}, segment role='{role}' [{seg_start}:{seg_end}]: " + f"expected loss_mask={expected}, got {seg_mask.tolist()}" + ) + + offset += doc_len + + +@pytest.mark.parametrize("vocab_size", [131072, 20000]) +@pytest.mark.parametrize( + "train_on_assistant_responses_only,train_on_thinking_traces,train_on_tool_calls", + [ + (False, False, False), + (True, False, False), + (True, True, False), + (True, False, True), + (True, True, True), + ], +) +def test_sft_dataset( + vocab_size, + train_on_assistant_responses_only, + train_on_thinking_traces, + train_on_tool_calls, + tmp_path_dist_ckpt, + sequence_length: int = 1500, + number_of_files: int = 10, + max_conversations_per_file: int = 20, +): + if torch.distributed.is_available(): + Utils.initialize_distributed() + if torch.distributed.get_rank() == 0: + compile_helpers() + torch.distributed.barrier() + else: + compile_helpers() + + tokenizer = build_tokenizer( + Namespace( + vocab_size=vocab_size, + tokenizer_type="NullTokenizer", + rank=0, + make_vocab_size_divisible_by=128, + tensor_model_parallel_size=1, + pad_vocab_size=False, + ) + ) + + chat_template_config = TestChatTemplateConfig() + + with TempNamedDir(tmp_path_dist_ckpt / "test_fast_builder", sync=True) as temp_dir: + if not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0: + file_prefixes = create_file_prefixes( + tokenizer, + chat_template_config, + number_of_files, + max_conversations_per_file, + os.path.join(temp_dir, "dataset"), + ) + else: + file_prefixes = [] + for i in range(number_of_files): + file_prefix_path = os.path.join(temp_dir, "dataset", f"file_{i}") + file_prefixes.append(file_prefix_path) + + if torch.distributed.is_initialized(): + torch.distributed.barrier() + + random.seed(1234) + + data_cache_path = os.path.join(temp_dir, "cache") + + args = Namespace( + seed=1234, + seq_length=sequence_length, + data_cache_path=data_cache_path, + split=None, + data_path=None, + train_data_path=file_prefixes[0:6], + valid_data_path=file_prefixes[6:9], + test_data_path=file_prefixes[9:10], + per_split_data_args_path=None, + data_args_path=None, + ) + + blend, blend_per_split = get_blend_and_blend_per_split(args) + + data_args = { + "random_seed": args.seed, + "sequence_length": args.seq_length, + "blend": blend, + "blend_per_split": blend_per_split, + "split": args.split, + "path_to_cache": args.data_cache_path, + "tokenizer": tokenizer, + "reset_position_ids": False, + "reset_attention_mask": False, + "eod_mask_loss": False, + "create_attention_mask": False, + "train_on_assistant_responses_only": train_on_assistant_responses_only, + "train_on_thinking_traces": train_on_thinking_traces, + "train_on_tool_calls": train_on_tool_calls, + "chat_template_config": chat_template_config, + } + config = SFTDatasetConfig(**data_args) + + train_ds, valid_ds, test_ds = BlendedMegatronDatasetBuilder( + SFTDataset, [100, 10, 10], lambda: True, config + ).build() + + # Shape invariant checks + loss mask verification + for sample_idx in [0, 1, -1]: + sample = train_ds[sample_idx] + tokens = sample['tokens'] + labels = sample['labels'] + loss_mask = sample['loss_mask'] + cu_seqlens = sample['cu_seqlens'] + + # Dtype checks + assert tokens.dtype == torch.int64 + assert labels.dtype == torch.int64 + assert loss_mask.dtype == torch.float32 + assert cu_seqlens.dtype == torch.int32 + + # Shape consistency: tokens, labels, loss_mask must have the same length + assert tokens.shape == labels.shape == loss_mask.shape + + # Packed length must not exceed sequence_length + assert tokens.shape[0] <= sequence_length + + # cu_seqlens[-1] == total tokens before the [:-1]/[1:] shift + assert tokens.shape[0] + 1 == cu_seqlens[-1] + + # Labels are correctly shifted: labels[i] == tokens[i+1] in the original sequence + # Reconstruct full sequence and verify + full_tokens = torch.cat([tokens[:1], labels]) + assert (full_tokens[:-1] == tokens).all(), "tokens should be full_tokens[:-1]" + assert (full_tokens[1:] == labels).all(), "labels should be full_tokens[1:]" + + # Token values are in valid vocab range + assert (tokens >= 0).all() and (tokens < vocab_size + 1).all() + assert (labels >= 0).all() and (labels < vocab_size + 1).all() + + # cu_seqlens invariants: starts at 0, monotonically increasing + assert cu_seqlens[0] == 0 + doc_lengths = cu_seqlens[1:] - cu_seqlens[:-1] + assert (doc_lengths > 0).all(), "Each document must have positive length" + assert ( + doc_lengths <= sequence_length + ).all(), "No document should exceed sequence_length" + + # Loss mask is binary (only 0.0 or 1.0) + assert ((loss_mask == 0.0) | (loss_mask == 1.0)).all(), "Loss mask must be binary" + + # Loss mask is non-trivial when train_on_assistant_responses_only=True + if train_on_assistant_responses_only: + assert (loss_mask == 0.0).any(), "Expected some masked positions" + assert (loss_mask == 1.0).any(), "Expected some unmasked positions" + + # Per-document loss mask verification via extract_segments + verify_loss_mask(sample, config) diff --git a/tests/unit_tests/tokenizers/test_tokenizer.py b/tests/unit_tests/tokenizers/test_tokenizer.py index 48432bda8fb..57a7d6962c4 100755 --- a/tests/unit_tests/tokenizers/test_tokenizer.py +++ b/tests/unit_tests/tokenizers/test_tokenizer.py @@ -276,13 +276,15 @@ def test_tiktoken_tokenizer(): def test_null_tokenizer(): metadata = {"library": "null-text"} - tokenizer = MegatronTokenizer.from_pretrained(metadata_path=metadata, vocab_size=131072) + vocab_size = 131072 + tokenizer = MegatronTokenizer.from_pretrained(metadata_path=metadata, vocab_size=vocab_size) - ids = tokenizer.tokenize("11 325 97") + text = "11 325 97" + ids = tokenizer.tokenize(text) - assert ids == [11, 325, 97] - assert tokenizer.vocab_size == 131072 - assert tokenizer.eod == 131071 + assert ids == [ord(c) % vocab_size for c in text] + assert tokenizer.vocab_size == vocab_size + assert tokenizer.eod == vocab_size - 1 assert tokenizer.pad == -1 @@ -295,8 +297,9 @@ def test_detokenize_skip_special_tokens_unsupported_backend(library, skip_specia tokenizer = MegatronTokenizer.from_pretrained( metadata_path={"library": library}, vocab_size=131072 ) - ids = tokenizer.tokenize("11 325 97") - expected = "11 325 97" + text = "11 325 97" + ids = tokenizer.tokenize(text) + expected = ' '.join(str(ord(c) % 131072) for c in text) elif library == "byte-level": tokenizer = MegatronTokenizer.from_pretrained( metadata_path={"library": library}, vocab_size=1024, _bos_id=3, special_tokens=[]