diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 46fa8bfb865..9d0f88a144c 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -21,6 +21,7 @@ Changelog *Megatron Framework (M-LM / M-Bridge)* - Add ``--sft_in_batch_packing``, ``--moe_aux_loss_coeff`` and ``--dsa_kernel_backend`` to ``examples/megatron_bridge/distill.py`` (``--dsa_kernel_backend`` also to ``quantize.py``): pack ``--sft_hf_dataset`` conversations per micro-batch instead of padding each to ``--seq_length``, override the MoE load-balancing loss weight (e.g. 0 for padded SFT data), and select fused DeepSeek sparse-attention kernels for long sequences. +- Add ``--document_boundaries`` to ``examples/megatron_bridge/distill.py`` to attend within each document of a pre-tokenized window, and ``tokenizer`` / ``max_length`` / ``pack_max_tokens`` to ``chat_jsonl`` blends to drop conversations with nothing to train and pair the rest for ``--sft_in_batch_packing``. Packed distillation of Gated DeltaNet models needs ``flash-linear-attention`` >= 0.5.2 (earlier versions compute it wrongly); install ``examples/megatron_bridge/requirements.txt`` in ``nemo:26.08``. - Add ``--num_layers_in_first_pipeline_stage``, ``--num_layers_in_last_pipeline_stage`` and ``--recompute_num_layers_per_stage`` to ``examples/megatron_bridge/distill.py`` to balance pipeline-parallel distillation: give the last stage, which also runs the output layers and KD loss, fewer layers, and recompute fewer layers on later stages. - Add ``--vllm_fakequant_export`` to ``examples/megatron_bridge/export_quantized_megatron_to_hf.py`` to evaluate formats vLLM has no kernels for (e.g. GGML IQ) as BF16 fake-quant checkpoints, and allow ``--ep_size`` > 1 for checkpoints built with ``--no_moe_grouped_gemm``. - GGML GPTQ checkpoints from ``examples/megatron_bridge`` keep their calibrated codes when the Megatron checkpoint is reloaded for QAD or export, and layerwise calibration now runs on Megatron vision-language models. diff --git a/examples/dataset/MEGATRON_DATA_PREP.md b/examples/dataset/MEGATRON_DATA_PREP.md index c920ed02567..6dbed0b9e78 100644 --- a/examples/dataset/MEGATRON_DATA_PREP.md +++ b/examples/dataset/MEGATRON_DATA_PREP.md @@ -221,6 +221,24 @@ rows, or runs out of them stops the run. Each source's selection is cached under the settings that choose its rows, so an interrupted run resumes and a changed configuration selects afresh. Nothing prunes `parts`; delete it once `train.jsonl` is written to reclaim the space. +To feed `distill.py --sft_in_batch_packing`, also count tokens the way Megatron-Bridge trains them: + +```yaml +tokenizer: Qwen/Qwen3.6-35B-A3B # The student's tokenizer. +max_length: 32768 # distill.py --seq_length; longer conversations are truncated. +loss_mode: assistant # Optional; as distill.py --sft_loss_mode. +pack_max_tokens: 33280 # Optional; order each split as pairs of at most this many tokens. +num_workers: 48 # Optional; tokenizer processes (default: all CPUs). +``` + +With `tokenizer` and `max_length`, a conversation with no trained token left after truncation is +dropped, since Megatron-Bridge fails on it during training. `pack_max_tokens` pairs long +conversations with short ones and writes each split pair by pair, so `--mbs 2` packs every pair +into one micro-batch; a conversation that fits with no other is dropped. With `--cp_size` > 1, +Megatron-Bridge pads each conversation to a multiple of `2 * cp_size`, so a pair can exceed the cap +by under `2 * cp_size` tokens per conversation. Tokenizing runs on the CPU: ~0.4 s per +conversation per process for GLM-5.3-Flash on Nemotron-Cascade-2 (4.4K tokens on average). + ## `reasoning_content` for Post-Training v3 Datasets v3 datasets include a `reasoning_content` field in assistant messages (chain-of-thought separate from diff --git a/examples/megatron_bridge/README.md b/examples/megatron_bridge/README.md index a1ccb5f9f36..66683526675 100644 --- a/examples/megatron_bridge/README.md +++ b/examples/megatron_bridge/README.md @@ -184,6 +184,25 @@ the last one (`last_turn`), or the whole conversation (`full`). With `--eval_ite `--sft_hf_validation` in the same format. This path needs a Megatron-Bridge with `DirectHFSFTDatasetConfig`. +Megatron-Bridge pads every conversation to `--seq_length` when `--ep_size` > 1, so for MoE models +add `--sft_in_batch_packing --mbs 2 --moe_aux_loss_coeff 0`: each micro-batch packs two +conversations into one unpadded THD sequence, and the MoE load-balancing loss, which would also +run over padding, is dropped. Order the rows as pairs that fit together with the chat-blend +`pack_max_tokens` option ([Dataset Preparation README](../dataset/MEGATRON_DATA_PREP.md#chat-jsonl-blends-for-assistant-only-loss)). +Megatron-Bridge holds every row in host memory on each rank, so keep `train.jsonl` to the rows +the run needs. + +With pre-tokenized data, `--document_boundaries` (`--mbs 1`, `--cp_size 1`) restricts attention to +each EOD-separated document of a window instead of the whole window. Tokenize the data with +`--append_eod`: the student tokenizer's EOS marks each document's end. Chat-template data has no +EOD between conversations, so distill.py rejects it. + +> [!IMPORTANT] +> Packed training (`--sft_in_batch_packing` or `--document_boundaries`) of Gated DeltaNet models +> (e.g. Qwen3.5 / Qwen3.6) needs `flash-linear-attention` >= 0.5.2: earlier versions compute the +> packed linear attention wrongly. `nemo:26.08` ships 0.4.2, so run +> `python -m pip install -r requirements.txt` in this directory first. + Teacher and student must share a tokenizer — distillation scores the teacher on the student's token ids, and the KD losses compare the two models' logits elementwise over the vocab dimension. diff --git a/examples/megatron_bridge/distill.py b/examples/megatron_bridge/distill.py index 1ac3e1ca7eb..9bd4e1e9aaa 100644 --- a/examples/megatron_bridge/distill.py +++ b/examples/megatron_bridge/distill.py @@ -58,6 +58,8 @@ except ImportError: HAS_DIRECT_HF_SFT = False from megatron.core import parallel_state +from megatron.core.datasets.gpt_dataset import GPTDataset +from megatron.core.datasets.indexed_dataset import IndexedDataset from megatron.core.datasets.utils import get_blend_from_list from megatron.core.distributed import DistributedDataParallelConfig from megatron.core.utils import unwrap_model @@ -214,6 +216,12 @@ def get_args(): nargs="+", help="List of tokenized data paths to load from (weight1 path1 weight2 path2 ...)", ) + parser.add_argument( + "--document_boundaries", + action="store_true", + help="Pre-tokenized data only: attend within each EOD-separated document of a sequence (THD " + "segments) instead of across the whole packed window. Needs --mbs 1 and --cp_size 1.", + ) parser.add_argument( "--data_path_to_cache", type=str, default=None, help="Path to cache the dataset indices" ) @@ -492,6 +500,9 @@ def get_args(): raise ValueError( "--sft_hf_validation / --sft_loss_mode / --sft_in_batch_packing require --sft_hf_dataset." ) + if args.document_boundaries and (args.sft or args.mbs != 1 or args.cp_size != 1): + # Each window carries its own number of documents, so windows cannot be stacked. + raise ValueError("--document_boundaries needs pre-tokenized data, --mbs 1 and --cp_size 1.") if args.sft and (args.data_paths or args.use_mock_data): raise ValueError( "--sft is mutually exclusive with --data_paths / --use_mock_data: the SFT branch wins " @@ -540,6 +551,47 @@ def _tokenizer_prepends_bos(args) -> bool: return tokenizer("x").input_ids[:1] == [tokenizer.bos_token_id] +def _document_eod_id(args) -> int: + """The token ending each document: the student tokenizer's EOS, which ``--append_eod`` appends. + + Fails if a data file's first document does not end with it, e.g. chat-template data, which has + no EOD between conversations, rather than silently attending across documents. + """ + eod = AutoTokenizer.from_pretrained( + args.student_hf_path, trust_remote_code=args.trust_remote_code + ).eos_token_id + if eod is None: + raise ValueError("--document_boundaries needs a student tokenizer with an EOS token.") + for prefix in get_blend_from_list(args.data_paths)[0] if args.data_paths else []: + # Bound to a name: a document is a view of the dataset's memory map, closed with it. + dataset = IndexedDataset(prefix) + if dataset[0][-1] != eod: + raise ValueError( + f"--document_boundaries: the first document of {prefix} does not end with the " + f"student tokenizer's EOS ({eod}). Tokenize text data with --append_eod." + ) + return eod + + +def _add_document_boundaries() -> None: + """Give each EOD-separated document of a pre-tokenized window its own THD segment.""" + getitem = GPTDataset.__getitem__ + + def getitem_with_boundaries(self, idx): + sample = getitem(self, idx) + tokens = sample["tokens"] + n = tokens.numel() + ends = (tokens == self.config.tokenizer.eod).nonzero(as_tuple=True)[0] + 1 + # On the tokens' (CPU) device: CUDA tensors from a dataloader worker are shared over IPC. + bounds = [tokens.new_zeros(1), ends[ends < n], tokens.new_tensor([n])] + cu_seqlens = torch.cat(bounds).to(torch.int32) + sample["cu_seqlens_q"] = cu_seqlens + sample["max_seqlen_q"] = (cu_seqlens[1:] - cu_seqlens[:-1]).max() + return sample + + GPTDataset.__getitem__ = getitem_with_boundaries + + def main(args: argparse.Namespace, owns_the_run: bool = True): # --student_hf_path stays as given. The student is built and exported from the local copy of # the whole checkpoint in student_local_checkpoint_path, since the HF export reads the source's @@ -547,6 +599,8 @@ def main(args: argparse.Namespace, owns_the_run: bool = True): args.student_hub_model_id, args.student_local_checkpoint_path = ensure_local_checkpoint( args.student_hf_path ) + # NullTokenizer's EOD defaults to vocab_size - 1, which pre-tokenized data never contains. + eod_id = _document_eod_id(args) if args.document_boundaries else None student_has_modelopt_state = args.student_megatron_path is not None and has_modelopt_state( args.student_megatron_path ) @@ -683,6 +737,8 @@ def _restore_student_hook(model_chunks): return model_chunks distill_provider.register_pre_wrap_hook(_restore_student_hook, prepend=True) + if args.document_boundaries: + _add_document_boundaries() if args.recompute_num_layers_per_stage is not None: # Set this pipeline stage's student recompute depth before the model is wrapped. @@ -714,7 +770,7 @@ def _per_stage_recompute_hook(model_chunks): "path_to_cache": args.data_path_to_cache, "random_seed": args.seed, "reset_attention_mask": False, - "reset_position_ids": False, + "reset_position_ids": args.document_boundaries, "eod_mask_loss": False, "num_dataset_builder_threads": 1, "data_sharding": True, @@ -829,7 +885,9 @@ def _per_stage_recompute_hook(model_chunks): ) if args.sft else TokenizerConfig( - tokenizer_type="NullTokenizer", vocab_size=distill_provider.vocab_size + tokenizer_type="NullTokenizer", + vocab_size=distill_provider.vocab_size, + null_tokenizer_eod_id=eod_id, ) ), checkpoint=CheckpointConfig( diff --git a/examples/megatron_bridge/requirements.txt b/examples/megatron_bridge/requirements.txt new file mode 100644 index 00000000000..e8e6ce84c91 --- /dev/null +++ b/examples/megatron_bridge/requirements.txt @@ -0,0 +1,4 @@ +# nemo:26.08 ships flash-linear-attention 0.4.2, whose Gated DeltaNet kernel computes packed (THD) sequences +# wrongly (--sft_in_batch_packing, --document_boundaries); 0.5.2 is the first correct release. +# TODO: remove once the examples move to nemo:26.10, which ships 0.5.2. +flash-linear-attention>=0.5.2 diff --git a/modelopt/torch/utils/plugins/prepare_megatron_data_blend.py b/modelopt/torch/utils/plugins/prepare_megatron_data_blend.py index 2bc7aacb8c9..b2f6bae779e 100644 --- a/modelopt/torch/utils/plugins/prepare_megatron_data_blend.py +++ b/modelopt/torch/utils/plugins/prepare_megatron_data_blend.py @@ -20,6 +20,7 @@ import hashlib import json import math +import multiprocessing as mp import os import random import re @@ -178,7 +179,10 @@ def _prepare_sources( def _to_chat( - row: dict[str, Any], content_field: str, reasoning_content: ReasoningContentMode = "native" + row: dict[str, Any], + content_field: str, + reasoning_content: ReasoningContentMode = "native", + loss_mode: str = "assistant", ) -> dict[str, Any] | None: """Convert a source row to ``{"messages": [...]}``, or ``None`` if it has nothing to train on. @@ -224,8 +228,10 @@ def _to_chat( if content.strip() ] chat = {"messages": turns} - # Nothing in a row without an assistant turn is trained under assistant-only loss. - if not any(turn.get("role") == "assistant" for turn in chat["messages"]): + # Only loss_mode full trains a row without an assistant turn. + if loss_mode != "full" and not any( + turn.get("role") == "assistant" for turn in chat["messages"] + ): return None return chat @@ -260,6 +266,7 @@ def _take_chat_rows( shuffle: bool, seed: int, reasoning_content: ReasoningContentMode, + loss_mode: str, ) -> list[dict[str, Any]]: """Return the first ``want`` usable chat rows of a source. @@ -277,7 +284,7 @@ def _take_chat_rows( max_rejected = 10 * want + 100 try: for row in _iter_rows(source, output_dir, shuffle, seed): - chat = _to_chat(row, source["content_field"], reasoning_content) + chat = _to_chat(row, source["content_field"], reasoning_content, loss_mode) # Judge the source row's top-level media keys too, which the conversion does not keep. if chat is None or not is_text_only_chat_example({**row, **chat}): rejected += 1 @@ -326,6 +333,69 @@ def _part_path(parts_dir: Path, source: dict[str, Any], settings: dict[str, Any] return parts_dir / f"{source['hf_dataset'].replace('/', '--')}_{label}_{digest}.jsonl" +_SFT_TOKENIZER: dict[str, Any] = {} + + +def _init_sft_tokenizer(tokenizer: str, max_length: int, loss_mode: str) -> None: + from megatron.bridge.data.collators.sft import extract_skipped_token_ids + from megatron.bridge.data.conversation_processing import infer_assistant_mask_boundary_config + from transformers import AutoTokenizer + + # distill.py --sft_hf_dataset tokenizes with the student's HF tokenizer, as loaded here. + processor = AutoTokenizer.from_pretrained(tokenizer) + _SFT_TOKENIZER.update( + processor=processor, + boundary_config=infer_assistant_mask_boundary_config(processor), + skipped_tokens=extract_skipped_token_ids(processor), + max_length=max_length, + loss_mode=loss_mode, + ) + + +def _sft_length(row: dict[str, Any]) -> int: + """Return the row's token count as Megatron-Bridge trains it, or 0 if nothing would be trained.""" + from megatron.bridge.data.conversation_processing import tokenize_chat_example + + try: + tokenized = tokenize_chat_example(row, warn_on_all_masked=False, **_SFT_TOKENIZER) + except ValueError as e: + # No assistant span matched: Bridge raises from a dataloader worker. Other errors are setup + # errors (e.g. a tokenizer without a chat template) that would drop every row. + if "did not match any loss-contributing spans" not in str(e): + raise + return 0 + return tokenized.input_ids.numel() if tokenized.assistant_mask.any() else 0 + + +def _sft_lengths(rows: list[dict[str, Any]], config: dict[str, Any]) -> list[int]: + """Tokenize ``rows`` with Megatron-Bridge's chat SFT path in ``num_workers`` processes.""" + init_args = (config["tokenizer"], config["max_length"], config.get("loss_mode", "assistant")) + num_workers = int(config.get("num_workers", len(os.sched_getaffinity(0)))) + if num_workers <= 1: + _init_sft_tokenizer(*init_args) + return [_sft_length(row) for row in rows] + with mp.Pool(num_workers, initializer=_init_sft_tokenizer, initargs=init_args) as pool: + return pool.map(_sft_length, rows, chunksize=64) + + +def _pair_rows(rows: list[dict[str, Any]], lengths: list[int], cap: int, seed: int) -> list: + """Order rows as pairs of at most ``cap`` tokens, so ``--mbs 2`` packs each pair together. + + Pairs the longest remaining row with the shortest; a row that fits with no other is dropped. + The pairs are shuffled with ``seed``. + """ + order = sorted(range(len(rows)), key=lengths.__getitem__) + pairs, i, j = [], 0, len(order) - 1 + while i < j: + if lengths[order[i]] + lengths[order[j]] <= cap: + pairs.append((order[i], order[j])) + i += 1 + j -= 1 + random.Random(seed).shuffle(pairs) + print(f"Paired {2 * len(pairs)} of {len(rows)} rows into packs of <= {cap} tokens", flush=True) + return [rows[k] for pair in pairs for k in pair] + + def _chat_budgets(config: dict[str, Any]) -> list[int]: """Return each source's row budget; validate the inputs.""" weights = [float(source["weight"]) for source in config["sources"]] @@ -354,6 +424,12 @@ def prepare_chat_data_blend(config_path: Path) -> dict[str, int]: The rows are mixed with ``seed`` and ``min(1000, n // 20)`` of them are held out for validation. Each source is cached under ``/parts``, keyed by the settings that select its rows, so an interrupted run resumes. + + With ``tokenizer`` and ``max_length`` (``distill.py --seq_length``), rows are tokenized as + Megatron-Bridge trains them (``loss_mode``, default ``assistant``, as ``--sft_loss_mode``), in + ``num_workers`` processes, and rows with no trained token are dropped. ``pack_max_tokens`` + then orders each split as pairs of at most that many tokens for ``--sft_in_batch_packing + --mbs 2``. """ config = load_config(config_path) if config.get("output_format") != "chat_jsonl": @@ -361,9 +437,20 @@ def prepare_chat_data_blend(config_path: Path) -> dict[str, int]: "prepare_chat_data_blend writes output_format: chat_jsonl, " f"got {config.get('output_format')!r}." ) - ignored = {"tokenizer", "target_tokens"} & config.keys() - if ignored: - raise ValueError(f"output_format: chat_jsonl does not use {sorted(ignored)}.") + if "target_tokens" in config: + raise ValueError("output_format: chat_jsonl does not use ['target_tokens'].") + if ("tokenizer" in config) != ("max_length" in config): + raise ValueError("chat_jsonl tokenizes rows only with both tokenizer and max_length.") + if "pack_max_tokens" in config and "max_length" not in config: + raise ValueError("pack_max_tokens needs tokenizer and max_length to count tokens.") + for key in ("max_length", "pack_max_tokens", "num_workers"): + value = config.get(key, 1) + if isinstance(value, bool) or not isinstance(value, int) or value <= 0: + raise ValueError(f"{key} must be a positive integer, got {value!r}.") + if config.get("loss_mode", "assistant") not in ("assistant", "last_turn", "full"): + raise ValueError( + f"loss_mode must be assistant, last_turn or full: {config['loss_mode']!r}." + ) if any("max_samples" in source for source in config["sources"]): raise ValueError("output_format: chat_jsonl does not use max_samples; set target_samples.") fields = {source["content_field"] for source in config["sources"]} - {"messages", "text"} @@ -387,6 +474,9 @@ def prepare_chat_data_blend(config_path: Path) -> dict[str, int]: settings = {"shuffle": shuffle, "reasoning_content": reasoning_content} if shuffle: settings["seed"] = seed + loss_mode = config.get("loss_mode", "assistant") + if loss_mode == "full": # keeps rows without an assistant turn + settings["loss_mode"] = loss_mode rows: list[dict[str, Any]] = [] for source, want in zip(config["sources"], budgets): @@ -396,13 +486,24 @@ def prepare_chat_data_blend(config_path: Path) -> dict[str, int]: part = _part_path(parts_dir, source, settings) got = _read_jsonl(part) if got is None or len(got) < want: - got = _take_chat_rows(source, output_dir, want, shuffle, seed, reasoning_content) + got = _take_chat_rows( + source, output_dir, want, shuffle, seed, reasoning_content, loss_mode + ) _write_jsonl(part, got) rows.extend(got[:want]) random.Random(seed).shuffle(rows) + if "max_length" in config: + lengths = _sft_lengths(rows, config) + kept = [k for k, n in enumerate(lengths) if n] + print(f"Dropped {len(rows) - len(kept)} rows with no trained token", flush=True) + rows, lengths = [rows[k] for k in kept], [lengths[k] for k in kept] num_validation = min(1000, len(rows) // 20) splits = {"validation": rows[:num_validation], "train": rows[num_validation:]} + if "pack_max_tokens" in config: + cap = config["pack_max_tokens"] + splits["validation"] = _pair_rows(splits["validation"], lengths[:num_validation], cap, seed) + splits["train"] = _pair_rows(splits["train"], lengths[num_validation:], cap, seed) if not splits["train"] or not splits["validation"]: raise RuntimeError(f"Too few rows ({len(rows)}) for a train/validation split.") for name, split_rows in splits.items(): @@ -420,7 +521,16 @@ def prepare_megatron_data_blend(config_path: Path) -> list[tuple[float, str]]: f"Unknown output_format {output_format!r} for a tokenized blend; use megatron, or " "chat_jsonl with prepare_chat_data_blend." ) - chat_keys = {"target_samples", "shuffle", "seed", "reasoning_content"} & config.keys() + chat_keys = { + "target_samples", + "shuffle", + "seed", + "reasoning_content", + "max_length", + "loss_mode", + "pack_max_tokens", + "num_workers", + } & config.keys() if chat_keys: raise ValueError(f"Tokenized blends do not use {sorted(chat_keys)}.") output_dir = Path(config["output_dir"]) diff --git a/tests/examples/megatron_bridge/test_distill.py b/tests/examples/megatron_bridge/test_distill.py index 1fb574d686e..eae2e14df84 100644 --- a/tests/examples/megatron_bridge/test_distill.py +++ b/tests/examples/megatron_bridge/test_distill.py @@ -18,6 +18,7 @@ import json import sys from pathlib import Path +from types import SimpleNamespace import pytest import torch @@ -29,13 +30,15 @@ create_tiny_qwen3_moe_dir, get_tiny_tokenizer, ) +from megatron.core.datasets.gpt_dataset import GPTDataset from transformers import AutoModelForImageTextToText from modelopt.torch.puzzletron.anymodel import convert_model +from modelopt.torch.utils.plugins.megatron_preprocess_data import megatron_preprocess_data sys.path.insert(0, str(Path(__file__).resolve().parents[3] / "examples" / "megatron_bridge")) -from distill import HAS_DIRECT_HF_SFT, _hf_source +from distill import HAS_DIRECT_HF_SFT, _add_document_boundaries, _document_eod_id, _hf_source from export_distilled_megatron_to_hf import _get_checkpoint_export_paths @@ -43,11 +46,15 @@ def test_distill_llm_hf_export(tmp_path, num_gpus, fake_hub_checkpoint): """Distill with an inline HF export, then re-export every saved iteration standalone. - Both exports read the student from the copy ``ensure_local_checkpoint`` resolved (see - ``fake_hub_checkpoint``). distill.py keeps the local path on its command line, since get_args() - loads the tokenizer from it; the standalone export is passed the Hub ID. + Distillation runs with --document_boundaries, so each EOD-separated document of a window trains + as its own THD segment. Both exports read the student from the copy ``ensure_local_checkpoint`` + resolved (see ``fake_hub_checkpoint``). distill.py keeps the local path on its command line, + since get_args() loads the tokenizer from it; the standalone export is passed the Hub ID. """ - teacher_hf_path = create_tiny_qwen3_dir(tmp_path, with_tokenizer=True) + # head_dim=8: TE's unfused attention, the only backend for head_dim 2, rejects THD input. + teacher_hf_path = create_tiny_qwen3_dir( + tmp_path, with_tokenizer=True, head_dim=8, max_position_embeddings=2048 + ) student_export_arg = fake_hub_checkpoint.serve( teacher_hf_path, "distill.py", "export_distilled_megatron_to_hf.py", nproc=num_gpus ) @@ -56,13 +63,21 @@ def test_distill_llm_hf_export(tmp_path, num_gpus, fake_hub_checkpoint): distilled_hf_path = tmp_path / "distilled_hf" all_exports = tmp_path / "all_exports" distill_cmd_parts = extend_cmd_parts( - ["torchrun", f"--nproc_per_node={num_gpus}", "distill.py", "--use_mock_data"], + [ + "torchrun", + f"--nproc_per_node={num_gpus}", + "distill.py", + "--use_mock_data", + "--document_boundaries", + ], student_hf_path=teacher_hf_path, teacher_hf_path=teacher_hf_path, output_dir=distill_output_dir, tp_size=num_gpus, pp_size=1, - seq_length=16, + # Mock token ids cycle through the 128-token vocab, so the EOD id recurs every 128 tokens + # and each window trains as ~17 THD segments. + seq_length=2048, mbs=1, gbs=4, train_iters=train_iters, @@ -97,6 +112,53 @@ def test_distill_llm_hf_export(tmp_path, num_gpus, fake_hub_checkpoint): assert (all_exports / f"iter_{it:07d}" / fake_hub_checkpoint.marker).exists() +@pytest.mark.parametrize("append_eod", [True, False]) +def test_document_eod_id(tmp_path, append_eod): + """Document boundaries use the EOS --append_eod wrote, and fail on data without it.""" + tokenizer = get_tiny_tokenizer() + tokenizer.save_pretrained(tmp_path / "tokenizer") + jsonl = tmp_path / "docs.jsonl" + jsonl.write_text("".join(json.dumps({"text": f"document number {i}"}) + "\n" for i in range(4))) + (prefix,) = megatron_preprocess_data( + jsonl_paths=jsonl, + output_dir=tmp_path / "data", + tokenizer_name_or_path=str(tmp_path / "tokenizer"), + append_eod=append_eod, + ) + args = argparse.Namespace( + student_hf_path=str(tmp_path / "tokenizer"), + trust_remote_code=False, + data_paths=["1.0", prefix], + ) + if append_eod: + assert _document_eod_id(args) == tokenizer.eos_token_id + else: + with pytest.raises(ValueError, match="--append_eod"): + _document_eod_id(args) + + +@pytest.mark.parametrize( + ("tokens", "cu_seqlens"), + [ + ([1, 2, 9, 3, 4, 5, 9, 6], [0, 3, 7, 8]), # EOD mid-window + ([1, 2, 3, 9, 4, 5, 6, 9], [0, 4, 8]), # EOD as the last token + ([9, 1, 2, 3, 4, 5, 6, 7], [0, 1, 8]), # EOD as the first token + ([1, 2, 3, 4, 5, 6, 7, 8], [0, 8]), # no EOD + ], +) +def test_document_boundaries_cu_seqlens(monkeypatch, tokens, cu_seqlens): + """Each EOD ends a THD segment; the window's end closes the last one.""" + monkeypatch.setattr( + GPTDataset, "__getitem__", lambda self, idx: {"tokens": torch.tensor(tokens)} + ) + _add_document_boundaries() # monkeypatch restores the original __getitem__ on teardown + dataset = SimpleNamespace(config=SimpleNamespace(tokenizer=SimpleNamespace(eod=9))) + sample = GPTDataset.__getitem__(dataset, 0) + expected = torch.tensor(cu_seqlens, dtype=torch.int32) + assert torch.equal(sample["cu_seqlens_q"], expected) + assert sample["max_seqlen_q"] == (expected[1:] - expected[:-1]).max() + + def test_distill_uneven_pipeline_per_stage_recompute(tmp_path, num_gpus): """An uneven student/teacher layer split with a different recompute depth on each stage.""" if num_gpus < 2: diff --git a/tests/gpu_megatron/torch/utils/plugins/test_prepare_megatron_data_blend.py b/tests/gpu_megatron/torch/utils/plugins/test_prepare_megatron_data_blend.py index 6306ecef292..a037f6f459d 100644 --- a/tests/gpu_megatron/torch/utils/plugins/test_prepare_megatron_data_blend.py +++ b/tests/gpu_megatron/torch/utils/plugins/test_prepare_megatron_data_blend.py @@ -176,8 +176,11 @@ def test_to_chat_handles_all_row_shapes(): msgs = [{"role": "user", "content": "hi"}, {"role": "assistant", "content": "hello"}] assert blend_module._to_chat({"messages": msgs}, "messages") == {"messages": msgs} assert blend_module._to_chat({"messages": []}, "messages") is None - # Nothing would be trained under assistant-only loss. + # Nothing would be trained under assistant-only loss; loss_mode full trains the user turn. assert blend_module._to_chat({"messages": msgs[:1]}, "messages") is None + assert blend_module._to_chat({"messages": msgs[:1]}, "messages", loss_mode="full") == { + "messages": msgs[:1] + } # Tool definitions and a top-level system prompt reach the chat path. tools = [{"type": "function", "function": {"name": "add"}}] row = {"messages": msgs, "tools": tools, "system": "Be brief.", "id": 7} @@ -228,13 +231,13 @@ def test_take_chat_rows_skips_non_text_rows_and_fails_without_enough(tmp_path: P # Media is rejected inside message content and as a top-level key the conversion drops. rows = [_chat(0), _image_chat(1), {**_chat(9), "images": ["a.png"]}, _chat(2), _chat(3)] monkeypatch.setattr(blend_module, "_iter_rows", lambda *a: iter(rows)) - got = blend_module._take_chat_rows(source, tmp_path, 3, False, 0, "native") + got = blend_module._take_chat_rows(source, tmp_path, 3, False, 0, "native", "assistant") assert [c["messages"][0]["content"] for c in got] == ["q0", "q2", "q3"] images = [_image_chat(i) for i in range(200)] monkeypatch.setattr(blend_module, "_iter_rows", lambda *a: iter(images)) with pytest.raises(RuntimeError, match=r"only 0/2 usable rows \(gave up after 121 rejected"): - blend_module._take_chat_rows(source, tmp_path, 2, False, 0, "native") + blend_module._take_chat_rows(source, tmp_path, 2, False, 0, "native", "assistant") def test_chat_blend_splits_weighted_rows(tmp_path: Path, chat_blend, monkeypatch): @@ -281,8 +284,9 @@ def test_chat_blend_reuses_and_rebuilds_cached_parts(tmp_path: Path, chat_blend) ({}, {"shuffle": True}), ({"shuffle": True, "seed": 0}, {"shuffle": True, "seed": 1}), ({}, {"reasoning_content": "inline"}), + ({}, {"loss_mode": "full"}), ], - ids=["split", "files", "shuffle", "seed", "reasoning_content"], + ids=["split", "files", "shuffle", "seed", "reasoning_content", "loss_mode"], ) def test_chat_blend_reselects_after_a_selection_change(chat_blend, first, second): write_config, calls = chat_blend @@ -302,7 +306,11 @@ def test_chat_blend_reselects_after_a_selection_change(chat_blend, first, second ({"sources": [_source("a", weight=-1)]}, "nonnegative"), ({"sources": [_source("a", weight=0)]}, "not all zero"), ({"sources": [_source("a", weight=float("inf"))]}, "finite"), - ({"tokenizer": "x"}, "does not use"), + ({"tokenizer": "x"}, "both tokenizer and max_length"), + ({"max_length": 64}, "both tokenizer and max_length"), + ({"pack_max_tokens": 64}, "pack_max_tokens needs tokenizer and max_length"), + ({"tokenizer": "x", "max_length": 0}, "max_length must be a positive integer"), + ({"tokenizer": "x", "max_length": 64, "loss_mode": "last"}, "loss_mode must be"), ({"target_tokens": 1000}, "does not use"), ({"reasoning_content": "drop"}, "reasoning_content must be one of"), # A row with both columns must not silently be read from ``text``. @@ -379,3 +387,47 @@ def test_chat_blend_reads_files_sources(tmp_path: Path, monkeypatch): assert sorted(r["messages"][0]["content"] for r in rows) == sorted(f"q{i}" for i in range(40)) # The part is named after its first file. assert len(list((out / "parts").glob("org--a_a_*.jsonl"))) == 1 + + +def test_chat_blend_drops_untrainable_rows_and_pairs_the_rest( + tiny_tokenizer_path: str, tmp_path: Path, chat_blend, monkeypatch +): + def iter_rows(source, output_dir, shuffle, seed): + # Prompts of 0..99 words; past max_length the truncation leaves no assistant token. + yield from ( + { + "messages": [ + {"role": "user", "content": "w " * i}, + {"role": "assistant", "content": "a"}, + ] + } + for i in range(100) + ) + + monkeypatch.setattr(blend_module, "_iter_rows", iter_rows) + write_config, _ = chat_blend + config = {"tokenizer": tiny_tokenizer_path, "max_length": 64, "num_workers": 1} + sizes = blend_module.prepare_chat_data_blend(write_config(**config)) + out = tmp_path / "blend" + blend_module._init_sft_tokenizer(tiny_tokenizer_path, 64, "assistant") + rows = _read_jsonl(out / "train.jsonl") + _read_jsonl(out / "validation.jsonl") + assert 40 < len(rows) < 100 # the 50 + 50 selected rows, minus the long prompts + assert all(blend_module._sft_length(row) for row in rows) + # A setup error would drop every row, so it is raised instead. + monkeypatch.setitem(blend_module._SFT_TOKENIZER, "loss_mode", "unknown") + with pytest.raises(ValueError, match="loss_mode"): + blend_module._sft_length(rows[0]) + + # A cap of 2 * max_length fits any two rows, so only an odd row per split is left out. + paired = blend_module.prepare_chat_data_blend(write_config(**config, pack_max_tokens=128)) + assert paired == {split: n - n % 2 for split, n in sizes.items()} + + +def test_pair_rows_pairs_long_with_short_under_the_cap(): + lengths = {"a": 10, "b": 60, "c": 30, "d": 50, "e": 95, "f": 20} + rows = blend_module._pair_rows(list(lengths), list(lengths.values()), cap=80, seed=0) + # e fits with no row and c is left without a partner; rows (2k, 2k + 1) are a pack. + assert {frozenset(rows[k : k + 2]) for k in range(0, len(rows), 2)} == { + frozenset("ab"), + frozenset("df"), + }