diff --git a/examples/multimodal/train.py b/examples/multimodal/train.py index ba49e660445..3627d8838e1 100644 --- a/examples/multimodal/train.py +++ b/examples/multimodal/train.py @@ -27,7 +27,8 @@ is_pipeline_last_stage, ) from megatron.training import get_args, get_timers, get_tokenizer, pretrain -from megatron.training.utils import is_last_rank, get_batch_on_this_cp_rank +from megatron.core.utils import get_batch_on_this_cp_rank +from megatron.training.utils import is_last_rank def get_batch(data_iterator, image_token_index, img_seq_len): diff --git a/examples/post_training/modelopt/finetune.py b/examples/post_training/modelopt/finetune.py index f7f7c24f970..94309995f96 100755 --- a/examples/post_training/modelopt/finetune.py +++ b/examples/post_training/modelopt/finetune.py @@ -29,6 +29,8 @@ ) from utils import get_hf_tokenizer from model_provider import model_provider +from megatron.core.parallel_state import get_context_parallel_group + REMOVE_THINK_CHAT_TEMPLATE = ( "{% if '' in content %}{% set content = content.split('')[-1] %}{% endif %}" @@ -435,7 +437,7 @@ def get_batch(data_iterator): batch["hidden_states"] = feature_b["hidden_states"].transpose(0, 1)[:args.seq_length] # slice batch along sequence dimension for context parallelism - batch = get_batch_on_this_cp_rank(batch) + batch = get_batch_on_this_cp_rank(batch, is_hybrid_cp=False, cp_group=get_context_parallel_group()) return batch diff --git a/examples/post_training/modelopt/quantize.py b/examples/post_training/modelopt/quantize.py index dc4947038e5..64c429b36be 100644 --- a/examples/post_training/modelopt/quantize.py +++ b/examples/post_training/modelopt/quantize.py @@ -35,6 +35,7 @@ mtq_luts = None warnings.warn("luts is not installed. LUTs quantization configs will not be available.") +from megatron.core.parallel_state import get_context_parallel_group from megatron.core.utils import get_batch_on_this_cp_rank from megatron.post_training.arguments import add_modelopt_args from megatron.post_training.checkpointing import load_modelopt_checkpoint @@ -410,7 +411,9 @@ def _dataset_forward_loop_func(model): batch_size=args.calib_batch_size, ) for sample in tqdm(dataloader, disable=torch.distributed.get_rank()): - sample = get_batch_on_this_cp_rank(sample) + sample = get_batch_on_this_cp_rank( + sample, is_hybrid_cp=False, cp_group=get_context_parallel_group() + ) simple_generate(model, sample["input_ids"], osl=1, calibration_mode=True) unwrapped_model = unwrap_model(model)[0] 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/models/gpt/gpt_model.py b/megatron/core/models/gpt/gpt_model.py index 4fe641bb17b..b9a31fd72fd 100644 --- a/megatron/core/models/gpt/gpt_model.py +++ b/megatron/core/models/gpt/gpt_model.py @@ -143,7 +143,10 @@ def __init__( self.rotary_scaling = rope_scaling self.mtp_block_spec = mtp_block_spec self.mtp_process = mtp_block_spec is not None and mtp_on_this_rank( - self.config, ignore_virtual=False, vp_stage=vp_stage + layout=self.config.pipeline_model_parallel_layout, + mtp_num_layers=self.config.mtp_num_layers, + ignore_virtual=False, + vp_stage=vp_stage, ) if self.pre_process or self.mtp_process: diff --git a/megatron/core/models/mamba/mamba_model.py b/megatron/core/models/mamba/mamba_model.py index 444d6b86398..2634629d2c6 100644 --- a/megatron/core/models/mamba/mamba_model.py +++ b/megatron/core/models/mamba/mamba_model.py @@ -196,7 +196,12 @@ def __init__( # to split the hybrid layer pattern into pipeline stages before parsing the pattern for # the current pipeline stage. This could also enable MTP standalone (MTP in a pipeline # stage separate from loss) to be supported in the hybrid model. - and mtp_on_this_rank(self.config, ignore_virtual=False, vp_stage=self.vp_stage) + and mtp_on_this_rank( + layout=self.config.pipeline_model_parallel_layout, + mtp_num_layers=self.config.mtp_num_layers, + ignore_virtual=False, + vp_stage=self.vp_stage, + ) ) # megatron core pipelining currently depends on model type diff --git a/megatron/core/models/mimo/partition/utils.py b/megatron/core/models/mimo/partition/utils.py index 0b43e5548ff..592a6253b4a 100644 --- a/megatron/core/models/mimo/partition/utils.py +++ b/megatron/core/models/mimo/partition/utils.py @@ -235,7 +235,7 @@ def _apply_context_parallel( batch["attention_mask"] = attention_mask if packed_seq_params is None or getattr(packed_seq_params, 'qkv_format', 'sbhd') == 'sbhd': - batch = get_batch_on_this_cp_rank(batch) + batch = get_batch_on_this_cp_rank(batch, is_hybrid_cp=False, cp_group=self.cfg.cp_group) else: assert _HAVE_TEX and is_te_min_version("1.10.0"), ( "Please update Transformer Engine to >= 1.10 " diff --git a/megatron/core/models/multimodal/llava_model.py b/megatron/core/models/multimodal/llava_model.py index 86ce04521a7..bd050ce4128 100644 --- a/megatron/core/models/multimodal/llava_model.py +++ b/megatron/core/models/multimodal/llava_model.py @@ -726,9 +726,12 @@ def _process_embedding_token_parallel( batch["new_loss_mask"] = new_loss_mask # Distribute sequence across CP ranks if packed_seq_params is None or packed_seq_params.qkv_format == 'sbhd': - from megatron.training.utils import get_batch_on_this_cp_rank + from megatron.core.parallel_state import get_context_parallel_group + from megatron.core.utils import get_batch_on_this_cp_rank - batch = get_batch_on_this_cp_rank(batch) + batch = get_batch_on_this_cp_rank( + batch, is_hybrid_cp=False, cp_group=get_context_parallel_group() + ) else: assert HAVE_TEX and is_te_min_version( "1.10.0" diff --git a/megatron/core/tokenizers/text/libraries/abstract_tokenizer.py b/megatron/core/tokenizers/text/libraries/abstract_tokenizer.py index 360db03e5f2..009dbe53e6a 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 4e3387b125c..7a4a0366d6d 100644 --- a/megatron/core/tokenizers/text/libraries/huggingface_tokenizer.py +++ b/megatron/core/tokenizers/text/libraries/huggingface_tokenizer.py @@ -241,7 +241,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 4ddf77fc774..6e139c83ecc 100644 --- a/megatron/core/tokenizers/text/libraries/null_tokenizer.py +++ b/megatron/core/tokenizers/text/libraries/null_tokenizer.py @@ -15,10 +15,11 @@ def __init__(self, vocab_size): """ """ self._vocab_size_without_eod = int(vocab_size) self._eod_id = self._vocab_size_without_eod + self.pad_id = self._vocab_size_without_eod - 1 - 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_without_eod 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 8a418f2dd7f..ab26c83a77c 100644 --- a/megatron/core/tokenizers/text/libraries/sft_tokenizer.py +++ b/megatron/core/tokenizers/text/libraries/sft_tokenizer.py @@ -154,7 +154,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._tokenizer.apply_chat_template( [turn], tokenize=True, chat_template=self._prompt_config.custom_chat_template @@ -185,7 +188,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. @@ -193,7 +196,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 39228ad4afd..2be11d06001 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 0145ae353ca..c4a5a764d34 100644 --- a/megatron/core/tokenizers/text/text_tokenizer.py +++ b/megatron/core/tokenizers/text/text_tokenizer.py @@ -62,7 +62,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. @@ -73,7 +73,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]) -> str: """ diff --git a/megatron/core/tokenizers/utils/build_tokenizer.py b/megatron/core/tokenizers/utils/build_tokenizer.py index bf02451ae6c..3f14579afac 100644 --- a/megatron/core/tokenizers/utils/build_tokenizer.py +++ b/megatron/core/tokenizers/utils/build_tokenizer.py @@ -11,6 +11,8 @@ logger = logging.getLogger(__name__) +NULL_TOKENIZERS = {'NullTokenizer': 'null-text', 'NullMultimodalTokenizer': 'null-multimodal'} + def build_tokenizer(args, **kwargs): """Initialize tokenizer.""" @@ -67,10 +69,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/transformer/multi_token_prediction.py b/megatron/core/transformer/multi_token_prediction.py index 8fe7a2636b0..c4b2dce24c5 100755 --- a/megatron/core/transformer/multi_token_prediction.py +++ b/megatron/core/transformer/multi_token_prediction.py @@ -56,6 +56,8 @@ else: TESpecProvider = None +from megatron.core.transformer.pipeline_parallel_layer_layout import PipelineParallelLayerLayout + def tie_word_embeddings_state_dict( sharded_state_dict: ShardedStateDict, @@ -468,13 +470,16 @@ def get_mtp_layer_spec_for_backend( def mtp_on_this_rank( - config: TransformerConfig, ignore_virtual: Optional[bool] = True, vp_stage: Optional[int] = None + layout: PipelineParallelLayerLayout = None, + mtp_num_layers: Optional[int] = None, + ignore_virtual: Optional[bool] = True, + vp_stage: Optional[int] = None, ) -> bool: """ Check if there is MTP on the current rank. Behavior: - - If a custom pipeline model parallel layout is provided in the config: + - If a custom pipeline model parallel layout is provided: - If virtual pipeline parallelism is enabled (and `ignore_virtual` is False), checks whether any MTP layers are present on this (pp_rank, vp_stage) pair. - Otherwise, checks all virtual pipeline ranks of the current pipeline rank. Returns @@ -484,25 +489,24 @@ def mtp_on_this_rank( """ mtp_on_this_rank = False pp_rank = parallel_state.get_pipeline_model_parallel_rank() - if config.pipeline_model_parallel_layout is not None: + if layout is not None: # with custom PP layout, we support put MTP layers on any pipeline stage - layout = config.pipeline_model_parallel_layout.layout if ( not ignore_virtual and parallel_state.get_virtual_pipeline_model_parallel_world_size() is not None ): assert vp_stage is not None, "vp_stage must be passed if virtual pipeline is enabled" - num_layers_to_build = layout[pp_rank][vp_stage].count(LayerType.mtp) + num_layers_to_build = layout.layout[pp_rank][vp_stage].count(LayerType.mtp) mtp_on_this_rank = num_layers_to_build > 0 else: - for vpp_rank in range(len(layout[pp_rank])): - num_layers_to_build = layout[pp_rank][vpp_rank].count(LayerType.mtp) + for vpp_rank in range(len(layout.layout[pp_rank])): + num_layers_to_build = layout.layout[pp_rank][vpp_rank].count(LayerType.mtp) if num_layers_to_build > 0: mtp_on_this_rank = True break else: # without custom PP layout, we only support put all of MTP layers on the last pipeline stage - if config.mtp_num_layers is not None: + if mtp_num_layers is not None: mtp_on_this_rank = parallel_state.is_pipeline_last_stage( ignore_virtual=ignore_virtual, vp_stage=vp_stage ) diff --git a/megatron/core/utils.py b/megatron/core/utils.py index 58e86e09247..bcb563a345f 100644 --- a/megatron/core/utils.py +++ b/megatron/core/utils.py @@ -32,7 +32,6 @@ from megatron.core import config from megatron.core._rank_utils import log_single_rank from megatron.core.package_info import __version__ as mcore_version -from megatron.core.packed_seq_params import PackedSeqParams try: from torch.distributed._tensor import DTensor @@ -67,9 +66,13 @@ # Alias the PyTorch wrapper so we can call tex.* APIs import transformer_engine_torch as tex + from transformer_engine.pytorch.attention.dot_product_attention.context_parallel import ( + pad_thd_sequences_for_cp, + ) except ImportError: # TE isn’t installed or the torch wrapper is missing tex = None + pad_thd_sequences_for_cp = None try: _torch_version = PkgVersion(torch.__version__) @@ -82,6 +85,8 @@ _mamba_ssm_version = None _causal_conv1d_version = None +# TODO(asolergi-nv): Clean up imports! + @contextmanager def null_decorator(*args, **kwargs): @@ -1965,41 +1970,579 @@ def is_submodule(module, parent_module, strict=True): ######################## -### context parallel ### +###### sft utils ####### ######################## -def get_batch_on_this_cp_rank( - batch: Dict[str, Any], cp_group: Optional[torch.distributed.ProcessGroup] = None +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), while ``cu_seqlens``, ``cu_seqlens_padded``, and + ``max_seqlen`` remain 1-D. + + 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, cu_seqlens_padded = pad_thd_sequences_for_cp( + tokens, + labels, + cu_seqlens, + divisibility_factor, + padding_token_id=padding_token_id, + padding_label_id=padding_label_id, + ) + # NOTE(asolergi-nv): Pad loss_mask for CP divisibility (reuse TE's padding + # with 0 for masked positions) + _, loss_mask, _ = pad_thd_sequences_for_cp( + loss_mask, + loss_mask, + cu_seqlens, + divisibility_factor, + padding_token_id=0, + padding_label_id=0, + ) + cu_seqlens_padded = cu_seqlens_padded.to( + torch.int32 + ) # NOTE(asolergi-nv): pad_thd_sequences_for_cp uses torch.cumsum + # which promotes cu_seqlens_padded from int32 to int64 + 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, + 'cu_seqlens_padded': cu_seqlens_padded, + 'max_seqlen': max_seqlen, + } + return batch + + +######################## +### tensor parallel #### +######################## + + +def get_batch_on_this_tp_rank( + batch: dict[str, torch.Tensor], + is_sft: bool, + is_hybrid_cp: bool, + create_attention_mask_in_dataloader: bool, + broadcast_src_rank: int, + broadcast_group: torch.distributed.ProcessGroup, + cp_size: int, + tp_rank: int, + micro_batch_size: int, + seq_length: int, + mtp_on_this_rank: bool, + pipeline_model_parallel_size: int = 1, + is_pipeline_first_stage: bool = False, + is_pipeline_last_stage: bool = False, ): - """Slice batch input along sequence dimension into multiple chunks, - which are parallelized across GPUs in a context parallel group. + """Broadcast batch tensors from TP rank 0 to all other ranks in the TP group. + + TP rank 0 holds the fully preprocessed batch (from the dataloader or from + ``preprocess_sft_batch`` when SFT is enabled). This function broadcasts + every required tensor to the remaining TP ranks so that all ranks hold + identical data before the forward pass. The set of tensors broadcast depends + on the pipeline stage and whether SFT / hybrid-CP modes are active. + + For SFT and hybrid-CP, variable-length metadata (``cu_seqlens``, + ``cu_seqlens_padded``) is broadcast using a length-prefixed protocol: TP + rank 0 first sends the numel, then the tensor itself, so receivers can + allocate the correct buffer size. + + For hybrid-CP, the sequence length may differ per micro-batch (since it + depends on `local_cp_size`), so the actual sequence length is broadcast + before allocating receive buffers on non-zero TP ranks. Args: - batch (Dict[str, Any]): Input batch tensors. - cp_group (Optional[torch.distributed.ProcessGroup]): Context-parallel process group. - If provided, uses this group's size and rank. Otherwise, falls back to - the current context-parallel settings from parallel_state. + batch (dict[str, torch.Tensor]): The batch dict. On TP rank 0 this + contains the actual data; on other ranks it is ignored (receive + buffers are allocated internally). + is_sft (bool): Whether this is an SFT (supervised fine-tuning) run + using THD packed sequences. + is_hybrid_cp (bool): Whether hybrid context parallelism is enabled. + create_attention_mask_in_dataloader (bool): Whether the dataloader + creates an explicit attention mask tensor. + broadcast_src_rank (int): Global rank of the broadcast source (TP rank 0). + broadcast_group (torch.distributed.ProcessGroup): The TP process group + used for broadcasting. + cp_size (int): Context-parallel world size. + tp_rank (int): This rank's position within the TP group. + micro_batch_size (int): Micro-batch size (number of samples). + seq_length (int): Sequence length used for allocating receive buffers + (ignored under hybrid-CP where it is broadcast dynamically). + mtp_on_this_rank (bool): Whether Multi-Token Prediction layers are + active on this rank (affects which tensors are needed). + pipeline_model_parallel_size (int): Number of pipeline-parallel stages. + is_pipeline_first_stage (bool): Whether this rank is on the first PP stage. + is_pipeline_last_stage (bool): Whether this rank is on the last PP stage. + + Returns: + dict[str, torch.Tensor]: The batch dict with all tensors populated on + every TP rank. Keys include 'tokens', 'labels', 'loss_mask', + 'position_ids', 'attention_mask', 'cu_seqlens', 'cu_seqlens_padded', + 'max_seqlen', 'local_cp_size', and 'hybrid_cp_group'. """ + # TODO(asolergi-nv): Enable PP with sft + + def _broadcast(item): + if item is not None: + torch.distributed.broadcast(item, broadcast_src_rank, group=broadcast_group) + + if tp_rank == 0: + + def _broadcast_cu_seqlens(cu_seqlens): + dev = torch.cuda.current_device() + n = 0 if cu_seqlens is None else int(cu_seqlens.numel()) + n_tensor = torch.tensor(n, dtype=torch.int64, device=dev) + _broadcast(n_tensor) + + if n == 0: + buf = torch.empty(0, dtype=torch.int32, device=dev) + else: + assert isinstance( + cu_seqlens, torch.Tensor + ), f"Expected cu_seqlens to be a torch.Tensor, got {type(cu_seqlens)}" + assert ( + cu_seqlens.dtype == torch.int32 + ), f"Expected cu_seqlens to be of type torch.int32, got {cu_seqlens.dtype}" + buf = cu_seqlens + _broadcast(buf) + + if is_hybrid_cp: + hybrid_cp_seq_length = torch.tensor( + batch['tokens'].shape[1], dtype=torch.int32, device=torch.cuda.current_device() + ) + _broadcast(hybrid_cp_seq_length) + + if pipeline_model_parallel_size == 1 or mtp_on_this_rank: + _broadcast(batch['tokens']) + _broadcast(batch['labels']) + _broadcast(batch['loss_mask']) + _broadcast(batch['position_ids']) + if is_sft or is_hybrid_cp: + _broadcast_cu_seqlens(batch['cu_seqlens']) + _broadcast(batch['max_seqlen']) + if cp_size > 1: + _broadcast_cu_seqlens(batch['cu_seqlens_padded']) + if create_attention_mask_in_dataloader: + _broadcast(batch['attention_mask']) + if is_hybrid_cp: + _broadcast(batch['local_cp_size']) + + elif is_pipeline_first_stage: + _broadcast(batch['tokens']) + _broadcast(batch['position_ids']) + if is_sft or is_hybrid_cp: + _broadcast_cu_seqlens(batch['cu_seqlens']) + _broadcast(batch['max_seqlen']) + if cp_size > 1: + _broadcast_cu_seqlens(batch['cu_seqlens_padded']) + if create_attention_mask_in_dataloader: + _broadcast(batch['attention_mask']) + if is_hybrid_cp: + _broadcast(batch['local_cp_size']) + + elif is_pipeline_last_stage: + # Multi-Token Prediction (MTP) layers need tokens and position_ids to calculate + # embedding. Currently the Multi-Token Prediction (MTP) layers is fixed on the + # last stage, so we need to broadcast tokens and position_ids to all of the + # tensor parallel ranks on the last stage. + _broadcast(batch['labels']) + _broadcast(batch['loss_mask']) + if create_attention_mask_in_dataloader: + _broadcast(batch['attention_mask']) - # With causal masking, each token only attends to its prior tokens. Simply split - # sequence into CP chunks can result in severe load imbalance. That's to say, chunks - # at the end of sequence have bigger workload than others. To address this issue, - # we split sequence into 2*CP ranks. Assuming CP=2, we then get 4 chunks, chunk_0 - # and chunk_3 are assigned to GPU0, chunk_1 and chunk_2 are assigned to GPU1, so - # that we can get balanced workload among GPUs in a context parallel group. - # Determine CP topology either from provided group or from current context parallel state - if cp_group is not None: - cp_size = get_pg_size(cp_group) - cp_rank = get_pg_rank(cp_group) else: - cp_size = parallel_state.get_context_parallel_world_size() - cp_rank = parallel_state.get_context_parallel_rank() + if is_hybrid_cp: + hybrid_cp_seq_length = torch.tensor( + 0, dtype=torch.int32, device=torch.cuda.current_device() + ) + _broadcast(hybrid_cp_seq_length) + shape = (micro_batch_size, hybrid_cp_seq_length.item()) + else: + shape = (micro_batch_size, seq_length) + + tokens = torch.empty(shape, dtype=torch.int64, device=torch.cuda.current_device()) + labels = torch.empty(shape, dtype=torch.int64, device=torch.cuda.current_device()) + loss_mask = torch.empty(shape, dtype=torch.float32, device=torch.cuda.current_device()) + position_ids = torch.empty(shape, dtype=torch.int64, device=torch.cuda.current_device()) + cu_seqlens = None + cu_seqlens_padded = None + max_seqlen = None + attention_mask = None + local_cp_size = None + + if is_sft or is_hybrid_cp: + max_seqlen = torch.empty(1, dtype=torch.int32, device=torch.cuda.current_device()) + if create_attention_mask_in_dataloader: + attention_mask = torch.empty( + (micro_batch_size, 1, seq_length, seq_length), + dtype=torch.bool, + device=torch.cuda.current_device(), + ) + + if is_hybrid_cp: + local_cp_size = torch.empty(1, dtype=torch.int32, device=torch.cuda.current_device()) + + def _broadcast_cu_seqlens(): + dev = torch.cuda.current_device() + + n = torch.empty((), dtype=torch.int64, device=dev) + _broadcast(n) + n = int(n.item()) + + if n == 0: + cu_seqlens = torch.empty(0, dtype=torch.int32, device=dev) + else: + cu_seqlens = torch.empty(n, dtype=torch.int32, device=dev) + _broadcast(cu_seqlens) + + assert ( + cu_seqlens.numel() > 0 + ), f"Expected cu_seqlens to have more than 0 elements, got {cu_seqlens.numel()}" + + return cu_seqlens + + if pipeline_model_parallel_size == 1 or mtp_on_this_rank: + _broadcast(tokens) + _broadcast(labels) + _broadcast(loss_mask) + _broadcast(position_ids) + if is_sft or is_hybrid_cp: + cu_seqlens = _broadcast_cu_seqlens() + _broadcast(max_seqlen) + if cp_size > 1: + cu_seqlens_padded = _broadcast_cu_seqlens() + if create_attention_mask_in_dataloader: + _broadcast(attention_mask) + if is_hybrid_cp: + _broadcast(local_cp_size) + + elif is_pipeline_first_stage: + labels = None + loss_mask = None + + _broadcast(tokens) + _broadcast(position_ids) + if is_sft or is_hybrid_cp: + cu_seqlens = _broadcast_cu_seqlens() + _broadcast(max_seqlen) + if cp_size > 1: + cu_seqlens_padded = _broadcast_cu_seqlens() + if create_attention_mask_in_dataloader: + _broadcast(attention_mask) + if is_hybrid_cp: + _broadcast(local_cp_size) + + elif is_pipeline_last_stage: + # Multi-Token Prediction (MTP) layers need tokens and position_ids + # to calculate embedding. Currently the Multi-Token Prediction (MTP) layers + # is fixed on the last stage, so we need to broadcast tokens and position_ids + # to all of the tensor parallel ranks on the last stage. + tokens = None + position_ids = None + cu_seqlens = None + cu_seqlens_padded = None + max_seqlen = None + + _broadcast(labels) + _broadcast(loss_mask) + if create_attention_mask_in_dataloader: + _broadcast(attention_mask) + + batch = { + 'tokens': tokens, + 'labels': labels, + 'loss_mask': loss_mask, + 'position_ids': position_ids, + 'attention_mask': attention_mask, + 'cu_seqlens': cu_seqlens, + 'cu_seqlens_padded': cu_seqlens_padded, + 'max_seqlen': max_seqlen, + 'local_cp_size': local_cp_size, + 'hybrid_cp_group': None, + } + + return batch + + +######################## +### context parallel ### +######################## + + +def get_sft_batch_on_this_cp_rank( + batch: dict[str, torch.Tensor], cp_group: torch.distributed.ProcessGroup +): + """Partition an SFT packed-sequence batch across context-parallel ranks using THD indexing. + + For SFT workloads the batch contains multiple variable-length sub-sequences + packed contiguously (THD format). This function uses Transformer Engine's + ``thd_get_partitioned_indices`` to compute the token indices assigned to the + current CP rank and gathers only those tokens from every sequence-dimension + tensor in the batch. + + Metadata keys ('attention_mask', 'cu_seqlens', 'cu_seqlens_padded', + 'max_seqlen', 'local_cp_size', 'hybrid_cp_group') are left unchanged + because TE's attention kernels consume them directly. + + Args: + batch (dict[str, torch.Tensor]): Batch dict with tensors of shape + ``[micro_batch_size, seq_length, ...]``. + cp_group (torch.distributed.ProcessGroup): The context-parallel process + group. + + Returns: + dict[str, torch.Tensor]: The batch with sequence-dimension tensors + index-selected to this CP rank's partition. + """ + cp_size = torch.distributed.get_world_size(cp_group) + cp_rank = torch.distributed.get_rank(cp_group) + + if cp_size > 1: + index = tex.thd_get_partitioned_indices( + ( + batch["cu_seqlens_padded"] + if batch["cu_seqlens_padded"] is not None + else batch["cu_seqlens"] + ), + ( + batch["tokens"].size(1) if batch["tokens"] is not None else batch["labels"].size(1) + ), # NOTE(asolergi-nv): Labels to enable PP! + cp_size, + cp_rank, + ) + SEQUENCE_KEYS = ('tokens', 'labels', 'loss_mask', 'position_ids') + for key in SEQUENCE_KEYS: + if batch.get(key) is not None: + batch[key] = batch[key].index_select(1, index) + return batch + + +def get_pretrain_batch_on_this_cp_rank( + batch: dict[str, torch.Tensor], cp_group: torch.distributed.ProcessGroup +): + """Partition a pretraining batch across context-parallel ranks with load-balanced chunking. + + With causal masking, each token only attends to its prior tokens. Simply splitting + the sequence into CP chunks can result in severe load imbalance, as chunks at the + end of the sequence have bigger workloads than earlier ones. To address this, the + sequence is split into ``2 * cp_size`` chunks and assigned in a zigzag pattern: + for CP=2 the 4 chunks are assigned as (chunk_0, chunk_3) -> GPU 0 and + (chunk_1, chunk_2) -> GPU 1, balancing the workload across the CP group. + + All tensor-valued entries in the batch are partitioned along their sequence + dimension (``seq_dim=1`` by default, ``seq_dim=2`` for 'attention_mask'). + None-valued entries are left unchanged. + + Args: + batch (dict[str, torch.Tensor]): Batch dict with tensors of shape + ``[micro_batch_size, seq_length, ...]``. + cp_group (torch.distributed.ProcessGroup): The context-parallel process + group. + + Returns: + dict[str, torch.Tensor]: The batch with sequence-dimension tensors + sliced to this CP rank's zigzag partition. + """ + + cp_size = torch.distributed.get_world_size(cp_group) + cp_rank = torch.distributed.get_rank(cp_group) if cp_size > 1: for key, val in batch.items(): if val is not None: - seq_dim = 1 if key != 'attention_mask' else 2 + seq_dim = 2 if key == 'attention_mask' else 1 + if not isinstance(val, torch.Tensor) or val.dim() <= seq_dim: + # NOTE(asolergi-nv): HybridCP includes 1D metadata tensors + # like cu_seqlens, cu_seqlens_padded, max_seqlen, local_cp_size + continue val = val.view( *val.shape[0:seq_dim], 2 * cp_size, @@ -2016,98 +2559,55 @@ def get_batch_on_this_cp_rank( return batch -def get_thd_batch_on_this_cp_rank( +def get_batch_on_this_cp_rank( batch: Dict[str, Any], - cu_seqlens: torch.Tensor, - cu_seqlens_padded: torch.Tensor, - max_seqlen: torch.Tensor, - cp_size: Optional[int] = None, - cp_rank: Optional[int] = None, + is_hybrid_cp: bool, + cp_group: Optional[torch.distributed.ProcessGroup] = None, + hybrid_cp_group_func: Optional[Callable[[int], torch.distributed.ProcessGroup]] = None, ): - """Slice each sub-sample in a packed sample batch input along - sequence dimension into multiple chunks, which are parallelized - across GPUs in a context parallel group. - """ - packed_seq_params = PackedSeqParams( - qkv_format="thd", - cu_seqlens_q=cu_seqlens, - cu_seqlens_kv=cu_seqlens, - cu_seqlens_q_padded=cu_seqlens_padded, - cu_seqlens_kv_padded=cu_seqlens_padded, - max_seqlen_q=int(max_seqlen[0].item()), - max_seqlen_kv=int(max_seqlen[0].item()), - ) - - cp_size = parallel_state.get_context_parallel_world_size() if cp_size is None else cp_size - cp_rank = parallel_state.get_context_parallel_rank() if cp_rank is None else cp_rank - if cp_size > 1: # slice batch along sequence dimension for context parallelism - assert tex is not None and is_te_min_version("1.10.0"), ( - "Please update Transformer Engine to >= 1.10 to use " - "Context Parallel with THD format data" - ) - index = tex.thd_get_partitioned_indices( - cu_seqlens_padded, batch['tokens'].size(1), cp_size, cp_rank - ) - for key, data in batch.items(): - if key in {'attention_mask', 'cu_seqlens', 'cu_seqlens_padded', 'max_seqlen'}: - continue - batch[key] = data.index_select(1, index) - - return batch, packed_seq_params - - -################################ -### hybrid context parallel ### -################################ + """Dispatch batch partitioning across context-parallel ranks. + + Routes to the appropriate CP partitioning strategy based on the batch + contents and parallelism mode: + - **SFT (packed sequences)**: When ``cu_seqlens`` is present and + ``is_hybrid_cp`` is False, delegates to ``get_sft_batch_on_this_cp_rank`` + which uses THD index-based partitioning. + - **Hybrid CP**: When ``cu_seqlens`` is present and ``is_hybrid_cp`` is + True, creates a local hybrid CP group (via ``hybrid_cp_group_func``) + and delegates to ``get_pretrain_batch_on_this_cp_rank`` with that group. + - **Pretraining**: When ``cu_seqlens`` is None, delegates to + ``get_pretrain_batch_on_this_cp_rank`` with zigzag load-balanced + chunking. + Args: + batch (Dict[str, Any]): Input batch tensors. Must contain a + 'cu_seqlens' key (may be None for pretraining). + is_hybrid_cp (bool): Whether hybrid context parallelism is enabled. + cp_group (Optional[torch.distributed.ProcessGroup]): Context-parallel + process group used for SFT and pretraining CP partitioning. + hybrid_cp_group_func (Optional[Callable[[int], torch.distributed.ProcessGroup]]): + Factory function that returns a hybrid CP process group for a given + ``group_size``. Required when ``is_hybrid_cp`` is True. -def get_batch_on_this_hybrid_cp_rank( - batch: Dict[str, Any], - local_cp_size: int, - cp_group: Optional[torch.distributed.ProcessGroup] = None, -): - """Slice batch input along sequence dimension into multiple chunks, - which are parallelized across GPUs in a context parallel group. + Returns: + Dict[str, Any]: The batch with sequence-dimension tensors partitioned + to this CP rank. """ - assert local_cp_size is not None - if cp_group is None: - # Get the local cp group required for as defined by the HybridCPDataLoaderWrapper - if local_cp_size > 1: - cp_group = parallel_state.get_hybrid_data_context_parallel_groups( - group_size=local_cp_size - ) - else: - # If cp group is provided, it must match the local cp size - # as defined by the HybridCPDataLoaderWrapper - assert cp_group.size() == local_cp_size - - # Convert [seqlen] to [1, seqlen] similar to default collate_fn - # as hybrid_context_parallel dataloader wrapper does not go through default collate_fn - for key, data in batch.items(): - if key in ['attention_mask']: - continue - batch[key] = torch.stack([data], 0) - sample_length = batch['tokens'].shape[1] - # TODO(pmannan): Take care of padding tokens here if not divisible by cp_size*2 - # Create packed_seq_params for SBHD format with cp group information. - packed_seq_params = PackedSeqParams( - qkv_format="sbhd", - cu_seqlens_q=torch.tensor([0, sample_length], device="cuda", pin_memory=True), - cu_seqlens_kv=torch.tensor([0, sample_length], device="cuda", pin_memory=True), - cu_seqlens_q_padded=torch.tensor([0, sample_length], device="cuda", pin_memory=True), - cu_seqlens_kv_padded=torch.tensor([0, sample_length], device="cuda", pin_memory=True), - max_seqlen_q=sample_length, - max_seqlen_kv=sample_length, - local_cp_size=local_cp_size, - cp_group=cp_group, - ) - - if cp_group is not None and cp_group.size() > 1: - # When using hybrid_context_parallel, each sub-sample of a packed sample is - # required to be divisible by CP*DP*2 or CP*DP*TP*2 (if using sequence parallel) - batch = get_batch_on_this_cp_rank(batch, cp_group=cp_group) - return batch, packed_seq_params + if batch.get("cu_seqlens") is not None: # NOTE(asolergi-nv): SFT & HybridCP case + if is_hybrid_cp: + assert ( + batch['local_cp_size'] is not None + ), "local_cp_size is required for hybrid context parallel" + if batch['local_cp_size'].item() > 1: + hybrid_cp_group = hybrid_cp_group_func(group_size=batch['local_cp_size'].item()) + batch = get_pretrain_batch_on_this_cp_rank(batch, cp_group=hybrid_cp_group) + batch["hybrid_cp_group"] = hybrid_cp_group + else: + batch = get_sft_batch_on_this_cp_rank(batch, cp_group=cp_group) + else: # NOTE(asolergi-nv): Pretrain case + batch = get_pretrain_batch_on_this_cp_rank(batch, cp_group=cp_group) + return batch ###################### diff --git a/megatron/training/datasets/sft_dataset.py b/megatron/training/datasets/sft_dataset.py deleted file mode 100644 index 9de5d2a52fe..00000000000 --- a/megatron/training/datasets/sft_dataset.py +++ /dev/null @@ -1,192 +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 - - 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': cu_seqlens, - 'max_seqlen': max_seqlen, - } diff --git a/megatron/training/training.py b/megatron/training/training.py index aba11d04df9..9496598a2db 100644 --- a/megatron/training/training.py +++ b/megatron/training/training.py @@ -186,7 +186,8 @@ def set_startup_timestamps(program_start=None, main_entry=None): from megatron.training.initialize import initialize_megatron from megatron.training.initialize import write_args_to_tensorboard from megatron.training.initialize import set_jit_fusion_options -from megatron.training.utils import get_batch_on_this_cp_rank, get_batch_on_this_tp_rank, is_hybrid_model +from megatron.training.utils import is_hybrid_model +from megatron.core.utils import get_batch_on_this_cp_rank, get_batch_on_this_tp_rank from megatron.training.datasets.data_samplers import build_pretraining_data_loader from megatron.core.datasets.data_schedule import HybridCPDataLoaderWrapper from megatron.core.optimizer_param_scheduler import OptimizerParamScheduler @@ -197,7 +198,9 @@ def set_startup_timestamps(program_start=None, main_entry=None): from megatron.core.parallel_state import ( destroy_global_memory_buffer, destroy_model_parallel, - update_pg_timeout + get_context_parallel_group, + get_hybrid_data_context_parallel_groups, + update_pg_timeout, ) from megatron.core.inference.symmetric_memory import SymmetricMemoryManager from megatron.core.inference.unified_memory import create_unified_mempool @@ -1770,13 +1773,49 @@ def setup_model_and_optimizer( def dummy_train_step(data_iterator): """Single dummy training step.""" + args = get_args() + tp_rank = mpu.get_tensor_model_parallel_rank() + is_sft = getattr(args, 'sft', False) + is_hybrid_cp = args.hybrid_context_parallel + + BATCH_KEYS = [ + "tokens", "labels", "loss_mask", "position_ids", "attention_mask", + "cu_seqlens", "cu_seqlens_padded", "max_seqlen", "local_cp_size", + "hybrid_cp_group", + ] + num_microbatches = get_num_microbatches() rerun_state_machine = get_rerun_state_machine() while rerun_state_machine.should_run_forward_backward(data_iterator): for _ in range(num_microbatches): # Re-use methods used in get_batch() from pretrain_{gpt, mamba}.py. - batch = get_batch_on_this_tp_rank(data_iterator) - batch = get_batch_on_this_cp_rank(batch) + batch = {} + if tp_rank == 0: + batch = next(data_iterator) + for key in BATCH_KEYS: + batch[key] = batch[key].cuda(non_blocking=True) if key in batch and batch[key] is not None else None + batch = get_batch_on_this_tp_rank( + batch, + broadcast_src_rank=mpu.get_tensor_model_parallel_src_rank(), + broadcast_group=mpu.get_tensor_model_parallel_group(), + is_sft=is_sft, + is_hybrid_cp=is_hybrid_cp, + create_attention_mask_in_dataloader=args.create_attention_mask_in_dataloader, + cp_size=args.context_parallel_size, + tp_rank=tp_rank, + micro_batch_size=args.micro_batch_size, + seq_length=args.seq_length, + mtp_on_this_rank=False, + pipeline_model_parallel_size=args.pipeline_model_parallel_size, + is_pipeline_first_stage=mpu.is_pipeline_first_stage(), + is_pipeline_last_stage=mpu.is_pipeline_last_stage(), + ) + batch = get_batch_on_this_cp_rank( + batch, + is_hybrid_cp=is_hybrid_cp, + cp_group=get_context_parallel_group(), + hybrid_cp_group_func=get_hybrid_data_context_parallel_groups, + ) def train_step(forward_step_func, data_iterator, model, optimizer, opt_param_scheduler, config, forward_backward_func, iteration=None): diff --git a/megatron/training/utils.py b/megatron/training/utils.py index 843f3d74ec9..0c466a3aaa4 100644 --- a/megatron/training/utils.py +++ b/megatron/training/utils.py @@ -11,7 +11,7 @@ import torch -from megatron.core.msc_utils import MultiStorageClientFeature, open_file +from megatron.core.msc_utils import open_file from megatron.core._rank_utils import safe_get_rank as _safe_get_rank try: @@ -37,7 +37,6 @@ from megatron.core.datasets.utils import get_blend_from_list from megatron.core.tensor_parallel import param_is_not_tensor_parallel_duplicate from megatron.core.utils import ( - get_batch_on_this_cp_rank, get_data_parallel_group_if_dtensor, to_local_if_dtensor, unwrap_model, @@ -520,208 +519,6 @@ def get_blend_and_blend_per_split(args): return blend, blend_per_split -def get_batch_on_this_tp_rank(data_iterator, mtp_on_this_rank: bool = False): - - args = get_args() - - def _broadcast(item): - if item is not None: - torch.distributed.broadcast( - item, - mpu.get_tensor_model_parallel_src_rank(), - group=mpu.get_tensor_model_parallel_group(), - ) - - if mpu.get_tensor_model_parallel_rank() == 0: - - assert data_iterator is not None - data = next(data_iterator) - batch = { - 'tokens': data["tokens"].cuda(non_blocking=True), - 'labels': data["labels"].cuda(non_blocking=True), - 'loss_mask': data["loss_mask"].cuda(non_blocking=True), - 'attention_mask': ( - None - if "attention_mask" not in data - else data["attention_mask"].cuda(non_blocking=True) - ), - 'position_ids': data["position_ids"].cuda(non_blocking=True), - 'cu_seqlens': ( - None - if "cu_seqlens" not in data - else data["cu_seqlens"].cuda(non_blocking=True) - ), - 'max_seqlen': ( - None - if "max_seqlen" not in data - else data["max_seqlen"].cuda(non_blocking=True) - ), - 'local_cp_size': ( - None - if "local_cp_size" not in data - else data["local_cp_size"].cuda(non_blocking=True) - ), - } - - def _broadcast_cu_seqlens(cu_seqlens): - dev = torch.cuda.current_device() - n = 0 if cu_seqlens is None else int(cu_seqlens.numel()) - n_tensor = torch.tensor(n, dtype=torch.int64, device=dev) - _broadcast(n_tensor) - - if n == 0: - buf = torch.empty(0, dtype=torch.int32, device=dev) - else: - assert isinstance(cu_seqlens, torch.Tensor) - assert cu_seqlens.dtype == torch.int32 - assert cu_seqlens.shape[0] == 1, "micro-batch-size must be 1 for packing" - buf = cu_seqlens.to(device=dev, non_blocking=True).contiguous() - _broadcast(buf) - - if args.hybrid_context_parallel: - seq_len = torch.tensor(batch['tokens'].shape[0], dtype=torch.int32, device=torch.cuda.current_device()) - _broadcast(seq_len) - - if args.pipeline_model_parallel_size == 1 or mtp_on_this_rank: - _broadcast(batch['tokens']) - _broadcast(batch['labels']) - _broadcast(batch['loss_mask']) - _broadcast(batch['attention_mask']) - _broadcast(batch['position_ids']) - _broadcast_cu_seqlens(batch['cu_seqlens']) - _broadcast(batch['max_seqlen']) - _broadcast(batch['local_cp_size']) - - elif mpu.is_pipeline_first_stage(): - _broadcast(batch['tokens']) - _broadcast(batch['attention_mask']) - _broadcast(batch['position_ids']) - _broadcast_cu_seqlens(batch['cu_seqlens']) - _broadcast(batch['max_seqlen']) - - elif mpu.is_pipeline_last_stage(): - # Multi-Token Prediction (MTP) layers need tokens and position_ids to calculate embedding. - # Currently the Multi-Token Prediction (MTP) layers is fixed on the last stage, so we need - # to broadcast tokens and position_ids to all of the tensor parallel ranks on the last stage. - _broadcast(batch['labels']) - _broadcast(batch['loss_mask']) - _broadcast(batch['attention_mask']) - - else: - if args.hybrid_context_parallel: - seq_len = torch.tensor(0, dtype=torch.int32, device=torch.cuda.current_device()) - _broadcast(seq_len) - shape = (seq_len.item()) - else: - shape = (args.micro_batch_size, args.seq_length) - - tokens = torch.empty( - shape, - dtype=torch.int64, - device=torch.cuda.current_device(), - ) - labels = torch.empty( - shape, - dtype=torch.int64, - device=torch.cuda.current_device(), - ) - loss_mask = torch.empty( - shape, - dtype=torch.float32, - device=torch.cuda.current_device(), - ) - if args.create_attention_mask_in_dataloader: - shape_attention_mask = (args.micro_batch_size, 1, args.seq_length, args.seq_length) if not args.hybrid_context_parallel else (1, 1, shape[0], shape[0]) - attention_mask = torch.empty( - shape_attention_mask, - dtype=torch.bool, - device=torch.cuda.current_device(), - ) - else: - attention_mask = None - position_ids = torch.empty( - shape, - dtype=torch.int64, - device=torch.cuda.current_device(), - ) - cu_seqlens = None - if args.hybrid_context_parallel or args.sft: - max_seqlen = torch.empty( - 1, - dtype=torch.int32, - device=torch.cuda.current_device(), - ) - else: - max_seqlen = None - - local_cp_size = torch.empty( - 1, - dtype=torch.int32, - device=torch.cuda.current_device(), - ) if args.hybrid_context_parallel else None - - def _broadcast_cu_seqlens(): - dev = torch.cuda.current_device() - - n = torch.empty((), dtype=torch.int64, device=dev) - _broadcast(n) - n = int(n.item()) - - if n == 0: - cu_seqlens = torch.empty(0, dtype=torch.int32, device=dev) - else: - cu_seqlens = torch.empty((args.micro_batch_size, n), dtype=torch.int32, device=dev) - _broadcast(cu_seqlens) - - return cu_seqlens if n > 0 else None - - if args.pipeline_model_parallel_size == 1 or mtp_on_this_rank: - _broadcast(tokens) - _broadcast(labels) - _broadcast(loss_mask) - _broadcast(attention_mask) - _broadcast(position_ids) - cu_seqlens = _broadcast_cu_seqlens() - _broadcast(max_seqlen) - _broadcast(local_cp_size) - - elif mpu.is_pipeline_first_stage(): - labels = None - loss_mask = None - - _broadcast(tokens) - _broadcast(attention_mask) - _broadcast(position_ids) - cu_seqlens = _broadcast_cu_seqlens() - _broadcast(max_seqlen) - - elif mpu.is_pipeline_last_stage(): - # Multi-Token Prediction (MTP) layers need tokens and position_ids to calculate embedding. - # Currently the Multi-Token Prediction (MTP) layers is fixed on the last stage, so we need - # to broadcast tokens and position_ids to all of the tensor parallel ranks on the last stage. - tokens = None - position_ids = None - cu_seqlens = None - max_seqlen = None - - _broadcast(labels) - _broadcast(loss_mask) - _broadcast(attention_mask) - - batch = { - 'tokens': tokens, - 'labels': labels, - 'loss_mask': loss_mask, - 'attention_mask': attention_mask, - 'position_ids': position_ids, - 'cu_seqlens': cu_seqlens, - 'max_seqlen': max_seqlen, - 'local_cp_size': local_cp_size, - } - - return batch - - def update_use_dist_ckpt(args): args.use_dist_ckpt = args.ckpt_format != "torch" diff --git a/pretrain_gpt.py b/pretrain_gpt.py index 31eee0f4dc6..7ed87ffe42d 100644 --- a/pretrain_gpt.py +++ b/pretrain_gpt.py @@ -22,30 +22,30 @@ import torch from gpt_builders import gpt_builder -from megatron.core import parallel_state +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.enums import ModelType from megatron.core.packed_seq_params import PackedSeqParams +from megatron.core.parallel_state import get_context_parallel_group, get_hybrid_data_context_parallel_groups from megatron.core.models.gpt import GPTModel from megatron.core.rerun_state_machine import get_rerun_state_machine from megatron.core.tokenizers.utils.build_tokenizer import build_tokenizer -from megatron.core.utils import get_attr_wrapped_model, get_thd_batch_on_this_cp_rank, get_batch_on_this_hybrid_cp_rank, StragglerDetector +from megatron.core.utils import get_attr_wrapped_model, StragglerDetector, preprocess_sft_batch, get_batch_on_this_cp_rank, get_batch_on_this_tp_rank from megatron.training import ( get_args, get_timers, + get_tokenizer, inprocess_restart, pretrain, print_rank_0, set_startup_timestamps, ) -from megatron.training.datasets.sft_dataset import SFTDataset -from megatron.core.transformer.multi_token_prediction import mtp_on_this_rank, get_mtp_ranks +from megatron.core.transformer.multi_token_prediction import mtp_on_this_rank as mtp_on_this_rank_func, get_mtp_ranks from megatron.training.arguments import core_transformer_config_from_args +from megatron.core.datasets.sft_dataset import SFTDataset, SFTDatasetConfig, IGNORE_INDEX from megatron.training.datasets.fim_dataset import GPTFIMDataset, GPTFIMDatasetConfig from megatron.training.utils import ( - get_batch_on_this_cp_rank, - get_batch_on_this_tp_rank, get_blend_and_blend_per_split, is_first_or_last_pipeline_stage, ) @@ -63,104 +63,37 @@ def get_batch(data_iterator, vp_stage: Optional[int] = None): - """Generate a batch. - - Packed sequence support (SFT / ``--sft`` flag): - When ``args.sft`` is True, the dataset emits THD-format batches where - multiple sequences are concatenated into a single flat token tensor. - The batch includes ``cu_seqlens`` (cumulative sequence lengths, shape - ``[1, S+1]``) and ``max_seqlen`` (shape ``[1]``) that describe the - individual sequence boundaries. - - This function validates and squeezes those fields: - - ``cu_seqlens``: asserted to have shape ``[1, S+1]`` (micro-batch - size must be 1 for packing), then squeezed to ``[S+1]``. - - ``max_seqlen``: asserted to be 1-D; kept as a tensor and passed - to ``get_thd_batch_on_this_cp_rank`` which performs the final - scalar conversion internally. - - Pipeline stage handling: - - First/last PP stages: fetch the full batch (tokens + labels) and - route through ``get_thd_batch_on_this_cp_rank`` to produce a - ``PackedSeqParams`` object that carries ``cu_seqlens`` and - ``max_seqlen`` to the attention kernel. - - Middle PP stages: only ``cu_seqlens`` and ``max_seqlen`` are - needed for attention masking; all other fields are returned as - ``None`` with a ``PackedSeqParams`` built directly here. - - MTP ranks (``mtp_on_this_rank``) also receive the full batch, - regardless of pipeline stage. - - Difference from ``pretrain_mamba.py``: - - Return format: GPT returns a 6-tuple - ``(tokens, labels, loss_mask, attention_mask, position_ids, - packed_seq_params)`` where ``packed_seq_params`` is a - ``PackedSeqParams`` dataclass. Mamba returns 7 values via - ``batch.values()`` with ``cu_seqlens`` and ``max_seqlen`` as - separate dict entries (no ``PackedSeqParams`` wrapper). - - Middle-stage return: GPT returns ``(None×5, PackedSeqParams)``; - Mamba returns an ``empty_batch`` dict with ``cu_seqlens`` and - ``max_seqlen`` set. - - CP with packed sequences: GPT delegates to - ``get_thd_batch_on_this_cp_rank`` (MCore utility); Mamba - implements the ``tex.thd_get_partitioned_indices`` CP slicing - inline and does not call that helper. - - MTP: GPT passes ``mtp_on_this_rank`` to ``get_batch_on_this_tp_rank`` - and uses it to gate the early-return; Mamba has no MTP support. - - ``max_seqlen`` conversion: Mamba converts to a Python int scalar - before returning (``int(max_seqlen[0].item())``); GPT keeps it as - a tensor and lets ``get_thd_batch_on_this_cp_rank`` convert it, - except for the middle-stage ``PackedSeqParams`` where conversion - is done inline. - """ + """Generate a batch.""" + + BATCH_KEYS = ["tokens", "labels", "loss_mask", "position_ids", "attention_mask", "cu_seqlens", "cu_seqlens_padded", "max_seqlen", "local_cp_size", "hybrid_cp_group"] + args = get_args() config = core_transformer_config_from_args(args) - # TODO: this is pretty hacky, find a better way - is_packed_sequence = get_args().sft # SFT always uses packed sequence - if not is_first_or_last_pipeline_stage(vp_stage) and not is_packed_sequence and ( - (not mtp_on_this_rank(config, ignore_virtual=False, vp_stage=vp_stage))): - return None, None, None, None, None, None - - # get batches based on the TP rank you are on - batch = get_batch_on_this_tp_rank( - data_iterator, - mtp_on_this_rank=mtp_on_this_rank(config, ignore_virtual=False, vp_stage=vp_stage) - ) - cu_seqlens = batch.pop('cu_seqlens', None) - cu_seqlens_padded = batch.pop('cu_seqlens_padded', None) - max_seqlen = batch.pop('max_seqlen', None) - local_cp_size = batch.pop('local_cp_size', None) - if local_cp_size is not None: - local_cp_size = int(local_cp_size.item()) + 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 + create_attention_mask_in_dataloader = args.create_attention_mask_in_dataloader + mtp_on_this_rank = mtp_on_this_rank_func(layout=config.pipeline_model_parallel_layout, mtp_num_layers=config.mtp_num_layers, ignore_virtual=False, vp_stage=vp_stage) + is_hybrid_cp = args.hybrid_context_parallel - if cu_seqlens is not None: - assert ( - cu_seqlens.dim() == 2 and cu_seqlens.shape[0] == 1 - ), "micro-batch-size must be 1 for packing" - cu_seqlens = cu_seqlens[0] - assert max_seqlen.dim() == 1 - - # For middle pipeline stages with packed sequences, only cu_seqlens and - # max_seqlen are needed (for attention masking); skip the full batch. - if not is_first_or_last_pipeline_stage(vp_stage) and is_packed_sequence: - return None, None, None, None, None, PackedSeqParams( - cu_seqlens_q=cu_seqlens, - cu_seqlens_kv=cu_seqlens, - max_seqlen_q=int(max_seqlen[0].item()), - max_seqlen_kv=int(max_seqlen[0].item()), - qkv_format='thd', - ) + if not is_first_or_last_pipeline_stage(vp_stage) and not mtp_on_this_rank: + return [None for _ in BATCH_KEYS] - if cu_seqlens is None and local_cp_size is None: - # slice batch along sequence dimension for context parallelism - batch = get_batch_on_this_cp_rank(batch) # The implementation of this function is in MCore - packed_seq_params = None - elif local_cp_size is None: # Packed THD format - batch, packed_seq_params = get_thd_batch_on_this_cp_rank(batch, cu_seqlens, cu_seqlens_padded, max_seqlen) - else: # Hybrid CP format - batch, packed_seq_params = get_batch_on_this_hybrid_cp_rank(batch, local_cp_size) + 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) if key in batch and batch[key] is not None else None - return (*batch.values(), packed_seq_params) + batch = get_batch_on_this_tp_rank(batch, broadcast_src_rank=mpu.get_tensor_model_parallel_src_rank(), broadcast_group=mpu.get_tensor_model_parallel_group(), is_sft=is_sft, is_hybrid_cp=is_hybrid_cp, create_attention_mask_in_dataloader=create_attention_mask_in_dataloader, cp_size=cp_size, tp_rank=tp_rank, micro_batch_size=args.micro_batch_size, seq_length=args.seq_length, mtp_on_this_rank=mtp_on_this_rank, pipeline_model_parallel_size=args.pipeline_model_parallel_size, is_pipeline_first_stage=mpu.is_pipeline_first_stage(), is_pipeline_last_stage=mpu.is_pipeline_last_stage()) + batch = get_batch_on_this_cp_rank(batch, is_hybrid_cp=is_hybrid_cp, cp_group=get_context_parallel_group(), hybrid_cp_group_func=get_hybrid_data_context_parallel_groups) + return [batch[key] for key in sorted(batch.keys())] # define spiky loss as a loss that's 10x the max loss observed @@ -245,7 +178,34 @@ def forward_step(data_iterator, model: GPTModel, return_schedule_plan: bool = Fa global stimer with stimer(bdata=True): vp_stage = get_attr_wrapped_model(model, "vp_stage") - tokens, labels, loss_mask, attention_mask, position_ids, packed_seq_params = get_batch(data_iterator, vp_stage) + ( + attention_mask, + cu_seqlens, + cu_seqlens_padded, + hybrid_cp_group, + labels, + local_cp_size, + loss_mask, + max_seqlen, + position_ids, + tokens, + ) = get_batch(data_iterator, vp_stage) + + packed_seq_params = None + if cu_seqlens is not None: + cu_seqlens_for_params = cu_seqlens_padded if cu_seqlens_padded is not None else cu_seqlens + packed_seq_params = PackedSeqParams( + qkv_format="thd", + cu_seqlens_q=cu_seqlens_for_params, + cu_seqlens_kv=cu_seqlens_for_params, + cu_seqlens_q_padded=cu_seqlens_padded, + cu_seqlens_kv_padded=cu_seqlens_padded, + max_seqlen_q=max_seqlen, + max_seqlen_kv=max_seqlen, + local_cp_size=local_cp_size, + cp_group=hybrid_cp_group, + ) + timers('batch-generator').stop() with stimer: @@ -271,13 +231,13 @@ def forward_step(data_iterator, model: GPTModel, return_schedule_plan: bool = Fa def is_dataset_built_on_rank(vp_stage=None, is_packed_sequence=False): args = get_args() config = core_transformer_config_from_args(args) - if parallel_state.get_tensor_model_parallel_rank() != 0: + if mpu.get_tensor_model_parallel_rank() != 0: return False elif is_packed_sequence: return True return ( is_first_or_last_pipeline_stage(vp_stage) - or mtp_on_this_rank(config, ignore_virtual=False, vp_stage=vp_stage) + or mtp_on_this_rank_func(layout=config.pipeline_model_parallel_layout, mtp_num_layers=config.mtp_num_layers, ignore_virtual=False, vp_stage=vp_stage) ) @@ -343,6 +303,9 @@ def core_gpt_dataset_config_from_args(args): ) return GPTFIMDatasetConfig(**data_args) + if args.sft: + return SFTDatasetConfig(**data_args) + return GPTDatasetConfig(**data_args) diff --git a/pretrain_mamba.py b/pretrain_mamba.py index a30979af714..e200c6f9260 100644 --- a/pretrain_mamba.py +++ b/pretrain_mamba.py @@ -26,26 +26,24 @@ from megatron.core.datasets.gpt_dataset import GPTDataset, GPTDatasetConfig, MockGPTDataset from megatron.core.enums import ModelType from megatron.core.packed_seq_params import PackedSeqParams -from megatron.core.parallel_state import ( - get_context_parallel_rank, - get_context_parallel_world_size, -) +from megatron.core.parallel_state import get_context_parallel_group, get_hybrid_data_context_parallel_groups from megatron.core.models.mamba import MambaModel from megatron.core.rerun_state_machine import get_rerun_state_machine from megatron.core.tokenizers.utils.build_tokenizer import build_tokenizer -from megatron.core.utils import get_attr_wrapped_model, is_te_min_version, StragglerDetector +from megatron.core.utils import get_attr_wrapped_model, StragglerDetector, preprocess_sft_batch, get_batch_on_this_cp_rank, get_batch_on_this_tp_rank from megatron.training import ( get_args, get_timers, + get_tokenizer, inprocess_restart, pretrain, print_rank_0, set_startup_timestamps, ) -from megatron.training.datasets.sft_dataset import SFTDataset +from megatron.core.transformer.multi_token_prediction import mtp_on_this_rank as mtp_on_this_rank_func +from megatron.training.arguments import core_transformer_config_from_args +from megatron.core.datasets.sft_dataset import SFTDataset, SFTDatasetConfig, IGNORE_INDEX from megatron.training.utils import ( - get_batch_on_this_cp_rank, - get_batch_on_this_tp_rank, get_blend_and_blend_per_split, is_first_or_last_pipeline_stage, ) @@ -58,92 +56,41 @@ except ImportError: has_nvidia_modelopt = False -try: - # Register the TE CUDA kernels - import transformer_engine # pylint: disable=unused-import - - # Alias the PyTorch wrapper so we can call tex.* APIs - import transformer_engine_torch as tex -except ImportError: - # TE isn’t installed or the torch wrapper is missing - tex = None - stimer = StragglerDetector() def get_batch(data_iterator, vp_stage=None): """Generate a batch.""" - empty_batch = { - 'tokens': None, - 'labels': None, - 'loss_mask': None, - 'attention_mask': None, - 'position_ids': None, - 'cu_seqlens': None, - 'max_seqlen': None, - } - - # TODO(duncan): Is there a more efficient way to access is_packed_sequence here? - is_packed_sequence = get_args().sft # SFT always uses packed sequence - if not is_first_or_last_pipeline_stage(vp_stage) and not is_packed_sequence: - return empty_batch.values() + BATCH_KEYS = ["tokens", "labels", "loss_mask", "position_ids", "attention_mask", "cu_seqlens", "cu_seqlens_padded", "max_seqlen", "local_cp_size", "hybrid_cp_group"] - batch = get_batch_on_this_tp_rank(data_iterator) - - cu_seqlens = batch['cu_seqlens'] - # Unused at the moment - cu_seqlens_padded = batch.pop('cu_seqlens_padded', None) - # Support for Hybrid Context Parallel (Unused in this script) - local_cp_size = batch.pop('local_cp_size', None) - - if cu_seqlens is not None: - assert ( - cu_seqlens.dim() == 2 and cu_seqlens.shape[0] == 1 - ), "micro-batch-size must be 1 for packing" - cu_seqlens = cu_seqlens[0] - batch['cu_seqlens'] = cu_seqlens - - max_seqlen = batch['max_seqlen'] - assert max_seqlen.dim() == 1 - # TODO(duncan): can this be kept as a 0-D tensor? - batch['max_seqlen'] = int(max_seqlen[0].item()) - - if mpu.is_pipeline_first_stage(ignore_virtual=(vp_stage is None), vp_stage=vp_stage): - total_tokens = batch['tokens'].size(1) - elif mpu.is_pipeline_last_stage(ignore_virtual=(vp_stage is None), vp_stage=vp_stage): - total_tokens = batch['labels'].size(1) - else: # packed sequence - empty_batch['cu_seqlens'] = cu_seqlens - empty_batch['max_seqlen'] = max_seqlen - return empty_batch.values() - - if cu_seqlens is None: - # slice batch along sequence dimension for context parallelism - batch = get_batch_on_this_cp_rank(batch) # The implementation of this function is in MCore - else: # Packed THD format - cp_size = get_context_parallel_world_size() - if cp_size > 1: # slice batch along sequence dimension for context parallelism - assert tex is not None and is_te_min_version("1.10.0"), ( - "Please update Transformer Engine to >= 1.10 to use " - "Context Parallel with THD format data" - ) - cp_rank = get_context_parallel_rank() - index = tex.thd_get_partitioned_indices( - cu_seqlens, - total_tokens, - cp_size, - cp_rank, - ) - for key, data in batch.items(): - if key in {'attention_mask', 'cu_seqlens', 'max_seqlen'}: - continue - if data is not None: - # On first PP rank, labels and loss_mask can be None. - # On last PP rank, tokens and position_ids can be None. - batch[key] = data.index_select(1, index) - - return batch.values() + args = get_args() + 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 + create_attention_mask_in_dataloader = args.create_attention_mask_in_dataloader + mtp_on_this_rank = mtp_on_this_rank_func(layout=config.pipeline_model_parallel_layout, mtp_num_layers=config.mtp_num_layers, ignore_virtual=False, vp_stage=vp_stage) + is_hybrid_cp = args.hybrid_context_parallel + + if not is_first_or_last_pipeline_stage(vp_stage) and not mtp_on_this_rank: + return [None for _ in BATCH_KEYS] + + 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) if key in batch and batch[key] is not None else None + + batch = get_batch_on_this_tp_rank(batch, broadcast_src_rank=mpu.get_tensor_model_parallel_src_rank(), broadcast_group=mpu.get_tensor_model_parallel_group(), is_sft=is_sft, is_hybrid_cp=is_hybrid_cp, create_attention_mask_in_dataloader=create_attention_mask_in_dataloader, cp_size=cp_size, tp_rank=tp_rank, micro_batch_size=args.micro_batch_size, seq_length=args.seq_length, mtp_on_this_rank=mtp_on_this_rank, pipeline_model_parallel_size=args.pipeline_model_parallel_size, is_pipeline_first_stage=mpu.is_pipeline_first_stage(), is_pipeline_last_stage=mpu.is_pipeline_last_stage()) + batch = get_batch_on_this_cp_rank(batch, is_hybrid_cp=is_hybrid_cp, cp_group=get_context_parallel_group(), hybrid_cp_group_func=get_hybrid_data_context_parallel_groups) + return [batch[key] for key in sorted(batch.keys())] # define spiky loss as a loss that's 10x the max loss observed @@ -212,7 +159,7 @@ def forward_step(data_iterator, model: MambaModel): Args: data_iterator : Input data iterator - model (MambaModel): The GPT Model + model (MambaModel): The Mamba Model """ timers = get_timers() @@ -224,28 +171,31 @@ def forward_step(data_iterator, model: MambaModel): with stimer(bdata=True): vp_stage = get_attr_wrapped_model(model, "vp_stage") ( - tokens, - labels, - loss_mask, attention_mask, - position_ids, cu_seqlens, + cu_seqlens_padded, + hybrid_cp_group, + labels, + local_cp_size, + loss_mask, max_seqlen, + position_ids, + tokens, ) = get_batch(data_iterator, vp_stage) - if cu_seqlens is None: - packed_seq_params = None - else: - total_tokens = tokens.size(1) if tokens is not None else labels.size(1) + packed_seq_params = None + if cu_seqlens is not None: + cu_seqlens_for_params = cu_seqlens_padded if cu_seqlens_padded is not None else cu_seqlens packed_seq_params = PackedSeqParams( qkv_format="thd", - cu_seqlens_q=cu_seqlens, - cu_seqlens_kv=cu_seqlens, - cu_seqlens_q_padded=None, - cu_seqlens_kv_padded=None, + cu_seqlens_q=cu_seqlens_for_params, + cu_seqlens_kv=cu_seqlens_for_params, + cu_seqlens_q_padded=cu_seqlens_padded, + cu_seqlens_kv_padded=cu_seqlens_padded, max_seqlen_q=max_seqlen, max_seqlen_kv=max_seqlen, - total_tokens=total_tokens, + local_cp_size=local_cp_size, + cp_group=hybrid_cp_group, ) timers('batch-generator').stop() @@ -265,12 +215,16 @@ def forward_step(data_iterator, model: MambaModel): def is_dataset_built_on_rank(vp_stage=None, is_packed_sequence=False): + args = get_args() + config = core_transformer_config_from_args(args) if mpu.get_tensor_model_parallel_rank() != 0: return False elif is_packed_sequence: return True - else: - return is_first_or_last_pipeline_stage(vp_stage) + return ( + is_first_or_last_pipeline_stage(vp_stage) + or mtp_on_this_rank_func(layout=config.pipeline_model_parallel_layout, mtp_num_layers=config.mtp_num_layers, ignore_virtual=False, vp_stage=vp_stage) + ) def core_gpt_dataset_config_from_args(args): @@ -286,28 +240,38 @@ def core_gpt_dataset_config_from_args(args): 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, - 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_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, + } + + 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_builder.py b/tests/unit_tests/data/test_builder.py index 59f73911c82..536235c69bc 100644 --- a/tests/unit_tests/data/test_builder.py +++ b/tests/unit_tests/data/test_builder.py @@ -53,7 +53,7 @@ def create_file_prefixes(tokenizer, number_of_files, maximum_number_of_documents file_prefix_path + ".bin", dtype=DType.optimal_dtype(tokenizer.vocab_size) ) number_of_documents = random.randint(10, maximum_number_of_documents) - for j in range(number_of_documents): + for _ in range(number_of_documents): number_of_tokens = random.randint(50, 100) tokenized_doc = [ str(random.randint(0, tokenizer.vocab_size - 1)) for _ in range(number_of_tokens) diff --git a/tests/unit_tests/data/test_get_batch.py b/tests/unit_tests/data/test_get_batch.py new file mode 100644 index 00000000000..07ddbec15fa --- /dev/null +++ b/tests/unit_tests/data/test_get_batch.py @@ -0,0 +1,587 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +import os +import sys + +import pytest +import torch + +from megatron.core import mpu +from megatron.core.num_microbatches_calculator import destroy_num_microbatches_calculator +from megatron.training.arguments import parse_args, validate_args +from megatron.training.global_vars import destroy_global_vars, set_global_variables +from pretrain_mamba import get_batch +from tests.unit_tests.test_utilities import Utils + + +def initialize_test_environment( + tp_size: int, + cp_size: int, + seq_length: int, + micro_batch_size: int, + global_batch_size: int = 1, + sft: bool = False, + hybrid_context_parallel: bool = False, + max_seqlen_per_cp_rank: int = 1024, + create_attention_mask: bool = False, +): + destroy_global_vars() + destroy_num_microbatches_calculator() + + sys.argv = ['test_get_batch.py'] + args = parse_args() + args.seq_length = seq_length + args.tensor_model_parallel_size = tp_size + args.sequence_parallel = True if tp_size > 1 else False + args.pipeline_model_parallel_size = 1 + args.context_parallel_size = cp_size + args.hybrid_context_parallel = hybrid_context_parallel + args.max_seqlen_per_cp_rank = max_seqlen_per_cp_rank + args.sft = sft + args.micro_batch_size = micro_batch_size + args.create_attention_mask_in_dataloader = create_attention_mask + args.global_batch_size = global_batch_size + args.calculate_per_token_loss = True + args.vocab_size = 1024 + args.tokenizer_type = "NullTokenizer" + args.num_layers = 4 + args.hidden_size = 512 + args.num_attention_heads = 8 + args.max_position_embeddings = seq_length + + os.environ['CUDA_DEVICE_MAX_CONNECTIONS'] = '1' + os.environ['NCCL_NVLS_ENABLE'] = '0' # NOTE(asolergi-nv): Without this, NCCL crashes + + validate_args(args) + set_global_variables(args, True) + + Utils.initialize_model_parallel( + tensor_model_parallel_size=tp_size, + context_parallel_size=cp_size, + hybrid_context_parallel=hybrid_context_parallel, + ) + return args + + +def create_sft_data_iterator(max_seq_length: int = 1024): + min_len = max(1, int(0.1 * max_seq_length)) + max_len = max(2, int(0.4 * max_seq_length)) + candidate_lengths = [torch.randint(min_len, max_len + 1, (1,)).item() for _ in range(10)] + + lengths = [] + total = 0 + for l in candidate_lengths: + if total + l >= max_seq_length: + break + lengths.append(l) + total += l + + num_real_tokens = sum(lengths) + assert ( + num_real_tokens < max_seq_length + ), f"Sum of lengths {num_real_tokens} is greater than max_seq_length {max_seq_length}" + text = torch.randint(0, 10000, (1, num_real_tokens + 1), dtype=torch.int64) + tokens = text[:, :-1].contiguous() + labels = text[:, 1:].contiguous() + + cu_seqlens = torch.cat( + ( + torch.zeros(1, dtype=torch.int32), + torch.cumsum(torch.tensor(lengths, dtype=torch.int64), dim=0).to(torch.int32), + ) + ) + + loss_mask = torch.ones((1, num_real_tokens), dtype=torch.float) + batch = {"tokens": tokens, "labels": labels, "loss_mask": loss_mask, "cu_seqlens": cu_seqlens} + return iter([batch]), num_real_tokens + + +@pytest.mark.parametrize("tp_size", [1, 2, 4]) +@pytest.mark.parametrize("cp_size", [1, 2, 4]) +@pytest.mark.parametrize("seq_length", [1024, 2048, 4096]) +def test_sft_batch(tp_size, cp_size, seq_length): + if cp_size * tp_size > torch.cuda.device_count(): + pytest.skip( + f"Skipping test because cp_size * tp_size > torch.cuda.device_count() ({cp_size * tp_size} > {torch.cuda.device_count()})" + ) + + global_batch_size = int(os.environ.get("WORLD_SIZE", 1)) // (tp_size * cp_size) + initialize_test_environment( + tp_size, + cp_size, + seq_length, + micro_batch_size=1, + global_batch_size=global_batch_size, + sft=True, + ) + + data_iterator = None + num_real_tokens = 0 + if mpu.get_tensor_model_parallel_rank() == 0: + data_iterator, num_real_tokens = create_sft_data_iterator(seq_length) + + ( + attention_mask, + cu_seqlens, + cu_seqlens_padded, + hybrid_cp_group, + labels, + local_cp_size, + loss_mask, + max_seqlen, + position_ids, + tokens, + ) = get_batch(data_iterator) + + # Presence checks + assert tokens is not None + assert labels is not None + assert loss_mask is not None + assert position_ids is not None + assert cu_seqlens is not None + assert max_seqlen is not None + assert attention_mask is None + assert hybrid_cp_group is None + assert local_cp_size is None + + # Shape: preprocess_sft_batch pads to seq_length; THD CP slicing gives seq_length // cp_size per rank + seq_len_per_rank = seq_length // cp_size + assert tokens.shape == ( + 1, + seq_len_per_rank, + ), f"Expected tokens shape (1, {seq_len_per_rank}), got {tokens.shape}" + assert labels.shape == ( + 1, + seq_len_per_rank, + ), f"Expected labels shape (1, {seq_len_per_rank}), got {labels.shape}" + assert loss_mask.shape == ( + 1, + seq_len_per_rank, + ), f"Expected loss_mask shape (1, {seq_len_per_rank}), got {loss_mask.shape}" + assert position_ids.shape == ( + 1, + seq_len_per_rank, + ), f"Expected position_ids shape (1, {seq_len_per_rank}), got {position_ids.shape}" + + # Dtype checks + assert tokens.dtype == torch.int64 + assert labels.dtype == torch.int64 + assert loss_mask.dtype == torch.float32 + assert position_ids.dtype == torch.int64 + + # cu_seqlens: 1D int32, starts at 0, ends at seq_length (padded by preprocess_sft_batch) + assert cu_seqlens.dim() == 1 + assert cu_seqlens.dtype == torch.int32 + assert cu_seqlens[0].item() == 0 + assert cu_seqlens[-1].item() == seq_length + assert cu_seqlens.shape[0] >= 2 # at least one sequence + + # max_seqlen: scalar, positive, within seq_length + assert max_seqlen.shape == (1,) + assert max_seqlen.dtype == torch.int32 + assert 0 < max_seqlen.item() <= seq_length + + if cp_size > 1: + assert cu_seqlens_padded is not None + assert cu_seqlens_padded.dim() == 1 + assert cu_seqlens_padded.dtype == torch.int32 + assert cu_seqlens_padded[0].item() == 0 + assert cu_seqlens_padded[-1].item() == seq_length + assert cu_seqlens_padded.shape == cu_seqlens.shape + + # Compute the divisibility factor (mirrors preprocess_sft_batch logic) + sp = tp_size > 1 + divisibility_factor = cp_size * 2 + if tp_size > 1 and sp: + divisibility_factor *= tp_size + + # Compute the segment lengths from cu_seqlens and cu_seqlens_padded + orig_seg_lengths = cu_seqlens[1:] - cu_seqlens[:-1] + padded_seg_lengths = cu_seqlens_padded[1:] - cu_seqlens_padded[:-1] + num_segments = len(orig_seg_lengths) + + # cu_seqlens_padded: every segment must be divisible by divisibility_factor + for i, seg_len in enumerate(padded_seg_lengths): + assert ( + seg_len.item() % divisibility_factor == 0 + ), f"Padded segment {i} length {seg_len.item()} not divisible by {divisibility_factor}" + + # cu_seqlens_padded segments >= cu_seqlens segments (padding only adds tokens). + # The last segment is excluded because pad_or_truncate_thd_tensors replaces + # the final entry of both cu_seqlens and cu_seqlens_padded with seq_length. + # Since cu_seqlens_padded[-2] >= cu_seqlens[-2] (from CP padding), the last + # padded segment (seq_length - cu_seqlens_padded[-2]) can be smaller than the + # last original segment (seq_length - cu_seqlens[-2]). + for i in range(num_segments - 1): + assert ( + padded_seg_lengths[i].item() >= orig_seg_lengths[i].item() + ), f"Segment {i}: padded length {padded_seg_lengths[i].item()} < original length {orig_seg_lengths[i].item()}" + + # loss_mask: must be binary (0.0 or 1.0) + assert ((loss_mask == 0.0) | (loss_mask == 1.0)).all(), "loss_mask must be binary" + + # Intra-sample CP padding validation: for each padded segment on this + # CP rank, verify that position_ids and loss_mask are consistent with + # the padding introduced by pad_thd_sequences_for_cp. + cp_rank = mpu.get_context_parallel_rank() + per_rank_seg_lens = (padded_seg_lengths // cp_size).tolist() + + offset = 0 + for i in range(num_segments): + seg_len = per_rank_seg_lens[i] + if seg_len == 0: + continue + seg_pos = position_ids[0, offset : offset + seg_len] + seg_loss = loss_mask[0, offset : offset + seg_len] + + # thd_get_partitioned_indices uses zigzag load-balanced partitioning: + # each padded segment is split into 2*cp_size chunks, and rank k + # gets chunk k and chunk (2*cp_size - 1 - k). + padded_seg_len = padded_seg_lengths[i].item() + num_chunks = 2 * cp_size + chunk_size = padded_seg_len // num_chunks + chunk0_start = cp_rank * chunk_size + chunk1_start = (num_chunks - 1 - cp_rank) * chunk_size + expected_pos = torch.cat( + [ + torch.arange( + chunk0_start, + chunk0_start + chunk_size, + dtype=torch.int64, + device=seg_pos.device, + ), + torch.arange( + chunk1_start, + chunk1_start + chunk_size, + dtype=torch.int64, + device=seg_pos.device, + ), + ] + ) + assert torch.equal(seg_pos, expected_pos), ( + f"Segment {i}: expected zigzag position_ids " + f"[{chunk0_start}..{chunk0_start + chunk_size - 1}, " + f"{chunk1_start}..{chunk1_start + chunk_size - 1}] but got " + f"[{seg_pos[0].item()}, ..., {seg_pos[-1].item()}] on CP rank {cp_rank}" + ) + + # For non-last segments, cu_seqlens entries are unchanged by + # pad_or_truncate_thd_tensors, so orig_seg_lengths[i] is the true + # sub-sequence length. This lets us make precise assertions: + # position_id >= orig_len => intra-sample CP padding => loss_mask == 0 + # position_id < orig_len => real token => loss_mask == 1 + # (The last segment absorbs end-of-sequence padding so its + # orig_seg_lengths entry is inflated -- skip the precise check.) + if i < num_segments - 1: + orig_len = orig_seg_lengths[i].item() + padding_mask = seg_pos >= orig_len + if padding_mask.any(): + assert (seg_loss[padding_mask] == 0.0).all(), ( + f"Segment {i}: intra-sample padding tokens (pos >= {orig_len}) " + f"must have loss_mask=0, CP rank {cp_rank}" + ) + real_mask = seg_pos < orig_len + if real_mask.any(): + assert (seg_loss[real_mask] == 1.0).all(), ( + f"Segment {i}: real tokens (pos < {orig_len}) " + f"must have loss_mask=1, CP rank {cp_rank}" + ) + + offset += seg_len + + assert ( + offset == seq_len_per_rank + ), f"Total per-rank offset {offset} != expected {seq_len_per_rank}" + else: + assert cu_seqlens_padded is None + + Utils.destroy_model_parallel() + + +def create_pretrain_data_iterator( + seq_length: int = 1024, micro_batch_size: int = 1, create_attention_mask: bool = False +): + text = torch.randint(0, 10000, (micro_batch_size, seq_length + 1), dtype=torch.int64) + tokens = text[:, :-1].contiguous() + labels = text[:, 1:].contiguous() + position_ids = ( + torch.arange(seq_length, dtype=torch.long) + .unsqueeze(0) + .expand(micro_batch_size, -1) + .contiguous() + ) + loss_mask = torch.ones((micro_batch_size, seq_length), dtype=torch.float) + + batch = { + "tokens": tokens, + "labels": labels, + "loss_mask": loss_mask, + "position_ids": position_ids, + } + + if create_attention_mask: + batch["attention_mask"] = torch.tril( + torch.ones((micro_batch_size, 1, seq_length, seq_length)) + ).bool() + + return iter([batch]) + + +@pytest.mark.parametrize("tp_size", [1, 2, 4]) +@pytest.mark.parametrize("cp_size", [1, 2, 4]) +@pytest.mark.parametrize("seq_length", [1024, 2048, 4096]) +@pytest.mark.parametrize("create_attention_mask", [True, False]) +@pytest.mark.parametrize("micro_batch_size", [1, 2, 4]) +def test_pretrain_batch(tp_size, cp_size, seq_length, create_attention_mask, micro_batch_size): + if cp_size * tp_size > torch.cuda.device_count(): + pytest.skip( + f"Skipping test because cp_size * tp_size > torch.cuda.device_count() ({cp_size * tp_size} > {torch.cuda.device_count()})" + ) + dp_size = int(os.environ.get("WORLD_SIZE", 1)) // (tp_size * cp_size) + global_batch_size = micro_batch_size * dp_size + initialize_test_environment( + tp_size, + cp_size, + seq_length, + micro_batch_size, + global_batch_size=global_batch_size, + sft=False, + create_attention_mask=create_attention_mask, + ) + + data_iterator = None + if mpu.get_tensor_model_parallel_rank() == 0: + data_iterator = create_pretrain_data_iterator( + seq_length, + micro_batch_size=micro_batch_size, + create_attention_mask=create_attention_mask, + ) + + ( + attention_mask, + cu_seqlens, + cu_seqlens_padded, + hybrid_cp_group, + labels, + local_cp_size, + loss_mask, + max_seqlen, + position_ids, + tokens, + ) = get_batch(data_iterator) + + # Presence checks + assert tokens is not None + assert labels is not None + 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 + assert local_cp_size is None + + # Shape: pretrain CP slicing takes 2 non-contiguous chunks → seq_length // cp_size tokens per rank + seq_len_per_rank = seq_length // cp_size + assert tokens.shape == ( + micro_batch_size, + seq_len_per_rank, + ), f"Expected tokens shape ({micro_batch_size}, {seq_len_per_rank}), got {tokens.shape}" + assert labels.shape == ( + micro_batch_size, + seq_len_per_rank, + ), f"Expected labels shape ({micro_batch_size}, {seq_len_per_rank}), got {labels.shape}" + assert loss_mask.shape == ( + micro_batch_size, + seq_len_per_rank, + ), f"Expected loss_mask shape ({micro_batch_size}, {seq_len_per_rank}), got {loss_mask.shape}" + assert position_ids.shape == ( + micro_batch_size, + seq_len_per_rank, + ), f"Expected position_ids shape ({micro_batch_size}, {seq_len_per_rank}), got {position_ids.shape}" + + # Dtype checks + assert tokens.dtype == torch.int64 + assert labels.dtype == torch.int64 + assert loss_mask.dtype == torch.float32 + assert position_ids.dtype == torch.int64 + + # Pretrain loss_mask is all-ones (no masking in the dataloader) + assert loss_mask.sum().item() == micro_batch_size * seq_len_per_rank + + if create_attention_mask: + assert attention_mask is not None + # attention_mask input shape (B, 1, S, S); seq_dim=2 splits the query dim → (B, 1, S // cp_size, S) + assert attention_mask.shape == ( + micro_batch_size, + 1, + seq_len_per_rank, + seq_length, + ), f"Expected attention_mask shape ({micro_batch_size}, 1, {seq_len_per_rank}, {seq_length}), got {attention_mask.shape}" + assert attention_mask.dtype == torch.bool + else: + assert attention_mask is None + + Utils.destroy_model_parallel() + + +def create_hybrid_cp_data_iterator(seq_length: int = 1024, cp_size: int = 1): + # Pack n_seqs equal-length sequences; total length must be divisible by 2 * cp_size for CP splitting + n_seqs = max(2, 2 * cp_size) + align = max(1, 2 * cp_size) + seq_len_each = (seq_length // n_seqs // align) * align + if seq_len_each == 0: + seq_len_each = align + total_seq_len = n_seqs * seq_len_each + + text = torch.randint(0, 10000, (1, total_seq_len + 1), dtype=torch.int64) + tokens = text[:, :-1].contiguous() # (1, total_seq_len) + labels = text[:, 1:].contiguous() # (1, total_seq_len) + loss_mask = torch.ones((1, total_seq_len), dtype=torch.float32) + position_ids = torch.cat( + [torch.arange(seq_len_each, dtype=torch.int64) for _ in range(n_seqs)] + ).unsqueeze( + 0 + ) # (1, total_seq_len) + + cu_seqlens = torch.cat( + [ + torch.zeros(1, dtype=torch.int32), + torch.cumsum(torch.tensor([seq_len_each] * n_seqs, dtype=torch.int64), dim=0).to( + torch.int32 + ), + ] + ) + max_seqlen = torch.tensor([seq_len_each], dtype=torch.int32) + local_cp_size_tensor = torch.tensor([cp_size], dtype=torch.int32) + + batch = { + "tokens": tokens, + "labels": labels, + "loss_mask": loss_mask, + "position_ids": position_ids, + "cu_seqlens": cu_seqlens, + "max_seqlen": max_seqlen, + "local_cp_size": local_cp_size_tensor, + } + + if cp_size > 1: + batch["cu_seqlens_padded"] = cu_seqlens.clone() + + return iter([batch]) + + +@pytest.mark.parametrize("tp_size", [1, 2, 4]) +@pytest.mark.parametrize("cp_size", [2, 4, 8]) +@pytest.mark.parametrize("seq_length", [1024]) +@pytest.mark.parametrize("create_attention_mask", [False]) +def test_hybrid_cp_batch(tp_size, cp_size, seq_length, create_attention_mask): + if cp_size * tp_size > torch.cuda.device_count(): + pytest.skip( + f"Skipping test because cp_size * tp_size > torch.cuda.device_count() ({cp_size * tp_size} > {torch.cuda.device_count()})" + ) + + initialize_test_environment( + tp_size, + cp_size, + seq_length, + 1, + 16, + sft=False, + hybrid_context_parallel=True, + create_attention_mask=create_attention_mask, + ) + + data_iterator = None + if mpu.get_tensor_model_parallel_rank() == 0: + data_iterator = create_hybrid_cp_data_iterator(seq_length, cp_size=cp_size) + + ( + attention_mask, + cu_seqlens, + cu_seqlens_padded, + hybrid_cp_group, + labels, + local_cp_size, + loss_mask, + max_seqlen, + position_ids, + tokens, + ) = get_batch(data_iterator) + + # Presence checks + assert tokens is not None + assert labels is not None + assert loss_mask is not None + assert position_ids is not None + assert attention_mask is None # HybridCP does not use attention mask from dataloader + assert cu_seqlens is not None # HybridCP always has cu_seqlens + assert max_seqlen is not None # HybridCP always has max_seqlen + assert local_cp_size is not None # HybridCP always has local_cp_size + + # Data iterator parameters (must match create_hybrid_cp_data_iterator) + n_seqs = max(2, 2 * cp_size) + align = max(1, 2 * cp_size) + seq_len_each = (seq_length // n_seqs // align) * align + if seq_len_each == 0: + seq_len_each = align + total_seq_len = n_seqs * seq_len_each # equals seq_length for the test parameters + + # Shape: HybridCP CP splitting gives total_seq_len // cp_size tokens per rank + seq_len_per_rank = total_seq_len // cp_size + assert tokens.shape == ( + 1, + seq_len_per_rank, + ), f"Expected tokens shape (1, {seq_len_per_rank}), got {tokens.shape}" + assert labels.shape == ( + 1, + seq_len_per_rank, + ), f"Expected labels shape (1, {seq_len_per_rank}), got {labels.shape}" + assert loss_mask.shape == ( + 1, + seq_len_per_rank, + ), f"Expected loss_mask shape (1, {seq_len_per_rank}), got {loss_mask.shape}" + assert position_ids.shape == ( + 1, + seq_len_per_rank, + ), f"Expected position_ids shape (1, {seq_len_per_rank}), got {position_ids.shape}" + + # Dtype checks + assert tokens.dtype == torch.int64 + assert labels.dtype == torch.int64 + assert loss_mask.dtype == torch.float32 + assert position_ids.dtype == torch.int64 + + # Loss mask is all-ones (no masking in the HybridCP pretrain dataloader) + assert loss_mask.sum().item() == seq_len_per_rank + + # cu_seqlens: 1D int32, [0, seq_len_each, 2*seq_len_each, ..., total_seq_len] + assert cu_seqlens.shape == ( + n_seqs + 1, + ), f"Expected cu_seqlens shape ({n_seqs + 1},), got {cu_seqlens.shape}" + assert cu_seqlens.dtype == torch.int32 + assert cu_seqlens[0].item() == 0 + assert cu_seqlens[-1].item() == total_seq_len + + # max_seqlen: scalar int32 equal to the per-sequence length in the iterator + assert max_seqlen.shape == (1,) + assert max_seqlen.dtype == torch.int32 + assert max_seqlen.item() == seq_len_each + + # local_cp_size: scalar int32 equal to cp_size + assert local_cp_size.shape == (1,) + assert local_cp_size.dtype == torch.int32 + assert local_cp_size.item() == cp_size + + if cp_size > 1: + assert cu_seqlens_padded is not None + assert cu_seqlens_padded.shape == (n_seqs + 1,) + assert cu_seqlens_padded.dtype == torch.int32 + assert cu_seqlens_padded[0].item() == 0 + assert cu_seqlens_padded[-1].item() == total_seq_len + assert hybrid_cp_group is not None + 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..5dc3bbf5c01 --- /dev/null +++ b/tests/unit_tests/data/test_sft_dataset.py @@ -0,0 +1,398 @@ +# 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, + ) + ) + + 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/models/test_mimo_partition.py b/tests/unit_tests/models/test_mimo_partition.py index 1527fb92935..3db10e81f51 100644 --- a/tests/unit_tests/models/test_mimo_partition.py +++ b/tests/unit_tests/models/test_mimo_partition.py @@ -401,7 +401,7 @@ def test_only_non_none_tensors_added_to_batch(self): sharded = {'embeddings': embeddings[:, :4, :]} captured = {} - def mock_fn(batch): + def mock_fn(batch, **kwargs): captured.update(batch) return sharded diff --git a/tests/unit_tests/pipeline_parallel/test_pipeline_layout.py b/tests/unit_tests/pipeline_parallel/test_pipeline_layout.py index e7fb8159af2..b84e29ec574 100644 --- a/tests/unit_tests/pipeline_parallel/test_pipeline_layout.py +++ b/tests/unit_tests/pipeline_parallel/test_pipeline_layout.py @@ -71,7 +71,12 @@ def initialize_gpt_model( else: layer_spec = layer_spec_fn() - if with_mtp and mtp_on_this_rank(transformer_config, ignore_virtual=False, vp_stage=i): + if with_mtp and mtp_on_this_rank( + layout=transformer_config.pipeline_model_parallel_layout, + mtp_num_layers=transformer_config.mtp_num_layers, + ignore_virtual=False, + vp_stage=i, + ): if is_moe: transformer_layer_spec_for_mtp = gpt_te_spec(transformer_config) else: diff --git a/tests/unit_tests/tokenizers/test_tokenizer.py b/tests/unit_tests/tokenizers/test_tokenizer.py index f38674c6329..60fab2e6071 100755 --- a/tests/unit_tests/tokenizers/test_tokenizer.py +++ b/tests/unit_tests/tokenizers/test_tokenizer.py @@ -246,12 +246,14 @@ 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 == 131073 + assert ids == [ord(c) % vocab_size for c in text] + assert tokenizer.vocab_size == vocab_size + 1 def test_bytelevel_tokenizer(): diff --git a/tests/unit_tests/transformer/moe/test_upcycling.py b/tests/unit_tests/transformer/moe/test_upcycling.py index 097356d6c78..533c9848a42 100644 --- a/tests/unit_tests/transformer/moe/test_upcycling.py +++ b/tests/unit_tests/transformer/moe/test_upcycling.py @@ -15,10 +15,19 @@ ) from megatron.core.models.gpt.gpt_model import GPTModel from megatron.core.num_microbatches_calculator import destroy_num_microbatches_calculator +from megatron.core.parallel_state import ( + get_context_parallel_group, + get_hybrid_data_context_parallel_groups, +) from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed from megatron.core.transformer.moe import upcycling_utils from megatron.core.transformer.moe.experts import SequentialMLP, TEGroupedMLP -from megatron.core.utils import get_te_version, is_te_min_version +from megatron.core.utils import ( + get_batch_on_this_cp_rank, + get_batch_on_this_tp_rank, + get_te_version, + is_te_min_version, +) from megatron.training.arguments import core_transformer_config_from_args, parse_args, validate_args from megatron.training.global_vars import ( destroy_global_vars, @@ -27,11 +36,7 @@ set_global_variables, ) from megatron.training.training import get_model, setup_model_and_optimizer -from megatron.training.utils import ( - get_batch_on_this_cp_rank, - get_batch_on_this_tp_rank, - unwrap_model, -) +from megatron.training.utils import unwrap_model from tests.unit_tests.test_utilities import Utils if HAVE_TE: @@ -162,8 +167,56 @@ def get_batch(data_iterator): if (not mpu.is_pipeline_first_stage()) and (not mpu.is_pipeline_last_stage()): return None, None, None, None, None - batch = get_batch_on_this_tp_rank(data_iterator) - batch = get_batch_on_this_cp_rank(batch) + args = get_args() + tp_rank = mpu.get_tensor_model_parallel_rank() + + BATCH_KEYS = [ + "tokens", + "labels", + "loss_mask", + "position_ids", + "attention_mask", + "cu_seqlens", + "cu_seqlens_padded", + "max_seqlen", + "local_cp_size", + "hybrid_cp_group", + ] + + batch = {} + if tp_rank == 0: + batch = next(data_iterator) + for key in BATCH_KEYS: + batch[key] = ( + batch[key].cuda(non_blocking=True) + if key in batch and batch[key] is not None + else None + ) + + batch = get_batch_on_this_tp_rank( + batch, + broadcast_src_rank=mpu.get_tensor_model_parallel_src_rank(), + broadcast_group=mpu.get_tensor_model_parallel_group(), + is_sft=False, + is_hybrid_cp=False, + create_attention_mask_in_dataloader=getattr( + args, 'create_attention_mask_in_dataloader', True + ), + cp_size=args.context_parallel_size, + tp_rank=tp_rank, + micro_batch_size=args.micro_batch_size, + seq_length=args.seq_length, + mtp_on_this_rank=False, + pipeline_model_parallel_size=args.pipeline_model_parallel_size, + is_pipeline_first_stage=mpu.is_pipeline_first_stage(), + is_pipeline_last_stage=mpu.is_pipeline_last_stage(), + ) + batch = get_batch_on_this_cp_rank( + batch, + is_hybrid_cp=False, + cp_group=get_context_parallel_group(), + hybrid_cp_group_func=get_hybrid_data_context_parallel_groups, + ) return batch.values() diff --git a/tests/unit_tests/transformer/test_multi_token_prediction.py b/tests/unit_tests/transformer/test_multi_token_prediction.py index 57423da335b..7aa4cd99a57 100644 --- a/tests/unit_tests/transformer/test_multi_token_prediction.py +++ b/tests/unit_tests/transformer/test_multi_token_prediction.py @@ -27,7 +27,7 @@ roll_tensor, ) from megatron.core.transformer.transformer_config import TransformerConfig -from megatron.core.utils import is_te_min_version +from megatron.core.utils import get_batch_on_this_cp_rank, is_te_min_version, unwrap_model from megatron.training.arguments import core_transformer_config_from_args, parse_args, validate_args from megatron.training.checkpointing import load_checkpoint, save_checkpoint from megatron.training.global_vars import ( @@ -37,7 +37,6 @@ set_global_variables, ) from megatron.training.training import get_model, setup_model_and_optimizer -from megatron.training.utils import get_batch_on_this_cp_rank, unwrap_model from tests.unit_tests.dist_checkpointing import TempNamedDir from tests.unit_tests.test_utilities import Utils @@ -400,7 +399,9 @@ def set_ckpt_path(ckpt_path): load_checkpoint(gpt_model, optimizer, opt_param_scheduler, strict=False) batch["output_ref"] = output_ref # Get batch for current CP rank (handles CP tensor splitting) - batch = get_batch_on_this_cp_rank(batch) + batch = get_batch_on_this_cp_rank( + batch, is_hybrid_cp=False, cp_group=get_context_parallel_group() + ) tokens, labels, loss_mask, attention_mask, position_ids, output_ref = batch.values() output = gpt_model[0].forward( input_ids=tokens, @@ -880,7 +881,9 @@ def set_ckpt_path(ckpt_path): load_checkpoint(mamba_model, optimizer, opt_param_scheduler, strict=False) batch["output_ref"] = output_ref - batch = get_batch_on_this_cp_rank(batch) + batch = get_batch_on_this_cp_rank( + batch, is_hybrid_cp=False, cp_group=get_context_parallel_group() + ) tokens, labels, loss_mask, attention_mask, position_ids, output_ref = batch.values() output = mamba_model[0].forward( input_ids=tokens,