From eef01dcce11788f6cc2d4fc4d24ff3c27d3dec14 Mon Sep 17 00:00:00 2001 From: Yue Date: Wed, 14 Oct 2026 13:27:47 +0000 Subject: [PATCH 1/6] Add a chat JSONL output mode to prepare_megatron_data_blend Tokenized .bin/.idx blends carry no loss mask, so --data_paths training computes the loss on every token. output_format: chat_jsonl writes the same weighted sources as train.jsonl / validation.jsonl chat data for distill.py --sft --sft_hf_dataset, which can restrict the loss to assistant tokens (--sft_loss_mode assistant). target_samples, shuffle and seed control the row budget, in-order selection and the mix. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Yue --- CHANGELOG.rst | 1 + examples/dataset/MEGATRON_DATA_PREP.md | 35 +++ .../plugins/prepare_megatron_data_blend.py | 192 ++++++++++++++- .../test_prepare_megatron_data_blend.py | 220 ++++++++++++++++++ 4 files changed, 436 insertions(+), 12 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 10e57cb2f85..d495140506e 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -64,6 +64,7 @@ Changelog - Add chat-data distillation to ``examples/megatron_bridge/distill.py``: ``--sft --sft_hf_dataset`` takes a local chat ``.jsonl`` or a Hugging Face Hub dataset, renders each conversation with the model's chat template, and ``--sft_loss_mode`` selects which turns are trained. Requires a Megatron-Bridge with ``DirectHFSFTDatasetConfig``. - Add ``freeze_mtp`` to the Megatron ``DistillationConfig`` (``--freeze_mtp`` in ``examples/megatron_bridge/distill.py``): the student's MTP head keeps its input weights and its loss is skipped, so distillation trains the rest of the model on the distillation loss alone. +- Add ``output_format: chat_jsonl`` to ``prepare_megatron_data_blend``: it writes a weighted blend as ``train.jsonl`` / ``validation.jsonl`` chat data for ``distill.py --sft --sft_hf_dataset``, so the distillation loss can be restricted to assistant tokens. See the `Megatron data preparation guide `_. - Add mergeable LoRA QAD for fake-quantized dense Megatron/Transformer Engine linear layers. Enable it with ``--lora_rank`` in Megatron-Bridge PTQ, train through the normal QAD flow, and export a standalone quantized checkpoint without deployment adapters. - Add optional Top-P (nucleus) truncation to the Megatron top-k logits KD loss (now ``TopLogitsKLLoss``) via ``logit_kl_top_p`` and ``logit_kl_top_p_min_k`` in ``DistillationConfig``: after the global Top-K selection, only the smallest prefix whose cumulative teacher probability reaches ``top_p`` (with a floor of ``min_k`` entries) contributes to the KL, mirroring ``--logits-save-top-p`` / ``--logits-save-top-p-min-k`` in Megatron-LM's logits saver. - Add an end-to-end W4A4 NVFP4 PTQ and QAD tutorial for Qwen3.6-35B-A3B also covering evaluation and vLLM throughput benchmarking. See `examples/megatron_bridge/tutorials/Qwen3.6-35B-A3B/README.md `_ for details. diff --git a/examples/dataset/MEGATRON_DATA_PREP.md b/examples/dataset/MEGATRON_DATA_PREP.md index 91b89375907..4bfad9c3262 100644 --- a/examples/dataset/MEGATRON_DATA_PREP.md +++ b/examples/dataset/MEGATRON_DATA_PREP.md @@ -5,6 +5,7 @@ | From JSONL files | Tokenize local JSONL files | \[[Link](#from-jsonl-files)\] | | From Hugging Face Hub | Stream or download HF datasets and tokenize | \[[Link](#from-hugging-face-hub)\] | | Token-budgeted data blends | Prepare weighted subsets for fast experiments | \[[Link](#prepare-token-budgeted-data-blends)\] | +| Chat JSONL blends | Weighted chat data for assistant-only loss | \[[Link](#chat-jsonl-blends-for-assistant-only-loss)\] | | `reasoning_content` for Post-Training v3 | Control how chain-of-thought traces are handled | \[[Link](#reasoning_content-for-post-training-v3-datasets)\] | | Nemotron Pre/Post-Training Datasets | Ready-to-run commands for all Nemotron datasets | \[[Link](#ready-to-run-tokenization-commands)\] | @@ -172,6 +173,40 @@ The output contains tokenized Megatron `.bin`/`.idx` files, `data_blend.txt` wit and `config.yaml` recording how the blend was generated. The final token count can slightly exceed the target because the final document from each source is kept whole. +### Chat JSONL blends for assistant-only loss + +Tokenized `.bin`/`.idx` data carries no loss mask, so `--data_paths` training computes the loss on every token. +To compute it on assistant tokens only, set `output_format: chat_jsonl`. The same sources are then written, +untokenized, as `train.jsonl` and `validation.jsonl` with one `{"messages": [...]}` conversation per line, for +Megatron-Bridge's chat path (`examples/megatron_bridge/distill.py --sft --sft_hf_dataset /train.jsonl +--sft_hf_validation /validation.jsonl --sft_loss_mode assistant`), which applies the model's chat template +when it loads the data. Run it in an environment with Megatron-Bridge, whose text-only chat filter selects the rows. + +```yaml +output_dir: /path/to/chat_blend +output_format: chat_jsonl +target_samples: 36_800 # Optional; omit to keep every usable row. +shuffle: false # Optional; take the first rows of each split in order (default: buffer shuffle). +seed: 1234 # Optional; seeds the cross-source mix and the validation hold-out. +sources: + - hf_dataset: nvidia/Nemotron-Pretraining-SFT-v1 + config: Nemotron-SFT-General + split: train + content_field: text + weight: 20 + - hf_dataset: nvidia/Nemotron-SFT-Math-v3 + config: default + split: train + content_field: messages + weight: 17 +``` + +Each source contributes `round(target_samples * weight / sum(weights))` conversations that contain only text. +A `text` row becomes one assistant turn, or is split into turns if it uses the `User` / +`Assistant` chat markers. The conversations are mixed and `min(1000, n // 20)` of them are held out +for validation. A source that fails to load or runs out of usable rows stops the run. Each source is cached under +`/parts`, so an interrupted run resumes. + ## `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/modelopt/torch/utils/plugins/prepare_megatron_data_blend.py b/modelopt/torch/utils/plugins/prepare_megatron_data_blend.py index fa5875a5e2c..4409d4f832b 100644 --- a/modelopt/torch/utils/plugins/prepare_megatron_data_blend.py +++ b/modelopt/torch/utils/plugins/prepare_megatron_data_blend.py @@ -16,17 +16,29 @@ """Prepare a weighted Megatron data blend from a YAML configuration.""" import argparse +import json import os +import random +import re import shutil +from collections.abc import Iterator from pathlib import Path from typing import Any, cast import huggingface_hub import yaml +from datasets import load_dataset from .megatron_preprocess_data import megatron_preprocess_data -__all__ = ["prepare_megatron_data_blend"] +__all__ = ["prepare_chat_data_blend", "prepare_megatron_data_blend"] + +# A turn of a ``text`` row in the NeMo chat format: ``User\n...``. Any marker ends the +# turn, so a trailing delimiter or a label marker is not kept as content. +_CHAT_TURN = re.compile( + r"(User|Assistant|System)\n?(.*?)(?=|\Z)", re.DOTALL +) +_CHAT_ROLES = {"User": "user", "Assistant": "assistant", "System": "system"} def load_config(path: Path) -> dict[str, Any]: @@ -53,6 +65,10 @@ def load_config(path: Path) -> dict[str, Any]: optional ``target_tokens``. Each source has ``hf_dataset``, ``content_field``, and ``weight``; it uses ``split`` with optional ``config`` and ``max_samples``, or selects repository ``files``. + + With ``output_format: chat_jsonl``, the blend is written as chat JSONL instead (see + :func:`prepare_chat_data_blend`): ``tokenizer`` and ``target_tokens`` are not used, and the + optional ``target_samples``, ``shuffle`` and ``seed`` keys apply. """ with path.open(encoding="utf-8") as stream: return cast("dict[str, Any]", yaml.safe_load(stream)) @@ -76,6 +92,18 @@ def _copy_config(source: Path, destination: Path) -> None: shutil.copyfile(source, destination) +def _download_files(source: dict[str, Any], output_dir: Path) -> list[str]: + """Download the repository ``files`` of a source and return their local paths.""" + dataset = source["hf_dataset"] + raw_dir = output_dir.parent / "raw" / dataset.replace("/", "--") + return [ + huggingface_hub.hf_hub_download( + repo_id=dataset, filename=file, repo_type="dataset", local_dir=raw_dir + ) + for file in source["files"] + ] + + def _prepare_sources( sources: list[dict[str, Any]], output_dir: Path, @@ -104,17 +132,7 @@ def _prepare_sources( content_field = source["content_field"] input_args: dict[str, Any] if "files" in source: - raw_dir = output_dir.parent / "raw" / dataset.replace("/", "--") - paths = [ - huggingface_hub.hf_hub_download( - repo_id=dataset, - filename=file, - repo_type="dataset", - local_dir=raw_dir, - ) - for file in source["files"] - ] - input_args = {"jsonl_paths": paths} + input_args = {"jsonl_paths": _download_files(source, output_dir)} else: input_args = { "hf_dataset": dataset, @@ -146,9 +164,152 @@ def _prepare_sources( return blend +def _to_chat(row: dict[str, Any], content_field: str) -> dict[str, Any] | None: + """Convert a source row to ``{"messages": [...]}``, or ``None`` if it has no usable content. + + A ``text`` row in the NeMo chat format (``User`` / ``Assistant`` markers) is split + into turns; any other ``text`` row becomes one assistant turn, so all of it is trained. + """ + if content_field == "messages": + return {"messages": row["messages"]} if row.get("messages") else None + text = row.get("text") or "" + if not text: + return None + if " delimiters as assistant content. + return {"messages": turns} if turns else None + return {"messages": [{"role": "assistant", "content": text}]} + + +def _iter_rows(source: dict[str, Any], output_dir: Path, shuffle: bool) -> Iterator[dict[str, Any]]: + """Yield the rows of a source: its ``files`` in order, or its Hugging Face split, streamed.""" + if "files" in source: + for path in _download_files(source, output_dir): + with open(path, encoding="utf-8") as f: + yield from (json.loads(line) for line in f if line.strip()) + return + dataset = load_dataset( + source["hf_dataset"], source.get("config"), split=source["split"], streaming=True + ) + yield from dataset.shuffle(seed=42) if shuffle else dataset + + +def _take_chat_rows( + source: dict[str, Any], output_dir: Path, want: int | None, shuffle: bool +) -> list[dict[str, Any]]: + """Return the first ``want`` usable chat rows of a source (all of them if ``want`` is None). + + Fails if the source cannot be read or runs out of usable rows, so the blend cannot change + silently; gives up after ``10 * want`` rejected rows. + """ + # Megatron-Bridge's chat SFT path consumes this output; keep exactly the rows it accepts. + from megatron.bridge.data.conversation_processing import is_text_only_chat_example + + name = f"{source['hf_dataset']} {source.get('split') or source.get('files')}" + rows: list[dict[str, Any]] = [] + rejected = 0 + try: + for row in _iter_rows(source, output_dir, shuffle): + chat = _to_chat(row, source["content_field"]) + if chat is None or not is_text_only_chat_example(chat): + rejected += 1 + if want is not None and rejected > 10 * want: + break + continue + rows.append(chat) + if want is not None and len(rows) >= want: + break + except Exception as e: # e.g. a gated split, or a schema the streaming reader cannot cast + raise RuntimeError(f"{name}: failed to load") from e + if want is not None and len(rows) < want: + raise RuntimeError(f"{name}: only {len(rows)}/{want} usable rows") + print(f"{name}: {len(rows)} rows ({rejected} rejected)", flush=True) + return rows + + +def _write_jsonl(path: Path, rows: list[dict[str, Any]]) -> None: + """Write ``rows`` atomically, so an interrupted write never leaves a partial file.""" + tmp = path.with_name(path.name + ".tmp") + with tmp.open("w", encoding="utf-8") as f: + f.writelines(json.dumps(row) + "\n" for row in rows) + os.replace(tmp, path) + + +def _read_jsonl(path: Path) -> list[dict[str, Any]] | None: + """Read a cached part, or return ``None`` if it is missing or unreadable.""" + try: + with path.open(encoding="utf-8") as f: + return [json.loads(line) for line in f] + except (OSError, ValueError): + return None + + +def prepare_chat_data_blend(config_path: Path) -> dict[str, int]: + """Write a weighted chat blend as ``train.jsonl`` / ``validation.jsonl``; return their sizes. + + For Megatron-Bridge's chat SFT path (``distill.py --sft --sft_hf_dataset``), which renders each + conversation with the model's chat template and can restrict the loss to assistant tokens. + Each source contributes ``round(target_samples * weight / sum(weights))`` rows (every usable + row if ``target_samples`` is omitted) that the chat path accepts as text-only chat. Hugging Face + splits are buffer-shuffled unless ``shuffle: false``; ``files`` are read from the beginning. + The rows are mixed with ``seed`` and ``min(1000, n // 20)`` of them are held out for validation. + Each source is cached under ``/parts``, so an interrupted run resumes. + """ + config = load_config(config_path) + ignored = {"tokenizer", "target_tokens"} & config.keys() + if ignored: + raise ValueError(f"output_format: chat_jsonl does not use {sorted(ignored)}.") + 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.") + output_dir = Path(config["output_dir"]) + parts_dir = output_dir / "parts" + parts_dir.mkdir(parents=True, exist_ok=True) + target = config.get("target_samples") + shuffle = bool(config.get("shuffle", True)) + weight_sum = sum(float(source["weight"]) for source in config["sources"]) + + rows: list[dict[str, Any]] = [] + for index, source in enumerate(config["sources"]): + want = ( + None + if target is None + else max(1, round(int(target) * float(source["weight"]) / weight_sum)) + ) + label = source.get("split") or "files" + dataset = source["hf_dataset"].replace("/", "--") + part = parts_dir / f"{index:02d}_{dataset}_{source.get('config')}_{label}.jsonl" + got = _read_jsonl(part) + if got is None or (want is not None and len(got) < want): + got = _take_chat_rows(source, output_dir, want, shuffle) + _write_jsonl(part, got) + rows.extend(got[:want]) + + random.Random(int(config.get("seed", 1234))).shuffle(rows) + num_validation = min(1000, len(rows) // 20) + splits = {"validation": rows[:num_validation], "train": rows[num_validation:]} + 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(): + _write_jsonl(output_dir / f"{name}.jsonl", split_rows) + _copy_config(config_path, output_dir / "config.yaml") + return {name: len(split_rows) for name, split_rows in splits.items()} + + def prepare_megatron_data_blend(config_path: Path) -> list[tuple[float, str]]: """Download and tokenize the configured weighted data sources.""" config = load_config(config_path) + chat_keys = {"target_samples", "shuffle", "seed"} & config.keys() + if config.get("output_format", "megatron") != "megatron" or chat_keys: + raise ValueError( + "Tokenized blends do not use output_format: chat_jsonl, target_samples, shuffle or " + "seed; run prepare_chat_data_blend for chat JSONL." + ) output_dir = Path(config["output_dir"]) output_dir.mkdir(parents=True, exist_ok=True) target_tokens = config.get("target_tokens") @@ -165,6 +326,13 @@ def main() -> None: """Prepare a data blend from the supplied configuration.""" parser = _build_parser() args = parser.parse_args() + output_format = load_config(args.config).get("output_format", "megatron") + if output_format == "chat_jsonl": + sizes = prepare_chat_data_blend(args.config) + print(f"Prepared chat JSONL: {sizes}. See train.jsonl, validation.jsonl and config.yaml.") + return + if output_format != "megatron": + raise ValueError(f"Unknown output_format {output_format!r}; use megatron or chat_jsonl.") blend = prepare_megatron_data_blend(args.config) print(f"Prepared {len(blend)} data paths. See data_blend.txt and config.yaml in the output.") 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 f531be27ad4..ba5b6648881 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 @@ -13,8 +13,10 @@ # See the License for the specific language governing permissions and # limitations under the License. +import importlib import json import sys +from collections import Counter from pathlib import Path from unittest.mock import Mock @@ -24,6 +26,9 @@ from modelopt.torch.utils.plugins.prepare_megatron_data_blend import main +# The plugins package re-exports the function under the module's name, so import the module itself. +blend_module = importlib.import_module("modelopt.torch.utils.plugins.prepare_megatron_data_blend") + def _setup_test( tiny_tokenizer_path: str, tmp_path: Path, monkeypatch, target_tokens: int | None @@ -103,3 +108,218 @@ def test_prepare_megatron_data_blend_with_split_and_files_sources( f"competitive_programming_python_00_messages{token_suffixes[1]}", ] assert (output_dir / "config.yaml").read_bytes() == config_path.read_bytes() + + +def _chat(i: int) -> dict: + return { + "messages": [ + {"role": "user", "content": f"q{i}"}, + {"role": "assistant", "content": f"a{i}"}, + ] + } + + +def _image_chat(i: int) -> dict: + return { + "messages": [ + {"role": "user", "content": [{"type": "image", "image": "x.png"}]}, + {"role": "assistant", "content": f"a{i}"}, + ] + } + + +def _read_jsonl(path: Path) -> list[dict]: + return [json.loads(line) for line in path.read_text(encoding="utf-8").splitlines()] + + +@pytest.fixture +def chat_blend(tmp_path: Path, monkeypatch): + """Two equal-weight ``messages`` sources of 100 rows each; returns (config writer, stream counts).""" + calls = {"org/a": 0, "org/b": 0} + + def iter_rows(source, output_dir, shuffle): + calls[source["hf_dataset"]] += 1 + yield from (_chat(i) for i in range(100)) + + monkeypatch.setattr(blend_module, "_iter_rows", iter_rows) + + def write_config(**overrides) -> Path: + config = { + "output_dir": str(tmp_path / "blend"), + "output_format": "chat_jsonl", + "target_samples": 100, + "shuffle": False, + "seed": 0, + "sources": [ + { + "hf_dataset": f"org/{n}", + "split": "train", + "content_field": "messages", + "weight": 1, + } + for n in ("a", "b") + ], + } + config.update(overrides) + path = tmp_path / "config.yaml" + path.write_text(yaml.safe_dump(config), encoding="utf-8") + return path + + return write_config, calls + + +def test_to_chat_handles_all_row_shapes(): + msgs = [{"role": "user", "content": "hi"}] + assert blend_module._to_chat({"messages": msgs}, "messages") == {"messages": msgs} + assert blend_module._to_chat({"messages": []}, "messages") is None + + expected = { + "messages": [ + {"role": "user", "content": "hello"}, + {"role": "assistant", "content": "world"}, + ] + } + text = "System\nUser\nhello\nAssistant\nworld\n" + assert blend_module._to_chat({"text": text}, "text") == expected + # Any marker ends a turn, so a trailing or label marker is not trained as content. + text = "User\nhello\nAssistant\nworld\n" + assert blend_module._to_chat({"text": text}, "text") == expected + assert blend_module._to_chat({"text": "def f(): pass"}, "text") == { + "messages": [{"role": "assistant", "content": "def f(): pass"}] + } + assert blend_module._to_chat({"text": ""}, "text") is None + # Markers that parse to no turns: dropped, never trained as assistant text. + assert blend_module._to_chat({"text": "\nraw body\n"}, "text") is None + + +def test_take_chat_rows_skips_non_text_rows_and_gives_up(tmp_path: Path, monkeypatch): + source = {"hf_dataset": "org/a", "split": "train", "content_field": "messages"} + rows = [_chat(0), _image_chat(1), _chat(2), _chat(3)] + monkeypatch.setattr(blend_module, "_iter_rows", lambda *a: iter(rows)) + got = blend_module._take_chat_rows(source, tmp_path, want=3, shuffle=False) + assert [c["messages"][0]["content"] for c in got] == ["q0", "q2", "q3"] + + monkeypatch.setattr( + blend_module, "_iter_rows", lambda *a: iter([_image_chat(i) for i in range(50)]) + ) + with pytest.raises(RuntimeError, match="only 0/2 usable rows"): + blend_module._take_chat_rows(source, tmp_path, want=2, shuffle=False) + + +def test_chat_blend_splits_weighted_rows(tmp_path: Path, chat_blend, monkeypatch): + write_config, _ = chat_blend + config_path = write_config() + monkeypatch.setattr( + sys, "argv", ["prepare_megatron_data_blend.py", "--config", str(config_path)] + ) + main() + out = tmp_path / "blend" + assert len(_read_jsonl(out / "validation.jsonl")) == 5 + rows = _read_jsonl(out / "train.jsonl") + _read_jsonl(out / "validation.jsonl") + # Equal weights: each source contributes its first 50 rows, q0..q49. + assert Counter(r["messages"][0]["content"] for r in rows) == {f"q{i}": 2 for i in range(50)} + assert (out / "config.yaml").read_bytes() == config_path.read_bytes() + + +def test_chat_blend_reuses_and_rebuilds_cached_parts(tmp_path: Path, chat_blend): + write_config, calls = chat_blend + config_path = write_config() + blend_module.prepare_chat_data_blend(config_path) + first = (tmp_path / "blend" / "train.jsonl").read_text(encoding="utf-8") + blend_module.prepare_chat_data_blend(config_path) + assert calls == {"org/a": 1, "org/b": 1} # the second run streams nothing + + part = tmp_path / "blend" / "parts" / "00_org--a_None_train.jsonl" + part.write_text(part.read_text(encoding="utf-8")[:-20], encoding="utf-8") # cut mid-record + blend_module.prepare_chat_data_blend(config_path) + assert calls == {"org/a": 2, "org/b": 1} + assert (tmp_path / "blend" / "train.jsonl").read_text(encoding="utf-8") == first + + +def test_chat_blend_writes_parts_atomically(tmp_path: Path, chat_blend, monkeypatch): + write_config, _ = chat_blend + real_dumps = blend_module.json.dumps + calls = {"n": 0} + + def flaky_dumps(obj, *args, **kwargs): + calls["n"] += 1 + if calls["n"] == 75: # partway through the second source's part + raise KeyboardInterrupt + return real_dumps(obj, *args, **kwargs) + + monkeypatch.setattr(blend_module.json, "dumps", flaky_dumps) + with pytest.raises(KeyboardInterrupt): + blend_module.prepare_chat_data_blend(write_config()) + parts = tmp_path / "blend" / "parts" + assert not (parts / "01_org--b_None_train.jsonl").exists() + assert len(_read_jsonl(parts / "00_org--a_None_train.jsonl")) == 50 + + +def test_chat_blend_fails_on_a_bad_source(tmp_path: Path, chat_blend, monkeypatch): + write_config, _ = chat_blend + real_iter_rows = blend_module._iter_rows + + def iter_rows(source, output_dir, shuffle): + if source["hf_dataset"] == "org/b": + yield from real_iter_rows(source, output_dir, shuffle) + return + yield from (_chat(i) for i in range(10)) # yields some rows, then the reader fails + raise TypeError("cannot cast") + + monkeypatch.setattr(blend_module, "_iter_rows", iter_rows) + with pytest.raises(RuntimeError, match="org/a train: failed to load"): + blend_module.prepare_chat_data_blend(write_config()) + assert not (tmp_path / "blend" / "train.jsonl").exists() + + +def test_chat_blend_fails_on_a_short_or_tiny_blend(chat_blend): + write_config, _ = chat_blend + with pytest.raises(RuntimeError, match="org/a train: only 100/250 usable rows"): + blend_module.prepare_chat_data_blend(write_config(target_samples=500)) + with pytest.raises(RuntimeError, match="Too few rows"): + blend_module.prepare_chat_data_blend(write_config(target_samples=10)) + + +def test_chat_blend_reads_files_sources(tmp_path: Path, monkeypatch): + jsonl_path = tmp_path / "chat.jsonl" + jsonl_path.write_text("".join(json.dumps(_chat(i)) + "\n" for i in range(40)), encoding="utf-8") + monkeypatch.setattr(huggingface_hub, "hf_hub_download", Mock(return_value=str(jsonl_path))) + config_path = tmp_path / "config.yaml" + config_path.write_text( + yaml.safe_dump( + { + "output_dir": str(tmp_path / "blend"), + "output_format": "chat_jsonl", + "sources": [ + { + "hf_dataset": "org/a", + "files": ["data/chat.jsonl"], + "content_field": "messages", + "weight": 1, + } + ], + } + ), + encoding="utf-8", + ) + # No target_samples: every usable row, in file order before mixing. + assert blend_module.prepare_chat_data_blend(config_path) == {"validation": 2, "train": 38} + + +@pytest.mark.parametrize( + ("overrides", "match"), + [ + ({"tokenizer": "x"}, "does not use"), + ({"target_tokens": 1000}, "does not use"), + ], +) +def test_chat_blend_rejects_tokenizer_options(chat_blend, overrides, match): + write_config, _ = chat_blend + with pytest.raises(ValueError, match=match): + blend_module.prepare_chat_data_blend(write_config(**overrides)) + + +def test_tokenized_blend_rejects_chat_options(chat_blend): + write_config, _ = chat_blend + with pytest.raises(ValueError, match="Tokenized blends do not use"): + blend_module.prepare_megatron_data_blend(write_config()) From f4a96f644f5951fc0c3f9ad1dffb9d7db727775a Mon Sep 17 00:00:00 2001 From: Yue Date: Wed, 14 Oct 2026 13:55:12 +0000 Subject: [PATCH 2/6] Key cached chat parts by their selection and honor zero-row budgets Cache each source's rows under a hash of the inputs that select them (dataset, config, split, files, content_field, shuffle) and of whether the selection was capped, so a changed configuration or a removed target_samples selects afresh instead of reusing stale rows. Skip sources whose weight rounds to zero rows, and validate target_samples and the weights. Note why the Megatron-Bridge import stays local. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Yue --- examples/dataset/MEGATRON_DATA_PREP.md | 7 +- .../plugins/prepare_megatron_data_blend.py | 50 ++++++--- .../test_prepare_megatron_data_blend.py | 106 +++++++++++++++++- 3 files changed, 144 insertions(+), 19 deletions(-) diff --git a/examples/dataset/MEGATRON_DATA_PREP.md b/examples/dataset/MEGATRON_DATA_PREP.md index 4bfad9c3262..b5210487d61 100644 --- a/examples/dataset/MEGATRON_DATA_PREP.md +++ b/examples/dataset/MEGATRON_DATA_PREP.md @@ -201,11 +201,12 @@ sources: weight: 17 ``` -Each source contributes `round(target_samples * weight / sum(weights))` conversations that contain only text. +Each source contributes `round(target_samples * weight / sum(weights))` conversations that contain only text; a source that rounds to zero is skipped. A `text` row becomes one assistant turn, or is split into turns if it uses the `User` / `Assistant` chat markers. The conversations are mixed and `min(1000, n // 20)` of them are held out -for validation. A source that fails to load or runs out of usable rows stops the run. Each source is cached under -`/parts`, so an interrupted run resumes. +for validation. A source that fails to load or runs out of usable rows stops the run. Each source's selection is cached under +`/parts`, keyed by the settings that choose its rows, so an interrupted run resumes and a changed +configuration selects afresh. ## `reasoning_content` for Post-Training v3 Datasets diff --git a/modelopt/torch/utils/plugins/prepare_megatron_data_blend.py b/modelopt/torch/utils/plugins/prepare_megatron_data_blend.py index 4409d4f832b..bf1e65ddd91 100644 --- a/modelopt/torch/utils/plugins/prepare_megatron_data_blend.py +++ b/modelopt/torch/utils/plugins/prepare_megatron_data_blend.py @@ -16,7 +16,9 @@ """Prepare a weighted Megatron data blend from a YAML configuration.""" import argparse +import hashlib import json +import math import os import random import re @@ -208,7 +210,8 @@ def _take_chat_rows( Fails if the source cannot be read or runs out of usable rows, so the blend cannot change silently; gives up after ``10 * want`` rejected rows. """ - # Megatron-Bridge's chat SFT path consumes this output; keep exactly the rows it accepts. + # Megatron-Bridge's chat SFT path consumes this output; keep exactly the rows it accepts. Imported + # here, not at module scope, because the tokenizing mode does not need Megatron-Bridge installed. from megatron.bridge.data.conversation_processing import is_text_only_chat_example name = f"{source['hf_dataset']} {source.get('split') or source.get('files')}" @@ -250,13 +253,40 @@ def _read_jsonl(path: Path) -> list[dict[str, Any]] | None: return None +def _part_path(parts_dir: Path, source: dict[str, Any], shuffle: bool, capped: bool) -> Path: + """Return the cache file of a source's selected rows, keyed by every input that selects them. + + A capped part holds the first rows of the source and serves any smaller budget; an uncapped part + holds every usable row, so a budget removed between runs never reuses a capped selection. + """ + selection = {key: source.get(key) for key in ("hf_dataset", "config", "split", "files")} + selection.update(content_field=source["content_field"], shuffle=shuffle, capped=capped) + digest = hashlib.sha256(json.dumps(selection, sort_keys=True).encode()).hexdigest()[:16] + label = source.get("split") or "files" + return parts_dir / f"{source['hf_dataset'].replace('/', '--')}_{label}_{digest}.jsonl" + + +def _chat_budgets(config: dict[str, Any]) -> list[int | None]: + """Return each source's row budget, ``None`` meaning every usable row; validate the inputs.""" + weights = [float(source["weight"]) for source in config["sources"]] + if not all(math.isfinite(w) and w >= 0 for w in weights) or sum(weights) <= 0: + raise ValueError(f"Source weights must be finite, nonnegative and not all zero: {weights}.") + target = config.get("target_samples") + if target is None: + return [None] * len(weights) + if isinstance(target, bool) or not isinstance(target, int) or target <= 0: + raise ValueError(f"target_samples must be a positive integer, got {target!r}.") + return [round(target * w / sum(weights)) for w in weights] + + def prepare_chat_data_blend(config_path: Path) -> dict[str, int]: """Write a weighted chat blend as ``train.jsonl`` / ``validation.jsonl``; return their sizes. For Megatron-Bridge's chat SFT path (``distill.py --sft --sft_hf_dataset``), which renders each conversation with the model's chat template and can restrict the loss to assistant tokens. Each source contributes ``round(target_samples * weight / sum(weights))`` rows (every usable - row if ``target_samples`` is omitted) that the chat path accepts as text-only chat. Hugging Face + row if ``target_samples`` is omitted; a source rounded to zero rows is skipped) that the chat + path accepts as text-only chat. Hugging Face splits are buffer-shuffled unless ``shuffle: false``; ``files`` are read from the beginning. The rows are mixed with ``seed`` and ``min(1000, n // 20)`` of them are held out for validation. Each source is cached under ``/parts``, so an interrupted run resumes. @@ -270,20 +300,14 @@ def prepare_chat_data_blend(config_path: Path) -> dict[str, int]: output_dir = Path(config["output_dir"]) parts_dir = output_dir / "parts" parts_dir.mkdir(parents=True, exist_ok=True) - target = config.get("target_samples") shuffle = bool(config.get("shuffle", True)) - weight_sum = sum(float(source["weight"]) for source in config["sources"]) rows: list[dict[str, Any]] = [] - for index, source in enumerate(config["sources"]): - want = ( - None - if target is None - else max(1, round(int(target) * float(source["weight"]) / weight_sum)) - ) - label = source.get("split") or "files" - dataset = source["hf_dataset"].replace("/", "--") - part = parts_dir / f"{index:02d}_{dataset}_{source.get('config')}_{label}.jsonl" + for source, want in zip(config["sources"], _chat_budgets(config)): + if want == 0: + print(f"{source['hf_dataset']}: weight rounds to 0 rows, skipped", flush=True) + continue + part = _part_path(parts_dir, source, shuffle, capped=want is not None) got = _read_jsonl(part) if got is None or (want is not None and len(got) < want): got = _take_chat_rows(source, output_dir, want, shuffle) 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 ba5b6648881..354e0766209 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 @@ -229,13 +229,112 @@ def test_chat_blend_reuses_and_rebuilds_cached_parts(tmp_path: Path, chat_blend) blend_module.prepare_chat_data_blend(config_path) assert calls == {"org/a": 1, "org/b": 1} # the second run streams nothing - part = tmp_path / "blend" / "parts" / "00_org--a_None_train.jsonl" + (part,) = (tmp_path / "blend" / "parts").glob("org--a_*.jsonl") part.write_text(part.read_text(encoding="utf-8")[:-20], encoding="utf-8") # cut mid-record blend_module.prepare_chat_data_blend(config_path) assert calls == {"org/a": 2, "org/b": 1} assert (tmp_path / "blend" / "train.jsonl").read_text(encoding="utf-8") == first +@pytest.mark.parametrize( + "change", + [ + {"shuffle": True}, + { + "sources": [ + { + "hf_dataset": f"org/{n}", + "split": "test", + "content_field": "messages", + "weight": 1, + } + for n in ("a", "b") + ] + }, + ], + ids=["shuffle", "split"], +) +def test_chat_blend_does_not_reuse_parts_of_another_selection(chat_blend, change): + write_config, calls = chat_blend + blend_module.prepare_chat_data_blend(write_config()) + blend_module.prepare_chat_data_blend(write_config(**change)) + assert calls == {"org/a": 2, "org/b": 2} + + +def test_chat_blend_does_not_reuse_a_capped_part_without_a_budget(tmp_path: Path, chat_blend): + write_config, calls = chat_blend + assert blend_module.prepare_chat_data_blend(write_config()) == {"validation": 5, "train": 95} + # Every usable row of both sources, not the 50 per source cached by the capped run. + sizes = blend_module.prepare_chat_data_blend(write_config(target_samples=None)) + assert sizes == {"validation": 10, "train": 190} + assert calls == {"org/a": 2, "org/b": 2} + + +def test_chat_blend_skips_sources_rounded_to_zero_rows(tmp_path: Path, chat_blend): + write_config, calls = chat_blend + config = write_config( + sources=[ + {"hf_dataset": "org/a", "split": "train", "content_field": "messages", "weight": 999}, + {"hf_dataset": "org/b", "split": "train", "content_field": "messages", "weight": 1}, + ] + ) + assert blend_module.prepare_chat_data_blend(config) == {"validation": 5, "train": 95} + assert calls == {"org/a": 1, "org/b": 0} + + +@pytest.mark.parametrize( + ("overrides", "match"), + [ + ({"target_samples": 0}, "positive integer"), + ({"target_samples": 1.5}, "positive integer"), + ( + { + "sources": [ + { + "hf_dataset": "org/a", + "split": "train", + "content_field": "messages", + "weight": -1, + } + ] + }, + "nonnegative", + ), + ( + { + "sources": [ + { + "hf_dataset": "org/a", + "split": "train", + "content_field": "messages", + "weight": 0, + } + ] + }, + "not all zero", + ), + ( + { + "sources": [ + { + "hf_dataset": "org/a", + "split": "train", + "content_field": "messages", + "weight": float("nan"), + } + ] + }, + "finite", + ), + ], +) +def test_chat_blend_rejects_invalid_budgets(chat_blend, overrides, match): + write_config, calls = chat_blend + with pytest.raises(ValueError, match=match): + blend_module.prepare_chat_data_blend(write_config(**overrides)) + assert calls == {"org/a": 0, "org/b": 0} + + def test_chat_blend_writes_parts_atomically(tmp_path: Path, chat_blend, monkeypatch): write_config, _ = chat_blend real_dumps = blend_module.json.dumps @@ -251,8 +350,9 @@ def flaky_dumps(obj, *args, **kwargs): with pytest.raises(KeyboardInterrupt): blend_module.prepare_chat_data_blend(write_config()) parts = tmp_path / "blend" / "parts" - assert not (parts / "01_org--b_None_train.jsonl").exists() - assert len(_read_jsonl(parts / "00_org--a_None_train.jsonl")) == 50 + assert not list(parts.glob("org--b_*.jsonl")) # only the interrupted .tmp file was written + (part,) = parts.glob("org--a_*.jsonl") + assert len(_read_jsonl(part)) == 50 def test_chat_blend_fails_on_a_bad_source(tmp_path: Path, chat_blend, monkeypatch): From 4359544a0ae9ae5f3ecf28258d34bb6bfe5a0ec5 Mon Sep 17 00:00:00 2001 From: Yue Date: Wed, 14 Oct 2026 14:14:42 +0000 Subject: [PATCH 3/6] Harden chat JSONL blend inputs and match the tokenizing reader - Prepare messages with prepare_messages_for_chat_template: tool_calls arguments are normalized, and a reasoning_content setting (native by default, since the chat template is applied at training time; inline or strip) controls reasoning traces. - Reject content_field values other than messages / text, and text rows with a marker for an unknown role instead of cutting that turn out. - Fail on a source without usable rows, also when uncapped or cached empty, and say whether the source ran out or the reader gave up. - Read .jsonl.gz files and skip malformed records, as the tokenizing reader does; shuffle with the configured seed and key parts on it. - Import datasets where it is used. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Yue --- examples/dataset/MEGATRON_DATA_PREP.md | 24 +++- .../plugins/prepare_megatron_data_blend.py | 125 +++++++++++++----- .../test_prepare_megatron_data_blend.py | 125 +++++++++++++++++- 3 files changed, 230 insertions(+), 44 deletions(-) diff --git a/examples/dataset/MEGATRON_DATA_PREP.md b/examples/dataset/MEGATRON_DATA_PREP.md index b5210487d61..6f5816ca59e 100644 --- a/examples/dataset/MEGATRON_DATA_PREP.md +++ b/examples/dataset/MEGATRON_DATA_PREP.md @@ -187,7 +187,8 @@ output_dir: /path/to/chat_blend output_format: chat_jsonl target_samples: 36_800 # Optional; omit to keep every usable row. shuffle: false # Optional; take the first rows of each split in order (default: buffer shuffle). -seed: 1234 # Optional; seeds the cross-source mix and the validation hold-out. +seed: 1234 # Optional; seeds the shuffle, the cross-source mix and the validation hold-out. +reasoning_content: native # Optional; native (default), inline or strip. sources: - hf_dataset: nvidia/Nemotron-Pretraining-SFT-v1 config: Nemotron-SFT-General @@ -201,12 +202,21 @@ sources: weight: 17 ``` -Each source contributes `round(target_samples * weight / sum(weights))` conversations that contain only text; a source that rounds to zero is skipped. -A `text` row becomes one assistant turn, or is split into turns if it uses the `User` / -`Assistant` chat markers. The conversations are mixed and `min(1000, n // 20)` of them are held out -for validation. A source that fails to load or runs out of usable rows stops the run. Each source's selection is cached under -`/parts`, keyed by the settings that choose its rows, so an interrupted run resumes and a changed -configuration selects afresh. +Each source contributes `round(target_samples * weight / sum(weights))` conversations that contain only text; +a source that rounds to zero is skipped. `content_field` must be `messages` or `text`: + +- A `text` row becomes one assistant turn, or is split into turns if it uses the `User` / + `Assistant` / `System` chat markers. A row with a marker for any other role is dropped. +- `messages` are prepared as in the tokenizing mode (JSON-string `tool_calls` arguments become dicts), except + that `reasoning_content` stays on the message by default (`native`), because the chat template is applied + later, at training time. Use `inline` if your model's chat template ignores `reasoning_content`, so the + reasoning is kept as a `` block in `content` (see [below](#reasoning_content-for-post-training-v3-datasets)), + or `strip` to drop it. + +`files` may be `.jsonl` or `.jsonl.gz`; malformed records in them are skipped with a warning. The conversations are +mixed and `min(1000, n // 20)` of them are held out for validation. A source that fails to load, has no usable +rows, or runs out of them stops the run. Each source's selection is cached under `/parts`, keyed by +the settings that choose its rows, so an interrupted run resumes and a changed configuration selects afresh. ## `reasoning_content` for Post-Training v3 Datasets diff --git a/modelopt/torch/utils/plugins/prepare_megatron_data_blend.py b/modelopt/torch/utils/plugins/prepare_megatron_data_blend.py index bf1e65ddd91..eeca3978af7 100644 --- a/modelopt/torch/utils/plugins/prepare_megatron_data_blend.py +++ b/modelopt/torch/utils/plugins/prepare_megatron_data_blend.py @@ -16,6 +16,7 @@ """Prepare a weighted Megatron data blend from a YAML configuration.""" import argparse +import gzip import hashlib import json import math @@ -29,7 +30,12 @@ import huggingface_hub import yaml -from datasets import load_dataset + +from modelopt.torch.utils import ( + REASONING_CONTENT_MODES, + ReasoningContentMode, + prepare_messages_for_chat_template, +) from .megatron_preprocess_data import megatron_preprocess_data @@ -41,6 +47,8 @@ r"(User|Assistant|System)\n?(.*?)(?=|\Z)", re.DOTALL ) _CHAT_ROLES = {"User": "user", "Assistant": "assistant", "System": "system"} +_CHAT_MARKER = re.compile(r"") +_CHAT_TRAILING_MARKER = re.compile(r"\s*\Z") def load_config(path: Path) -> dict[str, Any]: @@ -70,7 +78,7 @@ def load_config(path: Path) -> dict[str, Any]: With ``output_format: chat_jsonl``, the blend is written as chat JSONL instead (see :func:`prepare_chat_data_blend`): ``tokenizer`` and ``target_tokens`` are not used, and the - optional ``target_samples``, ``shuffle`` and ``seed`` keys apply. + optional ``target_samples``, ``shuffle``, ``seed`` and ``reasoning_content`` keys apply. """ with path.open(encoding="utf-8") as stream: return cast("dict[str, Any]", yaml.safe_load(stream)) @@ -166,49 +174,79 @@ def _prepare_sources( return blend -def _to_chat(row: dict[str, Any], content_field: str) -> dict[str, Any] | None: +def _to_chat( + row: dict[str, Any], content_field: str, reasoning_content: ReasoningContentMode = "native" +) -> dict[str, Any] | None: """Convert a source row to ``{"messages": [...]}``, or ``None`` if it has no usable content. + ``messages`` go through :func:`prepare_messages_for_chat_template` with ``reasoning_content``. A ``text`` row in the NeMo chat format (``User`` / ``Assistant`` markers) is split into turns; any other ``text`` row becomes one assistant turn, so all of it is trained. """ if content_field == "messages": - return {"messages": row["messages"]} if row.get("messages") else None + messages = row.get("messages") + if not messages: + return None + messages = prepare_messages_for_chat_template(messages, reasoning_content=reasoning_content) + return {"messages": messages} text = row.get("text") or "" if not text: return None if " delimiters as assistant content. - return {"messages": turns} if turns else None + # would train the raw delimiters as assistant content. A marker without a known + # role (e.g. Tool) would silently cut its turn out of the conversation, so drop the row; only + # a bare delimiter at the very end needs no role. + roleless = len(_CHAT_MARKER.findall(text)) - len(matches) + if not turns or roleless > (1 if _CHAT_TRAILING_MARKER.search(text) else 0): + return None + return {"messages": turns} return {"messages": [{"role": "assistant", "content": text}]} -def _iter_rows(source: dict[str, Any], output_dir: Path, shuffle: bool) -> Iterator[dict[str, Any]]: +def _iter_rows( + source: dict[str, Any], output_dir: Path, shuffle: bool, seed: int +) -> Iterator[dict[str, Any]]: """Yield the rows of a source: its ``files`` in order, or its Hugging Face split, streamed.""" if "files" in source: for path in _download_files(source, output_dir): - with open(path, encoding="utf-8") as f: - yield from (json.loads(line) for line in f if line.strip()) + opener = gzip.open if path.endswith(".gz") else open + with opener(path, "rt", encoding="utf-8") as f: + for line in f: + if not line.strip(): + continue + try: + yield json.loads(line) + except ValueError as e: + # As in the tokenizing reader: some Nemotron files carry a malformed record. + print(f"[WARN] Skipping malformed JSONL record in {path}: {e}", flush=True) return + from datasets import load_dataset + dataset = load_dataset( source["hf_dataset"], source.get("config"), split=source["split"], streaming=True ) - yield from dataset.shuffle(seed=42) if shuffle else dataset + yield from dataset.shuffle(seed=seed) if shuffle else dataset def _take_chat_rows( - source: dict[str, Any], output_dir: Path, want: int | None, shuffle: bool + source: dict[str, Any], + output_dir: Path, + want: int | None, + shuffle: bool, + seed: int, + reasoning_content: ReasoningContentMode, ) -> list[dict[str, Any]]: """Return the first ``want`` usable chat rows of a source (all of them if ``want`` is None). - Fails if the source cannot be read or runs out of usable rows, so the blend cannot change - silently; gives up after ``10 * want`` rejected rows. + Fails if the source cannot be read, has no usable row, or runs out of them, so the blend cannot + change silently; gives up after ``10 * want`` rejected rows. """ # Megatron-Bridge's chat SFT path consumes this output; keep exactly the rows it accepts. Imported # here, not at module scope, because the tokenizing mode does not need Megatron-Bridge installed. @@ -218,8 +256,8 @@ def _take_chat_rows( rows: list[dict[str, Any]] = [] rejected = 0 try: - for row in _iter_rows(source, output_dir, shuffle): - chat = _to_chat(row, source["content_field"]) + for row in _iter_rows(source, output_dir, shuffle, seed): + chat = _to_chat(row, source["content_field"], reasoning_content) if chat is None or not is_text_only_chat_example(chat): rejected += 1 if want is not None and rejected > 10 * want: @@ -230,8 +268,11 @@ def _take_chat_rows( break except Exception as e: # e.g. a gated split, or a schema the streaming reader cannot cast raise RuntimeError(f"{name}: failed to load") from e - if want is not None and len(rows) < want: - raise RuntimeError(f"{name}: only {len(rows)}/{want} usable rows") + if not rows or (want is not None and len(rows) < want): + gave_up = want is not None and rejected > 10 * want + reason = f"gave up after {rejected} rejected rows" if gave_up else "source exhausted" + budget = "any" if want is None else want + raise RuntimeError(f"{name}: only {len(rows)}/{budget} usable rows ({reason})") print(f"{name}: {len(rows)} rows ({rejected} rejected)", flush=True) return rows @@ -253,14 +294,16 @@ def _read_jsonl(path: Path) -> list[dict[str, Any]] | None: return None -def _part_path(parts_dir: Path, source: dict[str, Any], shuffle: bool, capped: bool) -> Path: +def _part_path( + parts_dir: Path, source: dict[str, Any], settings: dict[str, Any], capped: bool +) -> Path: """Return the cache file of a source's selected rows, keyed by every input that selects them. A capped part holds the first rows of the source and serves any smaller budget; an uncapped part holds every usable row, so a budget removed between runs never reuses a capped selection. """ selection = {key: source.get(key) for key in ("hf_dataset", "config", "split", "files")} - selection.update(content_field=source["content_field"], shuffle=shuffle, capped=capped) + selection.update(settings, content_field=source["content_field"], capped=capped) digest = hashlib.sha256(json.dumps(selection, sort_keys=True).encode()).hexdigest()[:16] label = source.get("split") or "files" return parts_dir / f"{source['hf_dataset'].replace('/', '--')}_{label}_{digest}.jsonl" @@ -286,10 +329,14 @@ def prepare_chat_data_blend(config_path: Path) -> dict[str, int]: conversation with the model's chat template and can restrict the loss to assistant tokens. Each source contributes ``round(target_samples * weight / sum(weights))`` rows (every usable row if ``target_samples`` is omitted; a source rounded to zero rows is skipped) that the chat - path accepts as text-only chat. Hugging Face - splits are buffer-shuffled unless ``shuffle: false``; ``files`` are read from the beginning. - The rows are mixed with ``seed`` and ``min(1000, n // 20)`` of them are held out for validation. - Each source is cached under ``/parts``, so an interrupted run resumes. + path accepts as text-only chat. ``content_field`` must be ``messages`` or ``text``; ``messages`` + keep ``reasoning_content`` for the chat template by default (``reasoning_content: native``; set + ``inline`` for a template that ignores the field, or ``strip``). Hugging Face splits are + buffer-shuffled with ``seed`` unless ``shuffle: false``; ``files`` (``.jsonl`` or ``.jsonl.gz``) + are read from the beginning, skipping malformed records. 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. """ config = load_config(config_path) ignored = {"tokenizer", "target_tokens"} & config.keys() @@ -297,24 +344,40 @@ def prepare_chat_data_blend(config_path: Path) -> dict[str, int]: raise ValueError(f"output_format: chat_jsonl does not use {sorted(ignored)}.") 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"} + if fields: + raise ValueError( + f"output_format: chat_jsonl reads content_field messages or text: {fields}." + ) + reasoning_content = config.get("reasoning_content", "native") + if reasoning_content not in REASONING_CONTENT_MODES: + raise ValueError( + f"reasoning_content must be one of {sorted(REASONING_CONTENT_MODES)}, " + f"got {reasoning_content!r}." + ) output_dir = Path(config["output_dir"]) parts_dir = output_dir / "parts" parts_dir.mkdir(parents=True, exist_ok=True) shuffle = bool(config.get("shuffle", True)) + seed = int(config.get("seed", 1234)) + # Everything besides the source itself that decides which rows a part holds. + settings = {"shuffle": shuffle, "reasoning_content": reasoning_content} + if shuffle: + settings["seed"] = seed rows: list[dict[str, Any]] = [] for source, want in zip(config["sources"], _chat_budgets(config)): if want == 0: print(f"{source['hf_dataset']}: weight rounds to 0 rows, skipped", flush=True) continue - part = _part_path(parts_dir, source, shuffle, capped=want is not None) + part = _part_path(parts_dir, source, settings, capped=want is not None) got = _read_jsonl(part) - if got is None or (want is not None and len(got) < want): - got = _take_chat_rows(source, output_dir, want, shuffle) + if not got or (want is not None and len(got) < want): + got = _take_chat_rows(source, output_dir, want, shuffle, seed, reasoning_content) _write_jsonl(part, got) rows.extend(got[:want]) - random.Random(int(config.get("seed", 1234))).shuffle(rows) + random.Random(seed).shuffle(rows) num_validation = min(1000, len(rows) // 20) splits = {"validation": rows[:num_validation], "train": rows[num_validation:]} if not splits["train"] or not splits["validation"]: @@ -328,11 +391,11 @@ def prepare_chat_data_blend(config_path: Path) -> dict[str, int]: def prepare_megatron_data_blend(config_path: Path) -> list[tuple[float, str]]: """Download and tokenize the configured weighted data sources.""" config = load_config(config_path) - chat_keys = {"target_samples", "shuffle", "seed"} & config.keys() + chat_keys = {"target_samples", "shuffle", "seed", "reasoning_content"} & config.keys() if config.get("output_format", "megatron") != "megatron" or chat_keys: raise ValueError( - "Tokenized blends do not use output_format: chat_jsonl, target_samples, shuffle or " - "seed; run prepare_chat_data_blend for chat JSONL." + "Tokenized blends do not use output_format: chat_jsonl, target_samples, shuffle, seed " + "or reasoning_content; run prepare_chat_data_blend for chat JSONL." ) output_dir = Path(config["output_dir"]) output_dir.mkdir(parents=True, exist_ok=True) 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 354e0766209..b3f0fdd0f70 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 @@ -13,6 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import gzip import importlib import json import sys @@ -137,7 +138,7 @@ def chat_blend(tmp_path: Path, monkeypatch): """Two equal-weight ``messages`` sources of 100 rows each; returns (config writer, stream counts).""" calls = {"org/a": 0, "org/b": 0} - def iter_rows(source, output_dir, shuffle): + def iter_rows(source, output_dir, shuffle, seed): calls[source["hf_dataset"]] += 1 yield from (_chat(i) for i in range(100)) @@ -190,20 +191,47 @@ def test_to_chat_handles_all_row_shapes(): assert blend_module._to_chat({"text": ""}, "text") is None # Markers that parse to no turns: dropped, never trained as assistant text. assert blend_module._to_chat({"text": "\nraw body\n"}, "text") is None + # A role without a mapping would cut its turn out of the middle: the whole row is dropped. + text = "User\nq\nTool\nresult\nAssistant\na\n" + assert blend_module._to_chat({"text": text}, "text") is None + + +@pytest.mark.parametrize( + ("mode", "assistant"), + [ + ("native", {"role": "assistant", "content": "4", "reasoning_content": "2+2"}), + ("inline", {"role": "assistant", "content": "\n2+2\n\n4"}), + ("strip", {"role": "assistant", "content": "4"}), + ], +) +def test_to_chat_prepares_messages_like_the_tokenizing_mode(mode, assistant): + messages = [ + {"role": "user", "content": "2+2?"}, + {"role": "assistant", "content": "4", "reasoning_content": "2+2"}, + { + "role": "assistant", + "content": "", + "tool_calls": [{"function": {"name": "add", "arguments": '{"a": 2, "b": 2}'}}], + }, + ] + chat = blend_module._to_chat({"messages": messages}, "messages", mode) + assert chat["messages"][1] == assistant + # JSON-string tool-call arguments become dicts, as the Nemotron v3 chat templates expect. + assert chat["messages"][2]["tool_calls"][0]["function"]["arguments"] == {"a": 2, "b": 2} def test_take_chat_rows_skips_non_text_rows_and_gives_up(tmp_path: Path, monkeypatch): source = {"hf_dataset": "org/a", "split": "train", "content_field": "messages"} rows = [_chat(0), _image_chat(1), _chat(2), _chat(3)] monkeypatch.setattr(blend_module, "_iter_rows", lambda *a: iter(rows)) - got = blend_module._take_chat_rows(source, tmp_path, want=3, shuffle=False) + got = blend_module._take_chat_rows(source, tmp_path, 3, False, 0, "native") assert [c["messages"][0]["content"] for c in got] == ["q0", "q2", "q3"] monkeypatch.setattr( blend_module, "_iter_rows", lambda *a: iter([_image_chat(i) for i in range(50)]) ) - with pytest.raises(RuntimeError, match="only 0/2 usable rows"): - blend_module._take_chat_rows(source, tmp_path, want=2, shuffle=False) + with pytest.raises(RuntimeError, match=r"only 0/2 usable rows \(gave up after 21 rejected"): + blend_module._take_chat_rows(source, tmp_path, 2, False, 0, "native") def test_chat_blend_splits_weighted_rows(tmp_path: Path, chat_blend, monkeypatch): @@ -261,6 +289,59 @@ def test_chat_blend_does_not_reuse_parts_of_another_selection(chat_blend, change assert calls == {"org/a": 2, "org/b": 2} +@pytest.mark.parametrize( + ("first", "second"), + [ + ({"shuffle": True, "seed": 0}, {"shuffle": True, "seed": 1}), + ({}, {"reasoning_content": "inline"}), + ( + { + "sources": [ + { + "hf_dataset": "org/a", + "files": ["a.jsonl"], + "content_field": "messages", + "weight": 1, + } + ] + }, + { + "sources": [ + { + "hf_dataset": "org/a", + "files": ["b.jsonl"], + "content_field": "messages", + "weight": 1, + } + ] + }, + ), + ], + ids=["seed", "reasoning_content", "files"], +) +def test_chat_blend_reselects_after_a_settings_change(chat_blend, first, second): + write_config, calls = chat_blend + blend_module.prepare_chat_data_blend(write_config(**first)) + blend_module.prepare_chat_data_blend(write_config(**first)) # cached + blend_module.prepare_chat_data_blend(write_config(**second)) + assert calls["org/a"] == 2 + + +def test_chat_blend_fails_on_a_source_without_usable_rows(chat_blend, monkeypatch): + write_config, _ = chat_blend + + def iter_rows(source, output_dir, shuffle, seed): + make_row = _chat if source["hf_dataset"] == "org/a" else _image_chat + yield from (make_row(i) for i in range(100)) + + monkeypatch.setattr(blend_module, "_iter_rows", iter_rows) + # Without target_samples, an all-rejected source must not drop out of the blend silently. + with pytest.raises( + RuntimeError, match=r"org/b train: only 0/any usable rows \(source exhausted" + ): + blend_module.prepare_chat_data_blend(write_config(target_samples=None)) + + def test_chat_blend_does_not_reuse_a_capped_part_without_a_budget(tmp_path: Path, chat_blend): write_config, calls = chat_blend assert blend_module.prepare_chat_data_blend(write_config()) == {"validation": 5, "train": 95} @@ -359,9 +440,9 @@ def test_chat_blend_fails_on_a_bad_source(tmp_path: Path, chat_blend, monkeypatc write_config, _ = chat_blend real_iter_rows = blend_module._iter_rows - def iter_rows(source, output_dir, shuffle): + def iter_rows(source, output_dir, shuffle, seed): if source["hf_dataset"] == "org/b": - yield from real_iter_rows(source, output_dir, shuffle) + yield from real_iter_rows(source, output_dir, shuffle, seed) return yield from (_chat(i) for i in range(10)) # yields some rows, then the reader fails raise TypeError("cannot cast") @@ -411,6 +492,21 @@ def test_chat_blend_reads_files_sources(tmp_path: Path, monkeypatch): [ ({"tokenizer": "x"}, "does not use"), ({"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``. + ( + { + "sources": [ + { + "hf_dataset": "org/a", + "split": "train", + "content_field": "conversations", + "weight": 1, + } + ] + }, + "content_field messages or text", + ), ], ) def test_chat_blend_rejects_tokenizer_options(chat_blend, overrides, match): @@ -423,3 +519,20 @@ def test_tokenized_blend_rejects_chat_options(chat_blend): write_config, _ = chat_blend with pytest.raises(ValueError, match="Tokenized blends do not use"): blend_module.prepare_megatron_data_blend(write_config()) + + +def test_chat_blend_reads_gzip_files_and_skips_malformed_records(tmp_path: Path, monkeypatch): + plain, packed = tmp_path / "a.jsonl", tmp_path / "b.jsonl.gz" + lines = [json.dumps(_chat(i)) for i in range(20)] + plain.write_text( + "\n".join([*lines[:10], '{"messages": [', *lines[10:]]) + "\n", encoding="utf-8" + ) + with gzip.open(packed, "wt", encoding="utf-8") as f: + f.writelines(json.dumps(_chat(i)) + "\n" for i in range(20, 40)) + paths = {"data/a.jsonl": str(plain), "data/b.jsonl.gz": str(packed)} + monkeypatch.setattr( + huggingface_hub, "hf_hub_download", lambda repo_id, filename, **kwargs: paths[filename] + ) + source = {"hf_dataset": "org/a", "files": list(paths), "content_field": "messages"} + rows = list(blend_module._iter_rows(source, tmp_path / "blend", False, 0)) + assert [r["messages"][0]["content"] for r in rows] == [f"q{i}" for i in range(40)] From 4ecf57456f2498b0375984b917b3ab70ffe32458 Mon Sep 17 00:00:00 2001 From: Yue Date: Wed, 14 Oct 2026 18:07:11 +0000 Subject: [PATCH 4/6] Tighten chat JSONL row checks and consolidate the tests Drop text rows with an empty non-System turn or a role marker not followed by a newline, and any row without an assistant turn; keep a row's tools for the chat template; require target_samples for Hugging Face split sources; make each entry point reject the other's output_format. Tests consolidated per review. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Yue --- examples/dataset/MEGATRON_DATA_PREP.md | 12 +- .../plugins/prepare_megatron_data_blend.py | 106 +++-- .../test_prepare_megatron_data_blend.py | 384 ++++++------------ 3 files changed, 187 insertions(+), 315 deletions(-) diff --git a/examples/dataset/MEGATRON_DATA_PREP.md b/examples/dataset/MEGATRON_DATA_PREP.md index 6f5816ca59e..e22a855779c 100644 --- a/examples/dataset/MEGATRON_DATA_PREP.md +++ b/examples/dataset/MEGATRON_DATA_PREP.md @@ -185,7 +185,7 @@ when it loads the data. Run it in an environment with Megatron-Bridge, whose tex ```yaml output_dir: /path/to/chat_blend output_format: chat_jsonl -target_samples: 36_800 # Optional; omit to keep every usable row. +target_samples: 36_800 # Required with split sources; omit only for files sources, to keep every row. shuffle: false # Optional; take the first rows of each split in order (default: buffer shuffle). seed: 1234 # Optional; seeds the shuffle, the cross-source mix and the validation hold-out. reasoning_content: native # Optional; native (default), inline or strip. @@ -202,12 +202,13 @@ sources: weight: 17 ``` -Each source contributes `round(target_samples * weight / sum(weights))` conversations that contain only text; -a source that rounds to zero is skipped. `content_field` must be `messages` or `text`: +Each source contributes `round(target_samples * weight / sum(weights))` conversations that contain only text +and have an assistant turn; a source that rounds to zero is skipped. `content_field` must be `messages` or `text`: - A `text` row becomes one assistant turn, or is split into turns if it uses the `User` / - `Assistant` / `System` chat markers. A row with a marker for any other role is dropped. -- `messages` are prepared as in the tokenizing mode (JSON-string `tool_calls` arguments become dicts), except + `Assistant` / `System` chat markers. A row with a marker for any other role (such as a SteerLM + label line) or with an empty turn other than `System` is dropped. +- `messages` and a row's `tools` are prepared as in the tokenizing mode (JSON-string `tool_calls` arguments become dicts), except that `reasoning_content` stays on the message by default (`native`), because the chat template is applied later, at training time. Use `inline` if your model's chat template ignores `reasoning_content`, so the reasoning is kept as a `` block in `content` (see [below](#reasoning_content-for-post-training-v3-datasets)), @@ -217,6 +218,7 @@ a source that rounds to zero is skipped. `content_field` must be `messages` or ` mixed and `min(1000, n // 20)` of them are held out for validation. A source that fails to load, has no usable rows, or runs out of them stops the run. Each source's selection is cached under `/parts`, keyed by 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. ## `reasoning_content` for Post-Training v3 Datasets diff --git a/modelopt/torch/utils/plugins/prepare_megatron_data_blend.py b/modelopt/torch/utils/plugins/prepare_megatron_data_blend.py index eeca3978af7..87ff2544d80 100644 --- a/modelopt/torch/utils/plugins/prepare_megatron_data_blend.py +++ b/modelopt/torch/utils/plugins/prepare_megatron_data_blend.py @@ -42,9 +42,11 @@ __all__ = ["prepare_chat_data_blend", "prepare_megatron_data_blend"] # A turn of a ``text`` row in the NeMo chat format: ``User\n...``. Any marker ends the -# turn, so a trailing delimiter or a label marker is not kept as content. +# turn, so a trailing delimiter is not kept as content. The role must be followed by a newline, a +# marker or the end, so ``Userguide`` is no role at all. _CHAT_TURN = re.compile( - r"(User|Assistant|System)\n?(.*?)(?=|\Z)", re.DOTALL + r"(User|Assistant|System)(?:\n|(?=|\Z))(.*?)(?=|\Z)", + re.DOTALL, ) _CHAT_ROLES = {"User": "user", "Assistant": "assistant", "System": "system"} _CHAT_MARKER = re.compile(r"") @@ -177,37 +179,49 @@ def _prepare_sources( def _to_chat( row: dict[str, Any], content_field: str, reasoning_content: ReasoningContentMode = "native" ) -> dict[str, Any] | None: - """Convert a source row to ``{"messages": [...]}``, or ``None`` if it has no usable content. + """Convert a source row to ``{"messages": [...]}``, or ``None`` if it has nothing to train on. - ``messages`` go through :func:`prepare_messages_for_chat_template` with ``reasoning_content``. - A ``text`` row in the NeMo chat format (``User`` / ``Assistant`` markers) is split - into turns; any other ``text`` row becomes one assistant turn, so all of it is trained. + ``messages`` go through :func:`prepare_messages_for_chat_template` with ``reasoning_content``, + and a non-empty ``tools`` list is kept for the chat template. A ``text`` row in the NeMo chat + format (``User`` / ``Assistant`` markers) is split into turns; any other ``text`` + row becomes one assistant turn, so all of it is trained. """ if content_field == "messages": messages = row.get("messages") if not messages: return None messages = prepare_messages_for_chat_template(messages, reasoning_content=reasoning_content) - return {"messages": messages} - text = row.get("text") or "" - if not text: - return None - if " delimiters as assistant content. A marker without a known + # role (a Tool turn, or a SteerLM label line) or a known turn left empty would silently cut + # it out of the conversation, so drop the row. Only a bare delimiter at the very end needs no + # role, and only the System block, which real rows leave empty, may be empty. + roleless = len(_CHAT_MARKER.findall(text)) - len(matches) + if roleless > (1 if _CHAT_TRAILING_MARKER.search(text) else 0) or any( + not content.strip() and role != "System" for role, content in matches + ): + return None turns = [ {"role": _CHAT_ROLES[role], "content": content.strip()} for role, content in matches if content.strip() ] - # Markers that parse to no turns make the row malformed, not plain text: falling through - # would train the raw delimiters as assistant content. A marker without a known - # role (e.g. Tool) would silently cut its turn out of the conversation, so drop the row; only - # a bare delimiter at the very end needs no role. - roleless = len(_CHAT_MARKER.findall(text)) - len(matches) - if not turns or roleless > (1 if _CHAT_TRAILING_MARKER.search(text) else 0): - return None - return {"messages": turns} - return {"messages": [{"role": "assistant", "content": text}]} + 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"]): + return None + return chat def _iter_rows( @@ -267,7 +281,7 @@ def _take_chat_rows( if want is not None and len(rows) >= want: break except Exception as e: # e.g. a gated split, or a schema the streaming reader cannot cast - raise RuntimeError(f"{name}: failed to load") from e + raise RuntimeError(f"{name}: failed to read or convert") from e if not rows or (want is not None and len(rows) < want): gave_up = want is not None and rejected > 10 * want reason = f"gave up after {rejected} rejected rows" if gave_up else "source exhausted" @@ -305,7 +319,7 @@ def _part_path( selection = {key: source.get(key) for key in ("hf_dataset", "config", "split", "files")} selection.update(settings, content_field=source["content_field"], capped=capped) digest = hashlib.sha256(json.dumps(selection, sort_keys=True).encode()).hexdigest()[:16] - label = source.get("split") or "files" + label = source.get("split") or Path(source["files"][0]).name.split(".")[0] return parts_dir / f"{source['hf_dataset'].replace('/', '--')}_{label}_{digest}.jsonl" @@ -316,6 +330,9 @@ def _chat_budgets(config: dict[str, Any]) -> list[int | None]: raise ValueError(f"Source weights must be finite, nonnegative and not all zero: {weights}.") target = config.get("target_samples") if target is None: + if any("split" in source for source in config["sources"]): + # Every row of a split would be held in memory, which a large split does not fit. + raise ValueError("target_samples is required when a source is a Hugging Face split.") return [None] * len(weights) if isinstance(target, bool) or not isinstance(target, int) or target <= 0: raise ValueError(f"target_samples must be a positive integer, got {target!r}.") @@ -327,18 +344,23 @@ def prepare_chat_data_blend(config_path: Path) -> dict[str, int]: For Megatron-Bridge's chat SFT path (``distill.py --sft --sft_hf_dataset``), which renders each conversation with the model's chat template and can restrict the loss to assistant tokens. - Each source contributes ``round(target_samples * weight / sum(weights))`` rows (every usable - row if ``target_samples`` is omitted; a source rounded to zero rows is skipped) that the chat - path accepts as text-only chat. ``content_field`` must be ``messages`` or ``text``; ``messages`` - keep ``reasoning_content`` for the chat template by default (``reasoning_content: native``; set - ``inline`` for a template that ignores the field, or ``strip``). Hugging Face splits are - buffer-shuffled with ``seed`` unless ``shuffle: false``; ``files`` (``.jsonl`` or ``.jsonl.gz``) - are read from the beginning, skipping malformed records. 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. + Each source contributes ``round(target_samples * weight / sum(weights))`` rows (a source rounded + to zero rows is skipped) that the chat path accepts as text-only chat; ``target_samples`` may be + omitted only for ``files`` sources, to keep every usable row. ``content_field`` must be + ``messages`` or ``text``; ``messages`` keep ``reasoning_content`` for the chat template by + default (``reasoning_content: native``; set ``inline`` for a template that ignores the field, or + ``strip``). Hugging Face splits are buffer-shuffled with ``seed`` unless ``shuffle: false``; + ``files`` (``.jsonl`` or ``.jsonl.gz``) are read from the beginning, skipping malformed records. + 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. """ config = load_config(config_path) + if config.get("output_format") != "chat_jsonl": + raise ValueError( + "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)}.") @@ -355,6 +377,7 @@ def prepare_chat_data_blend(config_path: Path) -> dict[str, int]: f"reasoning_content must be one of {sorted(REASONING_CONTENT_MODES)}, " f"got {reasoning_content!r}." ) + budgets = _chat_budgets(config) output_dir = Path(config["output_dir"]) parts_dir = output_dir / "parts" parts_dir.mkdir(parents=True, exist_ok=True) @@ -366,7 +389,7 @@ def prepare_chat_data_blend(config_path: Path) -> dict[str, int]: settings["seed"] = seed rows: list[dict[str, Any]] = [] - for source, want in zip(config["sources"], _chat_budgets(config)): + for source, want in zip(config["sources"], budgets): if want == 0: print(f"{source['hf_dataset']}: weight rounds to 0 rows, skipped", flush=True) continue @@ -391,12 +414,15 @@ def prepare_chat_data_blend(config_path: Path) -> dict[str, int]: def prepare_megatron_data_blend(config_path: Path) -> list[tuple[float, str]]: """Download and tokenize the configured weighted data sources.""" config = load_config(config_path) - chat_keys = {"target_samples", "shuffle", "seed", "reasoning_content"} & config.keys() - if config.get("output_format", "megatron") != "megatron" or chat_keys: + output_format = config.get("output_format", "megatron") + if output_format != "megatron": raise ValueError( - "Tokenized blends do not use output_format: chat_jsonl, target_samples, shuffle, seed " - "or reasoning_content; run prepare_chat_data_blend for chat JSONL." + 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() + if chat_keys: + raise ValueError(f"Tokenized blends do not use {sorted(chat_keys)}.") output_dir = Path(config["output_dir"]) output_dir.mkdir(parents=True, exist_ok=True) target_tokens = config.get("target_tokens") @@ -413,13 +439,11 @@ def main() -> None: """Prepare a data blend from the supplied configuration.""" parser = _build_parser() args = parser.parse_args() - output_format = load_config(args.config).get("output_format", "megatron") - if output_format == "chat_jsonl": + # Each entry point rejects a config written for the other. + if load_config(args.config).get("output_format") == "chat_jsonl": sizes = prepare_chat_data_blend(args.config) print(f"Prepared chat JSONL: {sizes}. See train.jsonl, validation.jsonl and config.yaml.") return - if output_format != "megatron": - raise ValueError(f"Unknown output_format {output_format!r}; use megatron or chat_jsonl.") blend = prepare_megatron_data_blend(args.config) print(f"Prepared {len(blend)} data paths. See data_blend.txt and config.yaml in the output.") 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 b3f0fdd0f70..733a57905c0 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 @@ -129,14 +129,25 @@ def _image_chat(i: int) -> dict: } +def _source(name: str, **fields) -> dict: + source = { + "hf_dataset": f"org/{name}", + "split": "train", + "content_field": "messages", + "weight": 1, + } + source.update(fields) + return {k: v for k, v in source.items() if v is not None} + + def _read_jsonl(path: Path) -> list[dict]: return [json.loads(line) for line in path.read_text(encoding="utf-8").splitlines()] @pytest.fixture def chat_blend(tmp_path: Path, monkeypatch): - """Two equal-weight ``messages`` sources of 100 rows each; returns (config writer, stream counts).""" - calls = {"org/a": 0, "org/b": 0} + """Sources of 100 ``messages`` rows each; returns (config writer, stream counts per source).""" + calls = Counter() def iter_rows(source, output_dir, shuffle, seed): calls[source["hf_dataset"]] += 1 @@ -151,15 +162,7 @@ def write_config(**overrides) -> Path: "target_samples": 100, "shuffle": False, "seed": 0, - "sources": [ - { - "hf_dataset": f"org/{n}", - "split": "train", - "content_field": "messages", - "weight": 1, - } - for n in ("a", "b") - ], + "sources": [_source("a"), _source("b")], } config.update(overrides) path = tmp_path / "config.yaml" @@ -170,9 +173,20 @@ def write_config(**overrides) -> Path: def test_to_chat_handles_all_row_shapes(): - msgs = [{"role": "user", "content": "hi"}] + 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. + assert blend_module._to_chat({"messages": msgs[:1]}, "messages") is None + # Tool definitions reach the chat template, as in the tokenizing mode. + tools = [{"type": "function", "function": {"name": "add"}}] + chat = blend_module._to_chat({"messages": msgs, "tools": tools}, "messages") + assert chat == {"messages": msgs, "tools": tools} + # Unlike the tokenizing mode, reasoning_content stays on the message by default (native). + reasoning = [{"role": "assistant", "content": "4", "reasoning_content": "2+2"}] + assert blend_module._to_chat({"messages": reasoning}, "messages")["messages"] == reasoning + inline = blend_module._to_chat({"messages": reasoning}, "messages", "inline") + assert inline["messages"] == [{"role": "assistant", "content": "\n2+2\n\n4"}] expected = { "messages": [ @@ -182,9 +196,8 @@ def test_to_chat_handles_all_row_shapes(): } text = "System\nUser\nhello\nAssistant\nworld\n" assert blend_module._to_chat({"text": text}, "text") == expected - # Any marker ends a turn, so a trailing or label marker is not trained as content. - text = "User\nhello\nAssistant\nworld\n" - assert blend_module._to_chat({"text": text}, "text") == expected + # Any marker ends a turn, so a trailing marker is not trained as content. + assert blend_module._to_chat({"text": text + ""}, "text") == expected assert blend_module._to_chat({"text": "def f(): pass"}, "text") == { "messages": [{"role": "assistant", "content": "def f(): pass"}] } @@ -194,49 +207,39 @@ def test_to_chat_handles_all_row_shapes(): # A role without a mapping would cut its turn out of the middle: the whole row is dropped. text = "User\nq\nTool\nresult\nAssistant\na\n" assert blend_module._to_chat({"text": text}, "text") is None + for text in ( + # A SteerLM label line is no role either. + "User\nq\nAssistant\nquality:4\na\n", + # A role must end at a newline: this is no User turn. + "Userguide\nq\nAssistant\na\n", + # An empty known turn would be cut out of the middle of the conversation. + "User\nq1\nAssistant\nUser\nq2\nAssistant\na2\n", + # No assistant turn to train on. + "User\nq\nAssistant\n", + ): + assert blend_module._to_chat({"text": text}, "text") is None, text -@pytest.mark.parametrize( - ("mode", "assistant"), - [ - ("native", {"role": "assistant", "content": "4", "reasoning_content": "2+2"}), - ("inline", {"role": "assistant", "content": "\n2+2\n\n4"}), - ("strip", {"role": "assistant", "content": "4"}), - ], -) -def test_to_chat_prepares_messages_like_the_tokenizing_mode(mode, assistant): - messages = [ - {"role": "user", "content": "2+2?"}, - {"role": "assistant", "content": "4", "reasoning_content": "2+2"}, - { - "role": "assistant", - "content": "", - "tool_calls": [{"function": {"name": "add", "arguments": '{"a": 2, "b": 2}'}}], - }, - ] - chat = blend_module._to_chat({"messages": messages}, "messages", mode) - assert chat["messages"][1] == assistant - # JSON-string tool-call arguments become dicts, as the Nemotron v3 chat templates expect. - assert chat["messages"][2]["tool_calls"][0]["function"]["arguments"] == {"a": 2, "b": 2} - - -def test_take_chat_rows_skips_non_text_rows_and_gives_up(tmp_path: Path, monkeypatch): - source = {"hf_dataset": "org/a", "split": "train", "content_field": "messages"} +def test_take_chat_rows_skips_non_text_rows_and_fails_without_enough(tmp_path: Path, monkeypatch): + source = _source("a") rows = [_chat(0), _image_chat(1), _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") assert [c["messages"][0]["content"] for c in got] == ["q0", "q2", "q3"] - monkeypatch.setattr( - blend_module, "_iter_rows", lambda *a: iter([_image_chat(i) for i in range(50)]) - ) + images = [_image_chat(i) for i in range(50)] + monkeypatch.setattr(blend_module, "_iter_rows", lambda *a: iter(images)) with pytest.raises(RuntimeError, match=r"only 0/2 usable rows \(gave up after 21 rejected"): blend_module._take_chat_rows(source, tmp_path, 2, False, 0, "native") + # Without a budget, an all-rejected source must not drop out of the blend silently. + with pytest.raises(RuntimeError, match=r"only 0/any usable rows \(source exhausted"): + blend_module._take_chat_rows(source, tmp_path, None, False, 0, "native") def test_chat_blend_splits_weighted_rows(tmp_path: Path, chat_blend, monkeypatch): - write_config, _ = chat_blend - config_path = write_config() + write_config, calls = chat_blend + sources = [_source("a", weight=3), _source("b"), _source("c", weight=0)] + config_path = write_config(sources=sources) monkeypatch.setattr( sys, "argv", ["prepare_megatron_data_blend.py", "--config", str(config_path)] ) @@ -244,8 +247,10 @@ def test_chat_blend_splits_weighted_rows(tmp_path: Path, chat_blend, monkeypatch out = tmp_path / "blend" assert len(_read_jsonl(out / "validation.jsonl")) == 5 rows = _read_jsonl(out / "train.jsonl") + _read_jsonl(out / "validation.jsonl") - # Equal weights: each source contributes its first 50 rows, q0..q49. - assert Counter(r["messages"][0]["content"] for r in rows) == {f"q{i}": 2 for i in range(50)} + # 3:1 weights: the first 75 rows of org/a and 25 of org/b; org/c rounds to zero and is skipped. + expected = Counter({f"q{i}": 1 for i in range(75)}) + Counter(f"q{i}" for i in range(25)) + assert Counter(r["messages"][0]["content"] for r in rows) == expected + assert calls["org/c"] == 0 assert (out / "config.yaml").read_bytes() == config_path.read_bytes() @@ -264,275 +269,116 @@ def test_chat_blend_reuses_and_rebuilds_cached_parts(tmp_path: Path, chat_blend) assert (tmp_path / "blend" / "train.jsonl").read_text(encoding="utf-8") == first -@pytest.mark.parametrize( - "change", - [ - {"shuffle": True}, - { - "sources": [ - { - "hf_dataset": f"org/{n}", - "split": "test", - "content_field": "messages", - "weight": 1, - } - for n in ("a", "b") - ] - }, - ], - ids=["shuffle", "split"], -) -def test_chat_blend_does_not_reuse_parts_of_another_selection(chat_blend, change): - write_config, calls = chat_blend - blend_module.prepare_chat_data_blend(write_config()) - blend_module.prepare_chat_data_blend(write_config(**change)) - assert calls == {"org/a": 2, "org/b": 2} - - @pytest.mark.parametrize( ("first", "second"), [ + ({}, {"sources": [_source("a", split="test")]}), + ( + {"sources": [_source("a", split=None, files=["a.jsonl"])]}, + {"sources": [_source("a", split=None, files=["b.jsonl"])]}, + ), + ({}, {"shuffle": True}), ({"shuffle": True, "seed": 0}, {"shuffle": True, "seed": 1}), ({}, {"reasoning_content": "inline"}), - ( - { - "sources": [ - { - "hf_dataset": "org/a", - "files": ["a.jsonl"], - "content_field": "messages", - "weight": 1, - } - ] - }, - { - "sources": [ - { - "hf_dataset": "org/a", - "files": ["b.jsonl"], - "content_field": "messages", - "weight": 1, - } - ] - }, + ( # a capped part must not stand in for every row + {"sources": [_source("a", split=None, files=["a.jsonl"])]}, + {"sources": [_source("a", split=None, files=["a.jsonl"])], "target_samples": None}, ), ], - ids=["seed", "reasoning_content", "files"], + ids=["split", "files", "shuffle", "seed", "reasoning_content", "no_budget"], ) -def test_chat_blend_reselects_after_a_settings_change(chat_blend, first, second): +def test_chat_blend_reselects_after_a_selection_change(chat_blend, first, second): write_config, calls = chat_blend blend_module.prepare_chat_data_blend(write_config(**first)) - blend_module.prepare_chat_data_blend(write_config(**first)) # cached blend_module.prepare_chat_data_blend(write_config(**second)) assert calls["org/a"] == 2 -def test_chat_blend_fails_on_a_source_without_usable_rows(chat_blend, monkeypatch): - write_config, _ = chat_blend - - def iter_rows(source, output_dir, shuffle, seed): - make_row = _chat if source["hf_dataset"] == "org/a" else _image_chat - yield from (make_row(i) for i in range(100)) - - monkeypatch.setattr(blend_module, "_iter_rows", iter_rows) - # Without target_samples, an all-rejected source must not drop out of the blend silently. - with pytest.raises( - RuntimeError, match=r"org/b train: only 0/any usable rows \(source exhausted" - ): - blend_module.prepare_chat_data_blend(write_config(target_samples=None)) - - -def test_chat_blend_does_not_reuse_a_capped_part_without_a_budget(tmp_path: Path, chat_blend): - write_config, calls = chat_blend - assert blend_module.prepare_chat_data_blend(write_config()) == {"validation": 5, "train": 95} - # Every usable row of both sources, not the 50 per source cached by the capped run. - sizes = blend_module.prepare_chat_data_blend(write_config(target_samples=None)) - assert sizes == {"validation": 10, "train": 190} - assert calls == {"org/a": 2, "org/b": 2} - - -def test_chat_blend_skips_sources_rounded_to_zero_rows(tmp_path: Path, chat_blend): - write_config, calls = chat_blend - config = write_config( - sources=[ - {"hf_dataset": "org/a", "split": "train", "content_field": "messages", "weight": 999}, - {"hf_dataset": "org/b", "split": "train", "content_field": "messages", "weight": 1}, - ] - ) - assert blend_module.prepare_chat_data_blend(config) == {"validation": 5, "train": 95} - assert calls == {"org/a": 1, "org/b": 0} - - @pytest.mark.parametrize( ("overrides", "match"), [ ({"target_samples": 0}, "positive integer"), ({"target_samples": 1.5}, "positive integer"), - ( - { - "sources": [ - { - "hf_dataset": "org/a", - "split": "train", - "content_field": "messages", - "weight": -1, - } - ] - }, - "nonnegative", - ), - ( - { - "sources": [ - { - "hf_dataset": "org/a", - "split": "train", - "content_field": "messages", - "weight": 0, - } - ] - }, - "not all zero", - ), - ( - { - "sources": [ - { - "hf_dataset": "org/a", - "split": "train", - "content_field": "messages", - "weight": float("nan"), - } - ] - }, - "finite", - ), + # Every row of a Hugging Face split would be held in memory. + ({"target_samples": None}, "target_samples is required"), + ({"output_format": "megatron"}, "writes output_format: chat_jsonl"), + ({"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"), + ({"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``. + ({"sources": [_source("a", content_field="conversations")]}, "messages or text"), ], ) -def test_chat_blend_rejects_invalid_budgets(chat_blend, overrides, match): +def test_chat_blend_rejects_invalid_configs(chat_blend, overrides, match): write_config, calls = chat_blend with pytest.raises(ValueError, match=match): blend_module.prepare_chat_data_blend(write_config(**overrides)) - assert calls == {"org/a": 0, "org/b": 0} + assert not calls -def test_chat_blend_writes_parts_atomically(tmp_path: Path, chat_blend, monkeypatch): +@pytest.mark.parametrize( + ("overrides", "match"), + [ + ({}, "Unknown output_format 'chat_jsonl'"), + ({"output_format": "chat-jsonl"}, "Unknown output_format 'chat-jsonl'"), + ({"output_format": "megatron"}, r"Tokenized blends do not use \['seed', 'shuffle'"), + ], +) +def test_tokenized_blend_rejects_chat_options(chat_blend, overrides, match): write_config, _ = chat_blend - real_dumps = blend_module.json.dumps - calls = {"n": 0} - - def flaky_dumps(obj, *args, **kwargs): - calls["n"] += 1 - if calls["n"] == 75: # partway through the second source's part - raise KeyboardInterrupt - return real_dumps(obj, *args, **kwargs) - - monkeypatch.setattr(blend_module.json, "dumps", flaky_dumps) - with pytest.raises(KeyboardInterrupt): - blend_module.prepare_chat_data_blend(write_config()) - parts = tmp_path / "blend" / "parts" - assert not list(parts.glob("org--b_*.jsonl")) # only the interrupted .tmp file was written - (part,) = parts.glob("org--a_*.jsonl") - assert len(_read_jsonl(part)) == 50 + with pytest.raises(ValueError, match=match): + blend_module.prepare_megatron_data_blend(write_config(**overrides)) -def test_chat_blend_fails_on_a_bad_source(tmp_path: Path, chat_blend, monkeypatch): +def test_chat_blend_fails_on_a_bad_short_or_tiny_source(tmp_path: Path, chat_blend, monkeypatch): write_config, _ = chat_blend - real_iter_rows = blend_module._iter_rows + with pytest.raises(RuntimeError, match="org/a train: only 100/250 usable rows"): + blend_module.prepare_chat_data_blend(write_config(target_samples=500)) + with pytest.raises(RuntimeError, match="Too few rows"): + blend_module.prepare_chat_data_blend(write_config(target_samples=10)) def iter_rows(source, output_dir, shuffle, seed): - if source["hf_dataset"] == "org/b": - yield from real_iter_rows(source, output_dir, shuffle, seed) - return yield from (_chat(i) for i in range(10)) # yields some rows, then the reader fails raise TypeError("cannot cast") monkeypatch.setattr(blend_module, "_iter_rows", iter_rows) - with pytest.raises(RuntimeError, match="org/a train: failed to load"): + with pytest.raises(RuntimeError, match="org/a train: failed to read or convert"): blend_module.prepare_chat_data_blend(write_config()) assert not (tmp_path / "blend" / "train.jsonl").exists() -def test_chat_blend_fails_on_a_short_or_tiny_blend(chat_blend): - write_config, _ = chat_blend - with pytest.raises(RuntimeError, match="org/a train: only 100/250 usable rows"): - blend_module.prepare_chat_data_blend(write_config(target_samples=500)) - with pytest.raises(RuntimeError, match="Too few rows"): - blend_module.prepare_chat_data_blend(write_config(target_samples=10)) +def test_write_jsonl_is_atomic(tmp_path: Path): + path = tmp_path / "part.jsonl" + with pytest.raises(TypeError): + blend_module._write_jsonl(path, [_chat(0), {"bad": object()}]) + assert not path.exists() def test_chat_blend_reads_files_sources(tmp_path: Path, monkeypatch): - jsonl_path = tmp_path / "chat.jsonl" - jsonl_path.write_text("".join(json.dumps(_chat(i)) + "\n" for i in range(40)), encoding="utf-8") - monkeypatch.setattr(huggingface_hub, "hf_hub_download", Mock(return_value=str(jsonl_path))) - config_path = tmp_path / "config.yaml" - config_path.write_text( - yaml.safe_dump( - { - "output_dir": str(tmp_path / "blend"), - "output_format": "chat_jsonl", - "sources": [ - { - "hf_dataset": "org/a", - "files": ["data/chat.jsonl"], - "content_field": "messages", - "weight": 1, - } - ], - } - ), - encoding="utf-8", - ) - # No target_samples: every usable row, in file order before mixing. - assert blend_module.prepare_chat_data_blend(config_path) == {"validation": 2, "train": 38} - - -@pytest.mark.parametrize( - ("overrides", "match"), - [ - ({"tokenizer": "x"}, "does not use"), - ({"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``. - ( - { - "sources": [ - { - "hf_dataset": "org/a", - "split": "train", - "content_field": "conversations", - "weight": 1, - } - ] - }, - "content_field messages or text", - ), - ], -) -def test_chat_blend_rejects_tokenizer_options(chat_blend, overrides, match): - write_config, _ = chat_blend - with pytest.raises(ValueError, match=match): - blend_module.prepare_chat_data_blend(write_config(**overrides)) - - -def test_tokenized_blend_rejects_chat_options(chat_blend): - write_config, _ = chat_blend - with pytest.raises(ValueError, match="Tokenized blends do not use"): - blend_module.prepare_megatron_data_blend(write_config()) - - -def test_chat_blend_reads_gzip_files_and_skips_malformed_records(tmp_path: Path, monkeypatch): plain, packed = tmp_path / "a.jsonl", tmp_path / "b.jsonl.gz" lines = [json.dumps(_chat(i)) for i in range(20)] - plain.write_text( - "\n".join([*lines[:10], '{"messages": [', *lines[10:]]) + "\n", encoding="utf-8" - ) + # A malformed record is skipped with a warning, as in the tokenizing reader. + plain.write_text("\n".join([*lines[:10], '{"messages": [', *lines[10:]]) + "\n", "utf-8") with gzip.open(packed, "wt", encoding="utf-8") as f: f.writelines(json.dumps(_chat(i)) + "\n" for i in range(20, 40)) paths = {"data/a.jsonl": str(plain), "data/b.jsonl.gz": str(packed)} monkeypatch.setattr( huggingface_hub, "hf_hub_download", lambda repo_id, filename, **kwargs: paths[filename] ) - source = {"hf_dataset": "org/a", "files": list(paths), "content_field": "messages"} - rows = list(blend_module._iter_rows(source, tmp_path / "blend", False, 0)) - assert [r["messages"][0]["content"] for r in rows] == [f"q{i}" for i in range(40)] + config_path = tmp_path / "config.yaml" + config = { + "output_dir": str(tmp_path / "blend"), + "output_format": "chat_jsonl", + "sources": [_source("a", split=None, files=list(paths))], + } + config_path.write_text(yaml.safe_dump(config), encoding="utf-8") + # No target_samples: every usable row. + assert blend_module.prepare_chat_data_blend(config_path) == {"validation": 2, "train": 38} + out = tmp_path / "blend" + rows = _read_jsonl(out / "train.jsonl") + _read_jsonl(out / "validation.jsonl") + 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 From 157897c0661b671eb1d87ad773a748bc9a48e0a6 Mon Sep 17 00:00:00 2001 From: Yue Date: Wed, 14 Oct 2026 18:45:26 +0000 Subject: [PATCH 5/6] Check source-row media, keep system prompts and require target_samples Run Megatron-Bridge's text-only check on the source row's top-level keys too, keep a top-level system prompt for Bridge to prepend, always require target_samples since the selected rows are mixed in memory, and import load_dataset at module scope like megatron_preprocess_data. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Yue --- examples/dataset/MEGATRON_DATA_PREP.md | 5 +- .../plugins/prepare_megatron_data_blend.py | 65 +++++++++---------- .../test_prepare_megatron_data_blend.py | 23 +++---- 3 files changed, 42 insertions(+), 51 deletions(-) diff --git a/examples/dataset/MEGATRON_DATA_PREP.md b/examples/dataset/MEGATRON_DATA_PREP.md index e22a855779c..c920ed02567 100644 --- a/examples/dataset/MEGATRON_DATA_PREP.md +++ b/examples/dataset/MEGATRON_DATA_PREP.md @@ -185,7 +185,7 @@ when it loads the data. Run it in an environment with Megatron-Bridge, whose tex ```yaml output_dir: /path/to/chat_blend output_format: chat_jsonl -target_samples: 36_800 # Required with split sources; omit only for files sources, to keep every row. +target_samples: 36_800 # Total conversations; required, since the selected rows are mixed in memory. shuffle: false # Optional; take the first rows of each split in order (default: buffer shuffle). seed: 1234 # Optional; seeds the shuffle, the cross-source mix and the validation hold-out. reasoning_content: native # Optional; native (default), inline or strip. @@ -208,7 +208,8 @@ and have an assistant turn; a source that rounds to zero is skipped. `content_fi - A `text` row becomes one assistant turn, or is split into turns if it uses the `User` / `Assistant` / `System` chat markers. A row with a marker for any other role (such as a SteerLM label line) or with an empty turn other than `System` is dropped. -- `messages` and a row's `tools` are prepared as in the tokenizing mode (JSON-string `tool_calls` arguments become dicts), except +- `messages` and a row's `tools` are prepared as in the tokenizing mode (JSON-string `tool_calls` arguments + become dicts), and a top-level `system` prompt is kept for Megatron-Bridge to prepend. The one difference is that `reasoning_content` stays on the message by default (`native`), because the chat template is applied later, at training time. Use `inline` if your model's chat template ignores `reasoning_content`, so the reasoning is kept as a `` block in `content` (see [below](#reasoning_content-for-post-training-v3-datasets)), diff --git a/modelopt/torch/utils/plugins/prepare_megatron_data_blend.py b/modelopt/torch/utils/plugins/prepare_megatron_data_blend.py index 87ff2544d80..4ddc5505774 100644 --- a/modelopt/torch/utils/plugins/prepare_megatron_data_blend.py +++ b/modelopt/torch/utils/plugins/prepare_megatron_data_blend.py @@ -30,6 +30,7 @@ import huggingface_hub import yaml +from datasets import load_dataset from modelopt.torch.utils import ( REASONING_CONTENT_MODES, @@ -182,7 +183,7 @@ def _to_chat( """Convert a source row to ``{"messages": [...]}``, or ``None`` if it has nothing to train on. ``messages`` go through :func:`prepare_messages_for_chat_template` with ``reasoning_content``, - and a non-empty ``tools`` list is kept for the chat template. A ``text`` row in the NeMo chat + and non-empty ``tools`` and ``system`` fields are kept for the chat path. A ``text`` row in the NeMo chat format (``User`` / ``Assistant`` markers) is split into turns; any other ``text`` row becomes one assistant turn, so all of it is trained. """ @@ -192,9 +193,11 @@ def _to_chat( return None messages = prepare_messages_for_chat_template(messages, reasoning_content=reasoning_content) chat: dict[str, Any] = {"messages": messages} - # As in the tokenizing mode: the chat template renders the tool definitions into the prompt. - if row.get("tools"): - chat["tools"] = row["tools"] + # As in the tokenizing mode, the chat template renders the tool definitions into the prompt; + # Megatron-Bridge turns a top-level system prompt into the leading system message. + for key in ("tools", "system"): + if row.get(key): + chat[key] = row[key] else: text = row.get("text") or "" if not text: @@ -241,8 +244,6 @@ def _iter_rows( # As in the tokenizing reader: some Nemotron files carry a malformed record. print(f"[WARN] Skipping malformed JSONL record in {path}: {e}", flush=True) return - from datasets import load_dataset - dataset = load_dataset( source["hf_dataset"], source.get("config"), split=source["split"], streaming=True ) @@ -252,15 +253,15 @@ def _iter_rows( def _take_chat_rows( source: dict[str, Any], output_dir: Path, - want: int | None, + want: int, shuffle: bool, seed: int, reasoning_content: ReasoningContentMode, ) -> list[dict[str, Any]]: - """Return the first ``want`` usable chat rows of a source (all of them if ``want`` is None). + """Return the first ``want`` usable chat rows of a source. - Fails if the source cannot be read, has no usable row, or runs out of them, so the blend cannot - change silently; gives up after ``10 * want`` rejected rows. + Fails if the source cannot be read or runs out of usable rows, so the blend cannot change + silently; gives up after ``10 * want`` rejected rows. """ # Megatron-Bridge's chat SFT path consumes this output; keep exactly the rows it accepts. Imported # here, not at module scope, because the tokenizing mode does not need Megatron-Bridge installed. @@ -272,21 +273,21 @@ def _take_chat_rows( try: for row in _iter_rows(source, output_dir, shuffle, seed): chat = _to_chat(row, source["content_field"], reasoning_content) - if chat is None or not is_text_only_chat_example(chat): + # 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 - if want is not None and rejected > 10 * want: + if rejected > 10 * want: break continue rows.append(chat) - if want is not None and len(rows) >= want: + if len(rows) >= want: break except Exception as e: # e.g. a gated split, or a schema the streaming reader cannot cast raise RuntimeError(f"{name}: failed to read or convert") from e - if not rows or (want is not None and len(rows) < want): - gave_up = want is not None and rejected > 10 * want + if len(rows) < want: + gave_up = rejected > 10 * want reason = f"gave up after {rejected} rejected rows" if gave_up else "source exhausted" - budget = "any" if want is None else want - raise RuntimeError(f"{name}: only {len(rows)}/{budget} usable rows ({reason})") + raise RuntimeError(f"{name}: only {len(rows)}/{want} usable rows ({reason})") print(f"{name}: {len(rows)} rows ({rejected} rejected)", flush=True) return rows @@ -308,34 +309,29 @@ def _read_jsonl(path: Path) -> list[dict[str, Any]] | None: return None -def _part_path( - parts_dir: Path, source: dict[str, Any], settings: dict[str, Any], capped: bool -) -> Path: +def _part_path(parts_dir: Path, source: dict[str, Any], settings: dict[str, Any]) -> Path: """Return the cache file of a source's selected rows, keyed by every input that selects them. - A capped part holds the first rows of the source and serves any smaller budget; an uncapped part - holds every usable row, so a budget removed between runs never reuses a capped selection. + A part holds the first usable rows of the source, so it serves any budget it covers. """ selection = {key: source.get(key) for key in ("hf_dataset", "config", "split", "files")} - selection.update(settings, content_field=source["content_field"], capped=capped) + selection.update(settings, content_field=source["content_field"]) digest = hashlib.sha256(json.dumps(selection, sort_keys=True).encode()).hexdigest()[:16] label = source.get("split") or Path(source["files"][0]).name.split(".")[0] return parts_dir / f"{source['hf_dataset'].replace('/', '--')}_{label}_{digest}.jsonl" -def _chat_budgets(config: dict[str, Any]) -> list[int | None]: - """Return each source's row budget, ``None`` meaning every usable row; validate the inputs.""" +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"]] if not all(math.isfinite(w) and w >= 0 for w in weights) or sum(weights) <= 0: raise ValueError(f"Source weights must be finite, nonnegative and not all zero: {weights}.") + # Required: the selected rows are held in memory to be mixed, which a whole source may not fit. target = config.get("target_samples") - if target is None: - if any("split" in source for source in config["sources"]): - # Every row of a split would be held in memory, which a large split does not fit. - raise ValueError("target_samples is required when a source is a Hugging Face split.") - return [None] * len(weights) if isinstance(target, bool) or not isinstance(target, int) or target <= 0: - raise ValueError(f"target_samples must be a positive integer, got {target!r}.") + raise ValueError( + f"target_samples is required and must be a positive integer, got {target!r}." + ) return [round(target * w / sum(weights)) for w in weights] @@ -345,8 +341,7 @@ def prepare_chat_data_blend(config_path: Path) -> dict[str, int]: For Megatron-Bridge's chat SFT path (``distill.py --sft --sft_hf_dataset``), which renders each conversation with the model's chat template and can restrict the loss to assistant tokens. Each source contributes ``round(target_samples * weight / sum(weights))`` rows (a source rounded - to zero rows is skipped) that the chat path accepts as text-only chat; ``target_samples`` may be - omitted only for ``files`` sources, to keep every usable row. ``content_field`` must be + to zero rows is skipped) that the chat path accepts as text-only chat. ``content_field`` must be ``messages`` or ``text``; ``messages`` keep ``reasoning_content`` for the chat template by default (``reasoning_content: native``; set ``inline`` for a template that ignores the field, or ``strip``). Hugging Face splits are buffer-shuffled with ``seed`` unless ``shuffle: false``; @@ -393,9 +388,9 @@ def prepare_chat_data_blend(config_path: Path) -> dict[str, int]: if want == 0: print(f"{source['hf_dataset']}: weight rounds to 0 rows, skipped", flush=True) continue - part = _part_path(parts_dir, source, settings, capped=want is not None) + part = _part_path(parts_dir, source, settings) got = _read_jsonl(part) - if not got or (want is not None and len(got) < want): + if got is None or len(got) < want: got = _take_chat_rows(source, output_dir, want, shuffle, seed, reasoning_content) _write_jsonl(part, got) rows.extend(got[:want]) 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 733a57905c0..a8c38d3a7c7 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 @@ -178,10 +178,11 @@ def test_to_chat_handles_all_row_shapes(): assert blend_module._to_chat({"messages": []}, "messages") is None # Nothing would be trained under assistant-only loss. assert blend_module._to_chat({"messages": msgs[:1]}, "messages") is None - # Tool definitions reach the chat template, as in the tokenizing mode. + # Tool definitions and a top-level system prompt reach the chat path. tools = [{"type": "function", "function": {"name": "add"}}] - chat = blend_module._to_chat({"messages": msgs, "tools": tools}, "messages") - assert chat == {"messages": msgs, "tools": tools} + row = {"messages": msgs, "tools": tools, "system": "Be brief.", "id": 7} + chat = blend_module._to_chat(row, "messages") + assert chat == {"messages": msgs, "tools": tools, "system": "Be brief."} # Unlike the tokenizing mode, reasoning_content stays on the message by default (native). reasoning = [{"role": "assistant", "content": "4", "reasoning_content": "2+2"}] assert blend_module._to_chat({"messages": reasoning}, "messages")["messages"] == reasoning @@ -222,7 +223,8 @@ def test_to_chat_handles_all_row_shapes(): def test_take_chat_rows_skips_non_text_rows_and_fails_without_enough(tmp_path: Path, monkeypatch): source = _source("a") - rows = [_chat(0), _image_chat(1), _chat(2), _chat(3)] + # 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") assert [c["messages"][0]["content"] for c in got] == ["q0", "q2", "q3"] @@ -231,9 +233,6 @@ def test_take_chat_rows_skips_non_text_rows_and_fails_without_enough(tmp_path: P monkeypatch.setattr(blend_module, "_iter_rows", lambda *a: iter(images)) with pytest.raises(RuntimeError, match=r"only 0/2 usable rows \(gave up after 21 rejected"): blend_module._take_chat_rows(source, tmp_path, 2, False, 0, "native") - # Without a budget, an all-rejected source must not drop out of the blend silently. - with pytest.raises(RuntimeError, match=r"only 0/any usable rows \(source exhausted"): - blend_module._take_chat_rows(source, tmp_path, None, False, 0, "native") def test_chat_blend_splits_weighted_rows(tmp_path: Path, chat_blend, monkeypatch): @@ -280,12 +279,8 @@ 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"}), - ( # a capped part must not stand in for every row - {"sources": [_source("a", split=None, files=["a.jsonl"])]}, - {"sources": [_source("a", split=None, files=["a.jsonl"])], "target_samples": None}, - ), ], - ids=["split", "files", "shuffle", "seed", "reasoning_content", "no_budget"], + ids=["split", "files", "shuffle", "seed", "reasoning_content"], ) def test_chat_blend_reselects_after_a_selection_change(chat_blend, first, second): write_config, calls = chat_blend @@ -299,7 +294,7 @@ def test_chat_blend_reselects_after_a_selection_change(chat_blend, first, second [ ({"target_samples": 0}, "positive integer"), ({"target_samples": 1.5}, "positive integer"), - # Every row of a Hugging Face split would be held in memory. + # A whole source would be held in memory. ({"target_samples": None}, "target_samples is required"), ({"output_format": "megatron"}, "writes output_format: chat_jsonl"), ({"sources": [_source("a", weight=-1)]}, "nonnegative"), @@ -372,10 +367,10 @@ def test_chat_blend_reads_files_sources(tmp_path: Path, monkeypatch): config = { "output_dir": str(tmp_path / "blend"), "output_format": "chat_jsonl", + "target_samples": 40, "sources": [_source("a", split=None, files=list(paths))], } config_path.write_text(yaml.safe_dump(config), encoding="utf-8") - # No target_samples: every usable row. assert blend_module.prepare_chat_data_blend(config_path) == {"validation": 2, "train": 38} out = tmp_path / "blend" rows = _read_jsonl(out / "train.jsonl") + _read_jsonl(out / "validation.jsonl") From 312129bef99f3f75890af529384ab11800bfca40 Mon Sep 17 00:00:00 2001 From: Yue Date: Wed, 14 Oct 2026 21:20:52 +0000 Subject: [PATCH 6/6] Drop text rows with leading text and floor the give-up threshold A text row with content before its first chat marker would lose it, so the row is dropped. Give up after 10 * want + 100 rejected rows, so a small budget does not misjudge a source by its first few rows. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Yue --- .../plugins/prepare_megatron_data_blend.py | 23 +++++++++++-------- .../test_prepare_megatron_data_blend.py | 6 +++-- 2 files changed, 18 insertions(+), 11 deletions(-) diff --git a/modelopt/torch/utils/plugins/prepare_megatron_data_blend.py b/modelopt/torch/utils/plugins/prepare_megatron_data_blend.py index 4ddc5505774..2bc7aacb8c9 100644 --- a/modelopt/torch/utils/plugins/prepare_megatron_data_blend.py +++ b/modelopt/torch/utils/plugins/prepare_megatron_data_blend.py @@ -206,13 +206,16 @@ def _to_chat( return {"messages": [{"role": "assistant", "content": text}]} matches = _CHAT_TURN.findall(text) # Markers that parse to no turns make the row malformed, not plain text: falling through - # would train the raw delimiters as assistant content. A marker without a known - # role (a Tool turn, or a SteerLM label line) or a known turn left empty would silently cut - # it out of the conversation, so drop the row. Only a bare delimiter at the very end needs no - # role, and only the System block, which real rows leave empty, may be empty. + # would train the raw delimiters as assistant content. Text before the first + # marker, a marker without a known role (a Tool turn, or a SteerLM label line) or a known turn + # left empty would silently cut it out of the conversation, so drop the row. Only a bare + # delimiter at the very end needs no role, and only the System block, which real rows leave + # empty, may be empty. roleless = len(_CHAT_MARKER.findall(text)) - len(matches) - if roleless > (1 if _CHAT_TRAILING_MARKER.search(text) else 0) or any( - not content.strip() and role != "System" for role, content in matches + if ( + not _CHAT_MARKER.match(text.lstrip()) + or roleless > (1 if _CHAT_TRAILING_MARKER.search(text) else 0) + or any(not content.strip() and role != "System" for role, content in matches) ): return None turns = [ @@ -261,7 +264,8 @@ def _take_chat_rows( """Return the first ``want`` usable chat rows of a source. Fails if the source cannot be read or runs out of usable rows, so the blend cannot change - silently; gives up after ``10 * want`` rejected rows. + silently; gives up after ``10 * want + 100`` rejected rows, a floor that keeps a small budget from + misjudging a source by its first few rows. """ # Megatron-Bridge's chat SFT path consumes this output; keep exactly the rows it accepts. Imported # here, not at module scope, because the tokenizing mode does not need Megatron-Bridge installed. @@ -270,13 +274,14 @@ def _take_chat_rows( name = f"{source['hf_dataset']} {source.get('split') or source.get('files')}" rows: list[dict[str, Any]] = [] rejected = 0 + 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) # 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 - if rejected > 10 * want: + if rejected > max_rejected: break continue rows.append(chat) @@ -285,7 +290,7 @@ def _take_chat_rows( except Exception as e: # e.g. a gated split, or a schema the streaming reader cannot cast raise RuntimeError(f"{name}: failed to read or convert") from e if len(rows) < want: - gave_up = rejected > 10 * want + gave_up = rejected > max_rejected reason = f"gave up after {rejected} rejected rows" if gave_up else "source exhausted" raise RuntimeError(f"{name}: only {len(rows)}/{want} usable rows ({reason})") print(f"{name}: {len(rows)} rows ({rejected} rejected)", flush=True) 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 a8c38d3a7c7..6306ecef292 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 @@ -209,6 +209,8 @@ def test_to_chat_handles_all_row_shapes(): text = "User\nq\nTool\nresult\nAssistant\na\n" assert blend_module._to_chat({"text": text}, "text") is None for text in ( + # Text before the first marker would be cut from the conversation. + "preamble\nUser\nq\nAssistant\na\n", # A SteerLM label line is no role either. "User\nq\nAssistant\nquality:4\na\n", # A role must end at a newline: this is no User turn. @@ -229,9 +231,9 @@ def test_take_chat_rows_skips_non_text_rows_and_fails_without_enough(tmp_path: P got = blend_module._take_chat_rows(source, tmp_path, 3, False, 0, "native") assert [c["messages"][0]["content"] for c in got] == ["q0", "q2", "q3"] - images = [_image_chat(i) for i in range(50)] + 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 21 rejected"): + 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")