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..c920ed02567 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,54 @@ 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 # 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. +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 +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 (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), 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)), + 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. +Nothing prunes `parts`; delete it once `train.jsonl` is written to reclaim the space. + ## `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..2bc7aacb8c9 100644 --- a/modelopt/torch/utils/plugins/prepare_megatron_data_blend.py +++ b/modelopt/torch/utils/plugins/prepare_megatron_data_blend.py @@ -16,17 +16,42 @@ """Prepare a weighted Megatron data blend from a YAML configuration.""" import argparse +import gzip +import hashlib +import json +import math 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 modelopt.torch.utils import ( + REASONING_CONTENT_MODES, + ReasoningContentMode, + prepare_messages_for_chat_template, +) 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 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))(.*?)(?=|\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]: @@ -53,6 +78,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``, ``seed`` and ``reasoning_content`` keys apply. """ with path.open(encoding="utf-8") as stream: return cast("dict[str, Any]", yaml.safe_load(stream)) @@ -76,6 +105,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 +145,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 +177,252 @@ def _prepare_sources( return blend +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 nothing to train on. + + ``messages`` go through :func:`prepare_messages_for_chat_template` with ``reasoning_content``, + 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. + """ + if content_field == "messages": + messages = row.get("messages") + if not messages: + 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; + # 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: + return None + if " 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 ( + 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 = [ + {"role": _CHAT_ROLES[role], "content": content.strip()} + for role, content in matches + if content.strip() + ] + chat = {"messages": turns} + # Nothing in a row without an assistant turn is trained under assistant-only loss. + if not any(turn.get("role") == "assistant" for turn in chat["messages"]): + return None + return chat + + +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): + 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 + dataset = load_dataset( + source["hf_dataset"], source.get("config"), split=source["split"], streaming=True + ) + yield from dataset.shuffle(seed=seed) if shuffle else dataset + + +def _take_chat_rows( + source: dict[str, Any], + output_dir: Path, + want: int, + shuffle: bool, + seed: int, + reasoning_content: ReasoningContentMode, +) -> list[dict[str, Any]]: + """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 + 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. + 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 + 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 > max_rejected: + break + continue + rows.append(chat) + 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 len(rows) < 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) + 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 _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 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"]) + 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]: + """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 isinstance(target, bool) or not isinstance(target, int) or target <= 0: + 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] + + +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 (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. + """ + 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)}.") + 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}." + ) + budgets = _chat_budgets(config) + 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"], budgets): + if want == 0: + print(f"{source['hf_dataset']}: weight rounds to 0 rows, skipped", flush=True) + continue + part = _part_path(parts_dir, source, settings) + got = _read_jsonl(part) + if got is None or len(got) < want: + got = _take_chat_rows(source, output_dir, want, shuffle, seed, reasoning_content) + _write_jsonl(part, got) + rows.extend(got[:want]) + + 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"]: + 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) + output_format = config.get("output_format", "megatron") + if output_format != "megatron": + raise ValueError( + 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") @@ -165,6 +439,11 @@ def main() -> None: """Prepare a data blend from the supplied configuration.""" parser = _build_parser() args = parser.parse_args() + # 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 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..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 @@ -13,8 +13,11 @@ # See the License for the specific language governing permissions and # limitations under the License. +import gzip +import importlib import json import sys +from collections import Counter from pathlib import Path from unittest.mock import Mock @@ -24,6 +27,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 +109,273 @@ 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 _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): + """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 + 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": [_source("a"), _source("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"}, {"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 and a top-level system prompt reach the chat path. + tools = [{"type": "function", "function": {"name": "add"}}] + row = {"messages": msgs, "tools": tools, "system": "Be brief.", "id": 7} + 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 + inline = blend_module._to_chat({"messages": reasoning}, "messages", "inline") + assert inline["messages"] == [{"role": "assistant", "content": "\n2+2\n\n4"}] + + 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 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"}] + } + 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 + 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. + "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 + + +def test_take_chat_rows_skips_non_text_rows_and_fails_without_enough(tmp_path: Path, monkeypatch): + source = _source("a") + # 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"] + + images = [_image_chat(i) for i in range(200)] + monkeypatch.setattr(blend_module, "_iter_rows", lambda *a: iter(images)) + with pytest.raises(RuntimeError, match=r"only 0/2 usable rows \(gave up after 121 rejected"): + blend_module._take_chat_rows(source, tmp_path, 2, False, 0, "native") + + +def test_chat_blend_splits_weighted_rows(tmp_path: Path, chat_blend, monkeypatch): + 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)] + ) + main() + out = tmp_path / "blend" + assert len(_read_jsonl(out / "validation.jsonl")) == 5 + rows = _read_jsonl(out / "train.jsonl") + _read_jsonl(out / "validation.jsonl") + # 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() + + +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").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( + ("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"}), + ], + 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 + blend_module.prepare_chat_data_blend(write_config(**first)) + blend_module.prepare_chat_data_blend(write_config(**second)) + assert calls["org/a"] == 2 + + +@pytest.mark.parametrize( + ("overrides", "match"), + [ + ({"target_samples": 0}, "positive integer"), + ({"target_samples": 1.5}, "positive integer"), + # 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"), + ({"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_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 not calls + + +@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 + with pytest.raises(ValueError, match=match): + blend_module.prepare_megatron_data_blend(write_config(**overrides)) + + +def test_chat_blend_fails_on_a_bad_short_or_tiny_source(tmp_path: Path, chat_blend, monkeypatch): + 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 iter_rows(source, output_dir, shuffle, seed): + 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 read or convert"): + blend_module.prepare_chat_data_blend(write_config()) + assert not (tmp_path / "blend" / "train.jsonl").exists() + + +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): + plain, packed = tmp_path / "a.jsonl", tmp_path / "b.jsonl.gz" + lines = [json.dumps(_chat(i)) for i in range(20)] + # 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] + ) + config_path = tmp_path / "config.yaml" + 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") + 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