From 05046acf769e8928e8392208ff88f8f65d8bbcd0 Mon Sep 17 00:00:00 2001
From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Date: Wed, 21 Oct 2026 02:18:06 -0700
Subject: [PATCH 1/5] Prepare packed chat-SFT data and add per-document
 attention to Megatron-Bridge distillation

- chat_jsonl blends: tokenizer/max_length/loss_mode drop rows with no trained
  token; pack_max_tokens pairs rows for --sft_in_batch_packing --mbs 2.
- distill.py --document_boundaries: per-document THD segments for
  pre-tokenized data.
- examples/megatron_bridge/requirements.txt: flash-linear-attention>=0.5.2,
  since earlier versions compute packed Gated DeltaNet wrongly (nemo:26.08
  ships 0.4.2).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
---
 CHANGELOG.rst                                 |   1 +
 examples/dataset/MEGATRON_DATA_PREP.md        |  17 +++
 examples/megatron_bridge/README.md            |  17 +++
 examples/megatron_bridge/distill.py           |  33 +++++-
 examples/megatron_bridge/requirements.txt     |   4 +
 .../plugins/prepare_megatron_data_blend.py    | 103 +++++++++++++++++-
 .../examples/megatron_bridge/test_distill.py  |  34 ++++++
 .../test_prepare_megatron_data_blend.py       |  46 +++++++-
 8 files changed, 249 insertions(+), 6 deletions(-)
 create mode 100644 examples/megatron_bridge/requirements.txt

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index 9ee5a31b096..88055629488 100755
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -19,6 +19,7 @@ Changelog
 *Megatron Framework (M-LM / M-Bridge)*
 
 - Add ``--sft_in_batch_packing``, ``--moe_aux_loss_coeff`` and ``--dsa_kernel_backend`` to ``examples/megatron_bridge/distill.py`` (``--dsa_kernel_backend`` also to ``quantize.py``): pack ``--sft_hf_dataset`` conversations per micro-batch instead of padding each to ``--seq_length``, override the MoE load-balancing loss weight (e.g. 0 for padded SFT data), and select fused DeepSeek sparse-attention kernels for long sequences.
+- Add ``--document_boundaries`` to ``examples/megatron_bridge/distill.py`` to attend within each document of a pre-tokenized window, and ``tokenizer`` / ``max_length`` / ``pack_max_tokens`` to ``chat_jsonl`` blends to drop conversations with nothing to train and pair the rest for ``--sft_in_batch_packing``. Packed distillation of Gated DeltaNet models needs ``flash-linear-attention`` >= 0.5.2 (earlier versions compute it wrongly); install ``examples/megatron_bridge/requirements.txt`` in ``nemo:26.08``.
 - Add ``--vllm_fakequant_export`` to ``examples/megatron_bridge/export_quantized_megatron_to_hf.py`` to evaluate formats vLLM has no kernels for (e.g. GGML IQ) as BF16 fake-quant checkpoints, and allow ``--ep_size`` > 1 for checkpoints built with ``--no_moe_grouped_gemm``.
 - GGML GPTQ checkpoints from ``examples/megatron_bridge`` keep their calibrated codes when the Megatron checkpoint is reloaded for QAD or export, and layerwise calibration now runs on Megatron vision-language models.
 - The Megatron unified HF export (``export_mcore_gpt_to_hf``) now copies the source checkpoint's non-model files (tokenizer, processor, remote code, chat templates, ...) into the export unchanged, from local disk only. The ``examples/megatron_bridge`` scripts make a Hugging Face Hub model local first with ``modelopt.torch.export.ensure_local_checkpoint``.
diff --git a/examples/dataset/MEGATRON_DATA_PREP.md b/examples/dataset/MEGATRON_DATA_PREP.md
index c920ed02567..357203a4264 100644
--- a/examples/dataset/MEGATRON_DATA_PREP.md
+++ b/examples/dataset/MEGATRON_DATA_PREP.md
@@ -221,6 +221,23 @@ rows, or runs out of them stops the run. Each source's selection is cached under
 the settings that choose its rows, so an interrupted run resumes and a changed configuration selects afresh.
 Nothing prunes `parts`; delete it once `train.jsonl` is written to reclaim the space.
 
+To feed `distill.py --sft_in_batch_packing`, also count tokens the way Megatron-Bridge trains them:
+
+```yaml
+tokenizer: Qwen/Qwen3.6-35B-A3B  # The student's tokenizer.
+max_length: 32768       # distill.py --seq_length; longer conversations are truncated.
+loss_mode: assistant    # Optional; as distill.py --sft_loss_mode.
+pack_max_tokens: 33280  # Optional; order each split as pairs of at most this many tokens.
+num_workers: 48         # Optional; tokenizer processes (default: all CPUs).
+```
+
+With `tokenizer` and `max_length`, a conversation with no trained token left after truncation is
+dropped, since Megatron-Bridge fails on it during training. `pack_max_tokens` pairs long
+conversations with short ones and writes each split pair by pair, so `--mbs 2` packs every pair
+into one micro-batch; a conversation that fits with no other is dropped. Tokenizing runs on the
+CPU: ~0.4 s per conversation per process for GLM-5.3-Flash on Nemotron-Cascade-2 (4.4K tokens
+on average).
+
 ## `reasoning_content` for Post-Training v3 Datasets
 
 v3 datasets include a `reasoning_content` field in assistant messages (chain-of-thought separate from
diff --git a/examples/megatron_bridge/README.md b/examples/megatron_bridge/README.md
index db46db9e43c..382df67004e 100644
--- a/examples/megatron_bridge/README.md
+++ b/examples/megatron_bridge/README.md
@@ -184,6 +184,23 @@ the last one (`last_turn`), or the whole conversation (`full`). With `--eval_ite
 `--sft_hf_validation` in the same format. This path needs a Megatron-Bridge with
 `DirectHFSFTDatasetConfig`.
 
+Megatron-Bridge pads every conversation to `--seq_length` when `--ep_size` > 1, so for MoE models
+add `--sft_in_batch_packing --mbs 2 --moe_aux_loss_coeff 0`: each micro-batch packs two
+conversations into one unpadded THD sequence, and the MoE load-balancing loss, which would also
+run over padding, is dropped. Order the rows as pairs that fit together with the chat-blend
+`pack_max_tokens` option ([Dataset Preparation README](../dataset/MEGATRON_DATA_PREP.md#chat-jsonl-blends-for-assistant-only-loss)).
+Megatron-Bridge holds every row in host memory on each rank, so keep `train.jsonl` to the rows
+the run needs.
+
+With pre-tokenized data, `--document_boundaries` (`--mbs 1`, `--cp_size 1`) restricts attention to
+each EOD-separated document of a window instead of the whole window.
+
+> [!IMPORTANT]
+> Packed training (`--sft_in_batch_packing` or `--document_boundaries`) of Gated DeltaNet models
+> (e.g. Qwen3.5 / Qwen3.6) needs `flash-linear-attention` >= 0.5.2: earlier versions compute the
+> packed linear attention wrongly. `nemo:26.08` ships 0.4.2, so run
+> `python -m pip install -r requirements.txt` in this directory first.
+
 Teacher and student must share a tokenizer — distillation scores the teacher on the student's
 token ids, and the KD losses compare the two models' logits elementwise over the vocab dimension.
 
diff --git a/examples/megatron_bridge/distill.py b/examples/megatron_bridge/distill.py
index 8fbb3ab2fb1..c9d2c082845 100644
--- a/examples/megatron_bridge/distill.py
+++ b/examples/megatron_bridge/distill.py
@@ -57,6 +57,7 @@
     HAS_DIRECT_HF_SFT = True
 except ImportError:
     HAS_DIRECT_HF_SFT = False
+from megatron.core.datasets.gpt_dataset import GPTDataset
 from megatron.core.datasets.utils import get_blend_from_list
 from megatron.core.distributed import DistributedDataParallelConfig
 from megatron.core.utils import unwrap_model
@@ -201,6 +202,12 @@ def get_args():
         nargs="+",
         help="List of tokenized data paths to load from (weight1 path1 weight2 path2 ...)",
     )
+    parser.add_argument(
+        "--document_boundaries",
+        action="store_true",
+        help="Pre-tokenized data only: attend within each EOD-separated document of a sequence (THD "
+        "segments) instead of across the whole packed window. Needs --mbs 1 and --cp_size 1.",
+    )
     parser.add_argument(
         "--data_path_to_cache", type=str, default=None, help="Path to cache the dataset indices"
     )
@@ -449,6 +456,9 @@ def get_args():
         raise ValueError(
             "--sft_hf_validation / --sft_loss_mode / --sft_in_batch_packing require --sft_hf_dataset."
         )
+    if args.document_boundaries and (args.sft or args.mbs != 1 or args.cp_size != 1):
+        # Each window carries its own number of documents, so windows cannot be stacked.
+        raise ValueError("--document_boundaries needs pre-tokenized data, --mbs 1 and --cp_size 1.")
     if args.sft and (args.data_paths or args.use_mock_data):
         raise ValueError(
             "--sft is mutually exclusive with --data_paths / --use_mock_data: the SFT branch wins "
@@ -497,6 +507,25 @@ def _tokenizer_prepends_bos(args) -> bool:
     return tokenizer("x").input_ids[:1] == [tokenizer.bos_token_id]
 
 
+def _add_document_boundaries() -> None:
+    """Give each EOD-separated document of a pre-tokenized window its own THD segment."""
+    getitem = GPTDataset.__getitem__
+
+    def getitem_with_boundaries(self, idx):
+        sample = getitem(self, idx)
+        tokens = sample["tokens"]
+        n = tokens.numel()
+        ends = (tokens == self.config.tokenizer.eod).nonzero(as_tuple=True)[0] + 1
+        # On the tokens' (CPU) device: CUDA tensors from a dataloader worker are shared over IPC.
+        bounds = [tokens.new_zeros(1), ends[ends < n], tokens.new_tensor([n])]
+        cu_seqlens = torch.cat(bounds).to(torch.int32)
+        sample["cu_seqlens_q"] = cu_seqlens
+        sample["max_seqlen_q"] = (cu_seqlens[1:] - cu_seqlens[:-1]).max()
+        return sample
+
+    GPTDataset.__getitem__ = getitem_with_boundaries
+
+
 def main(args: argparse.Namespace, owns_the_run: bool = True):
     # --student_hf_path stays as given. The student is built and exported from the local copy of
     # the whole checkpoint in student_local_checkpoint_path, since the HF export reads the source's
@@ -638,6 +667,8 @@ def _restore_student_hook(model_chunks):
             return model_chunks
 
         distill_provider.register_pre_wrap_hook(_restore_student_hook, prepend=True)
+    if args.document_boundaries:
+        _add_document_boundaries()
 
     # Build optimizer and scheduler
     optimizer_config, scheduler_config = distributed_fused_adam_with_cosine_annealing(
@@ -653,7 +684,7 @@ def _restore_student_hook(model_chunks):
         "path_to_cache": args.data_path_to_cache,
         "random_seed": args.seed,
         "reset_attention_mask": False,
-        "reset_position_ids": False,
+        "reset_position_ids": args.document_boundaries,
         "eod_mask_loss": False,
         "num_dataset_builder_threads": 1,
         "data_sharding": True,
diff --git a/examples/megatron_bridge/requirements.txt b/examples/megatron_bridge/requirements.txt
new file mode 100644
index 00000000000..e8e6ce84c91
--- /dev/null
+++ b/examples/megatron_bridge/requirements.txt
@@ -0,0 +1,4 @@
+# nemo:26.08 ships flash-linear-attention 0.4.2, whose Gated DeltaNet kernel computes packed (THD) sequences
+# wrongly (--sft_in_batch_packing, --document_boundaries); 0.5.2 is the first correct release.
+# TODO: remove once the examples move to nemo:26.10, which ships 0.5.2.
+flash-linear-attention>=0.5.2
diff --git a/modelopt/torch/utils/plugins/prepare_megatron_data_blend.py b/modelopt/torch/utils/plugins/prepare_megatron_data_blend.py
index 2bc7aacb8c9..0a0d0ca7465 100644
--- a/modelopt/torch/utils/plugins/prepare_megatron_data_blend.py
+++ b/modelopt/torch/utils/plugins/prepare_megatron_data_blend.py
@@ -20,6 +20,7 @@
 import hashlib
 import json
 import math
+import multiprocessing as mp
 import os
 import random
 import re
@@ -326,6 +327,65 @@ def _part_path(parts_dir: Path, source: dict[str, Any], settings: dict[str, Any]
     return parts_dir / f"{source['hf_dataset'].replace('/', '--')}_{label}_{digest}.jsonl"
 
 
+_SFT_TOKENIZER: dict[str, Any] = {}
+
+
+def _init_sft_tokenizer(tokenizer: str, max_length: int, loss_mode: str) -> None:
+    from megatron.bridge.data.collators.sft import extract_skipped_token_ids
+    from megatron.bridge.data.conversation_processing import infer_assistant_mask_boundary_config
+    from transformers import AutoTokenizer
+
+    # distill.py --sft_hf_dataset tokenizes with the student's HF tokenizer, as loaded here.
+    processor = AutoTokenizer.from_pretrained(tokenizer)
+    _SFT_TOKENIZER.update(
+        processor=processor,
+        boundary_config=infer_assistant_mask_boundary_config(processor),
+        skipped_tokens=extract_skipped_token_ids(processor),
+        max_length=max_length,
+        loss_mode=loss_mode,
+    )
+
+
+def _sft_length(row: dict[str, Any]) -> int:
+    """Return the row's token count as Megatron-Bridge trains it, or 0 if nothing would be trained."""
+    from megatron.bridge.data.conversation_processing import tokenize_chat_example
+
+    try:
+        tokenized = tokenize_chat_example(row, warn_on_all_masked=False, **_SFT_TOKENIZER)
+    except ValueError:  # e.g. no assistant span matched: Bridge raises from a dataloader worker
+        return 0
+    return tokenized.input_ids.numel() if tokenized.assistant_mask.any() else 0
+
+
+def _sft_lengths(rows: list[dict[str, Any]], config: dict[str, Any]) -> list[int]:
+    """Tokenize ``rows`` with Megatron-Bridge's chat SFT path in ``num_workers`` processes."""
+    init_args = (config["tokenizer"], config["max_length"], config.get("loss_mode", "assistant"))
+    num_workers = int(config.get("num_workers", len(os.sched_getaffinity(0))))
+    if num_workers <= 1:
+        _init_sft_tokenizer(*init_args)
+        return [_sft_length(row) for row in rows]
+    with mp.Pool(num_workers, initializer=_init_sft_tokenizer, initargs=init_args) as pool:
+        return pool.map(_sft_length, rows, chunksize=64)
+
+
+def _pair_rows(rows: list[dict[str, Any]], lengths: list[int], cap: int, seed: int) -> list:
+    """Order rows as pairs of at most ``cap`` tokens, so ``--mbs 2`` packs each pair together.
+
+    Pairs the longest remaining row with the shortest; a row that fits with no other is dropped.
+    The pairs are shuffled with ``seed``.
+    """
+    order = sorted(range(len(rows)), key=lengths.__getitem__)
+    pairs, i, j = [], 0, len(order) - 1
+    while i < j:
+        if lengths[order[i]] + lengths[order[j]] <= cap:
+            pairs.append((order[i], order[j]))
+            i += 1
+        j -= 1
+    random.Random(seed).shuffle(pairs)
+    print(f"Paired {2 * len(pairs)} of {len(rows)} rows into packs of <= {cap} tokens", flush=True)
+    return [rows[k] for pair in pairs for k in pair]
+
+
 def _chat_budgets(config: dict[str, Any]) -> list[int]:
     """Return each source's row budget; validate the inputs."""
     weights = [float(source["weight"]) for source in config["sources"]]
@@ -354,6 +414,12 @@ def prepare_chat_data_blend(config_path: Path) -> dict[str, int]:
     The rows are mixed with ``seed`` and ``min(1000, n // 20)`` of them are held out for validation.
     Each source is cached under ``<output_dir>/parts``, keyed by the settings that select its rows,
     so an interrupted run resumes.
+
+    With ``tokenizer`` and ``max_length`` (``distill.py --seq_length``), rows are tokenized as
+    Megatron-Bridge trains them (``loss_mode``, default ``assistant``, as ``--sft_loss_mode``), in
+    ``num_workers`` processes, and rows with no trained token are dropped. ``pack_max_tokens``
+    then orders each split as pairs of at most that many tokens for ``--sft_in_batch_packing
+    --mbs 2``.
     """
     config = load_config(config_path)
     if config.get("output_format") != "chat_jsonl":
@@ -361,9 +427,20 @@ def prepare_chat_data_blend(config_path: Path) -> dict[str, int]:
             "prepare_chat_data_blend writes output_format: chat_jsonl, "
             f"got {config.get('output_format')!r}."
         )
-    ignored = {"tokenizer", "target_tokens"} & config.keys()
-    if ignored:
-        raise ValueError(f"output_format: chat_jsonl does not use {sorted(ignored)}.")
+    if "target_tokens" in config:
+        raise ValueError("output_format: chat_jsonl does not use ['target_tokens'].")
+    if ("tokenizer" in config) != ("max_length" in config):
+        raise ValueError("chat_jsonl tokenizes rows only with both tokenizer and max_length.")
+    if "pack_max_tokens" in config and "max_length" not in config:
+        raise ValueError("pack_max_tokens needs tokenizer and max_length to count tokens.")
+    for key in ("max_length", "pack_max_tokens", "num_workers"):
+        value = config.get(key, 1)
+        if isinstance(value, bool) or not isinstance(value, int) or value <= 0:
+            raise ValueError(f"{key} must be a positive integer, got {value!r}.")
+    if config.get("loss_mode", "assistant") not in ("assistant", "last_turn", "full"):
+        raise ValueError(
+            f"loss_mode must be assistant, last_turn or full: {config['loss_mode']!r}."
+        )
     if any("max_samples" in source for source in config["sources"]):
         raise ValueError("output_format: chat_jsonl does not use max_samples; set target_samples.")
     fields = {source["content_field"] for source in config["sources"]} - {"messages", "text"}
@@ -401,8 +478,17 @@ def prepare_chat_data_blend(config_path: Path) -> dict[str, int]:
         rows.extend(got[:want])
 
     random.Random(seed).shuffle(rows)
+    if "max_length" in config:
+        lengths = _sft_lengths(rows, config)
+        kept = [k for k, n in enumerate(lengths) if n]
+        print(f"Dropped {len(rows) - len(kept)} rows with no trained token", flush=True)
+        rows, lengths = [rows[k] for k in kept], [lengths[k] for k in kept]
     num_validation = min(1000, len(rows) // 20)
     splits = {"validation": rows[:num_validation], "train": rows[num_validation:]}
+    if "pack_max_tokens" in config:
+        cap = config["pack_max_tokens"]
+        splits["validation"] = _pair_rows(splits["validation"], lengths[:num_validation], cap, seed)
+        splits["train"] = _pair_rows(splits["train"], lengths[num_validation:], cap, seed)
     if not splits["train"] or not splits["validation"]:
         raise RuntimeError(f"Too few rows ({len(rows)}) for a train/validation split.")
     for name, split_rows in splits.items():
@@ -420,7 +506,16 @@ def prepare_megatron_data_blend(config_path: Path) -> list[tuple[float, str]]:
             f"Unknown output_format {output_format!r} for a tokenized blend; use megatron, or "
             "chat_jsonl with prepare_chat_data_blend."
         )
-    chat_keys = {"target_samples", "shuffle", "seed", "reasoning_content"} & config.keys()
+    chat_keys = {
+        "target_samples",
+        "shuffle",
+        "seed",
+        "reasoning_content",
+        "max_length",
+        "loss_mode",
+        "pack_max_tokens",
+        "num_workers",
+    } & config.keys()
     if chat_keys:
         raise ValueError(f"Tokenized blends do not use {sorted(chat_keys)}.")
     output_dir = Path(config["output_dir"])
diff --git a/tests/examples/megatron_bridge/test_distill.py b/tests/examples/megatron_bridge/test_distill.py
index 6d146957426..a40e446a34e 100644
--- a/tests/examples/megatron_bridge/test_distill.py
+++ b/tests/examples/megatron_bridge/test_distill.py
@@ -97,6 +97,40 @@ def test_distill_llm_hf_export(tmp_path, num_gpus, fake_hub_checkpoint):
             assert (all_exports / f"iter_{it:07d}" / fake_hub_checkpoint.marker).exists()
 
 
+def test_distill_llm_document_boundaries(tmp_path, num_gpus):
+    """--document_boundaries trains each EOD-separated document of a window as its own THD segment."""
+    # head_dim=8: TE's unfused attention, the only backend for head_dim 2, rejects THD input.
+    hf_path = create_tiny_qwen3_dir(
+        tmp_path, with_tokenizer=True, head_dim=8, max_position_embeddings=2048
+    )
+    distill_output_dir = tmp_path / "distill_output"
+    distill_cmd_parts = extend_cmd_parts(
+        [
+            "torchrun",
+            f"--nproc_per_node={num_gpus}",
+            "distill.py",
+            "--use_mock_data",
+            "--document_boundaries",
+        ],
+        student_hf_path=hf_path,
+        teacher_hf_path=hf_path,
+        output_dir=distill_output_dir,
+        tp_size=num_gpus,
+        pp_size=1,
+        # Mock documents are 1-4095 tokens, so most of these windows span a document boundary.
+        seq_length=2048,
+        mbs=1,
+        gbs=4,
+        train_iters=2,
+        lr_warmup_iters=1,
+        eval_interval=2,
+        eval_iters=0,
+        log_interval=1,
+    )
+    run_example_command(distill_cmd_parts, example_path="megatron_bridge")
+    assert (distill_output_dir / "checkpoints/iter_0000002").exists()
+
+
 def test_distill_llm_sft(tmp_path, num_gpus):
     """--sft distills from prompt-completion jsonl instead of pre-tokenized --data_paths."""
     teacher_hf_path = create_tiny_qwen3_dir(tmp_path, with_tokenizer=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 6306ecef292..db858f22a2d 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
@@ -302,7 +302,11 @@ def test_chat_blend_reselects_after_a_selection_change(chat_blend, first, second
         ({"sources": [_source("a", weight=-1)]}, "nonnegative"),
         ({"sources": [_source("a", weight=0)]}, "not all zero"),
         ({"sources": [_source("a", weight=float("inf"))]}, "finite"),
-        ({"tokenizer": "x"}, "does not use"),
+        ({"tokenizer": "x"}, "both tokenizer and max_length"),
+        ({"max_length": 64}, "both tokenizer and max_length"),
+        ({"pack_max_tokens": 64}, "pack_max_tokens needs tokenizer and max_length"),
+        ({"tokenizer": "x", "max_length": 0}, "max_length must be a positive integer"),
+        ({"tokenizer": "x", "max_length": 64, "loss_mode": "last"}, "loss_mode must be"),
         ({"target_tokens": 1000}, "does not use"),
         ({"reasoning_content": "drop"}, "reasoning_content must be one of"),
         # A row with both columns must not silently be read from ``text``.
@@ -379,3 +383,43 @@ def test_chat_blend_reads_files_sources(tmp_path: Path, monkeypatch):
     assert sorted(r["messages"][0]["content"] for r in rows) == sorted(f"q{i}" for i in range(40))
     # The part is named after its first file.
     assert len(list((out / "parts").glob("org--a_a_*.jsonl"))) == 1
+
+
+def test_chat_blend_drops_untrainable_rows_and_pairs_the_rest(
+    tiny_tokenizer_path: str, tmp_path: Path, chat_blend, monkeypatch
+):
+    def iter_rows(source, output_dir, shuffle, seed):
+        # Prompts of 0..99 words; past max_length the truncation leaves no assistant token.
+        yield from (
+            {
+                "messages": [
+                    {"role": "user", "content": "w " * i},
+                    {"role": "assistant", "content": "a"},
+                ]
+            }
+            for i in range(100)
+        )
+
+    monkeypatch.setattr(blend_module, "_iter_rows", iter_rows)
+    write_config, _ = chat_blend
+    config = {"tokenizer": tiny_tokenizer_path, "max_length": 64, "num_workers": 1}
+    sizes = blend_module.prepare_chat_data_blend(write_config(**config))
+    out = tmp_path / "blend"
+    blend_module._init_sft_tokenizer(tiny_tokenizer_path, 64, "assistant")
+    rows = _read_jsonl(out / "train.jsonl") + _read_jsonl(out / "validation.jsonl")
+    assert 40 < len(rows) < 100  # the 50 + 50 selected rows, minus the long prompts
+    assert all(blend_module._sft_length(row) for row in rows)
+
+    # A cap of 2 * max_length fits any two rows, so only an odd row per split is left out.
+    paired = blend_module.prepare_chat_data_blend(write_config(**config, pack_max_tokens=128))
+    assert paired == {split: n - n % 2 for split, n in sizes.items()}
+
+
+def test_pair_rows_pairs_long_with_short_under_the_cap():
+    lengths = {"a": 10, "b": 60, "c": 30, "d": 50, "e": 95, "f": 20}
+    rows = blend_module._pair_rows(list(lengths), list(lengths.values()), cap=80, seed=0)
+    # e fits with no row and c is left without a partner; rows (2k, 2k + 1) are a pack.
+    assert {frozenset(rows[k : k + 2]) for k in range(0, len(rows), 2)} == {
+        frozenset("ab"),
+        frozenset("df"),
+    }

From 24155ea36bcd8a3892c8686dd8e581a915991f59 Mon Sep 17 00:00:00 2001
From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Date: Wed, 21 Oct 2026 09:38:55 -0700
Subject: [PATCH 2/5] Use the student EOS as the document-boundary EOD; fold
 its test into the HF-export one

NullTokenizer's EOD defaulted to vocab_size - 1, which pre-tokenized data never
contains, so --document_boundaries never split a window. Pass the student
tokenizer's EOS (what --append_eod writes) and fail on data whose first document
does not end with it. Run the document-boundary distillation inside
test_distill_llm_hf_export instead of a separate end-to-end test.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
---
 examples/megatron_bridge/README.md            |  4 +-
 examples/megatron_bridge/distill.py           | 29 ++++++-
 .../examples/megatron_bridge/test_distill.py  | 77 ++++++++++---------
 3 files changed, 71 insertions(+), 39 deletions(-)

diff --git a/examples/megatron_bridge/README.md b/examples/megatron_bridge/README.md
index 382df67004e..04a375c0e0e 100644
--- a/examples/megatron_bridge/README.md
+++ b/examples/megatron_bridge/README.md
@@ -193,7 +193,9 @@ Megatron-Bridge holds every row in host memory on each rank, so keep `train.json
 the run needs.
 
 With pre-tokenized data, `--document_boundaries` (`--mbs 1`, `--cp_size 1`) restricts attention to
-each EOD-separated document of a window instead of the whole window.
+each EOD-separated document of a window instead of the whole window. Tokenize the data with
+`--append_eod`: the student tokenizer's EOS marks each document's end. Chat-template data has no
+EOD between conversations, so distill.py rejects it.
 
 > [!IMPORTANT]
 > Packed training (`--sft_in_batch_packing` or `--document_boundaries`) of Gated DeltaNet models
diff --git a/examples/megatron_bridge/distill.py b/examples/megatron_bridge/distill.py
index c9d2c082845..1788aac91a7 100644
--- a/examples/megatron_bridge/distill.py
+++ b/examples/megatron_bridge/distill.py
@@ -58,6 +58,7 @@
 except ImportError:
     HAS_DIRECT_HF_SFT = False
 from megatron.core.datasets.gpt_dataset import GPTDataset
+from megatron.core.datasets.indexed_dataset import IndexedDataset
 from megatron.core.datasets.utils import get_blend_from_list
 from megatron.core.distributed import DistributedDataParallelConfig
 from megatron.core.utils import unwrap_model
@@ -507,6 +508,28 @@ def _tokenizer_prepends_bos(args) -> bool:
     return tokenizer("x").input_ids[:1] == [tokenizer.bos_token_id]
 
 
+def _document_eod_id(args) -> int:
+    """The token ending each document: the student tokenizer's EOS, which ``--append_eod`` appends.
+
+    Fails if a data file's first document does not end with it, e.g. chat-template data, which has
+    no EOD between conversations, rather than silently attending across documents.
+    """
+    eod = AutoTokenizer.from_pretrained(
+        args.student_hf_path, trust_remote_code=args.trust_remote_code
+    ).eos_token_id
+    if eod is None:
+        raise ValueError("--document_boundaries needs a student tokenizer with an EOS token.")
+    for prefix in get_blend_from_list(args.data_paths)[0] if args.data_paths else []:
+        # Bound to a name: a document is a view of the dataset's memory map, closed with it.
+        dataset = IndexedDataset(prefix)
+        if dataset[0][-1] != eod:
+            raise ValueError(
+                f"--document_boundaries: the first document of {prefix} does not end with the "
+                f"student tokenizer's EOS ({eod}). Tokenize text data with --append_eod."
+            )
+    return eod
+
+
 def _add_document_boundaries() -> None:
     """Give each EOD-separated document of a pre-tokenized window its own THD segment."""
     getitem = GPTDataset.__getitem__
@@ -533,6 +556,8 @@ def main(args: argparse.Namespace, owns_the_run: bool = True):
     args.student_hub_model_id, args.student_local_checkpoint_path = ensure_local_checkpoint(
         args.student_hf_path
     )
+    # NullTokenizer's EOD defaults to vocab_size - 1, which pre-tokenized data never contains.
+    eod_id = _document_eod_id(args) if args.document_boundaries else None
     student_has_modelopt_state = args.student_megatron_path is not None and has_modelopt_state(
         args.student_megatron_path
     )
@@ -799,7 +824,9 @@ def _restore_student_hook(model_chunks):
             )
             if args.sft
             else TokenizerConfig(
-                tokenizer_type="NullTokenizer", vocab_size=distill_provider.vocab_size
+                tokenizer_type="NullTokenizer",
+                vocab_size=distill_provider.vocab_size,
+                null_tokenizer_eod_id=eod_id,
             )
         ),
         checkpoint=CheckpointConfig(
diff --git a/tests/examples/megatron_bridge/test_distill.py b/tests/examples/megatron_bridge/test_distill.py
index a40e446a34e..b030cd7fd7d 100644
--- a/tests/examples/megatron_bridge/test_distill.py
+++ b/tests/examples/megatron_bridge/test_distill.py
@@ -32,10 +32,11 @@
 from transformers import AutoModelForImageTextToText
 
 from modelopt.torch.puzzletron.anymodel import convert_model
+from modelopt.torch.utils.plugins.megatron_preprocess_data import megatron_preprocess_data
 
 sys.path.insert(0, str(Path(__file__).resolve().parents[3] / "examples" / "megatron_bridge"))
 
-from distill import HAS_DIRECT_HF_SFT, _hf_source
+from distill import HAS_DIRECT_HF_SFT, _document_eod_id, _hf_source
 from export_distilled_megatron_to_hf import _get_checkpoint_export_paths
 
 
@@ -43,11 +44,15 @@
 def test_distill_llm_hf_export(tmp_path, num_gpus, fake_hub_checkpoint):
     """Distill with an inline HF export, then re-export every saved iteration standalone.
 
-    Both exports read the student from the copy ``ensure_local_checkpoint`` resolved (see
-    ``fake_hub_checkpoint``). distill.py keeps the local path on its command line, since get_args()
-    loads the tokenizer from it; the standalone export is passed the Hub ID.
+    Distillation runs with --document_boundaries, so each EOD-separated document of a window trains
+    as its own THD segment. Both exports read the student from the copy ``ensure_local_checkpoint``
+    resolved (see ``fake_hub_checkpoint``). distill.py keeps the local path on its command line,
+    since get_args() loads the tokenizer from it; the standalone export is passed the Hub ID.
     """
-    teacher_hf_path = create_tiny_qwen3_dir(tmp_path, with_tokenizer=True)
+    # head_dim=8: TE's unfused attention, the only backend for head_dim 2, rejects THD input.
+    teacher_hf_path = create_tiny_qwen3_dir(
+        tmp_path, with_tokenizer=True, head_dim=8, max_position_embeddings=2048
+    )
     student_export_arg = fake_hub_checkpoint.serve(
         teacher_hf_path, "distill.py", "export_distilled_megatron_to_hf.py", nproc=num_gpus
     )
@@ -56,13 +61,20 @@ def test_distill_llm_hf_export(tmp_path, num_gpus, fake_hub_checkpoint):
     distilled_hf_path = tmp_path / "distilled_hf"
     all_exports = tmp_path / "all_exports"
     distill_cmd_parts = extend_cmd_parts(
-        ["torchrun", f"--nproc_per_node={num_gpus}", "distill.py", "--use_mock_data"],
+        [
+            "torchrun",
+            f"--nproc_per_node={num_gpus}",
+            "distill.py",
+            "--use_mock_data",
+            "--document_boundaries",
+        ],
         student_hf_path=teacher_hf_path,
         teacher_hf_path=teacher_hf_path,
         output_dir=distill_output_dir,
         tp_size=num_gpus,
         pp_size=1,
-        seq_length=16,
+        # Mock documents are 1-4095 tokens, so most of these windows span a document boundary.
+        seq_length=2048,
         mbs=1,
         gbs=4,
         train_iters=train_iters,
@@ -97,38 +109,29 @@ def test_distill_llm_hf_export(tmp_path, num_gpus, fake_hub_checkpoint):
             assert (all_exports / f"iter_{it:07d}" / fake_hub_checkpoint.marker).exists()
 
 
-def test_distill_llm_document_boundaries(tmp_path, num_gpus):
-    """--document_boundaries trains each EOD-separated document of a window as its own THD segment."""
-    # head_dim=8: TE's unfused attention, the only backend for head_dim 2, rejects THD input.
-    hf_path = create_tiny_qwen3_dir(
-        tmp_path, with_tokenizer=True, head_dim=8, max_position_embeddings=2048
+@pytest.mark.parametrize("append_eod", [True, False])
+def test_document_eod_id(tmp_path, append_eod):
+    """Document boundaries use the EOS --append_eod wrote, and fail on data without it."""
+    tokenizer = get_tiny_tokenizer()
+    tokenizer.save_pretrained(tmp_path / "tokenizer")
+    jsonl = tmp_path / "docs.jsonl"
+    jsonl.write_text("".join(json.dumps({"text": f"document number {i}"}) + "\n" for i in range(4)))
+    (prefix,) = megatron_preprocess_data(
+        jsonl_paths=jsonl,
+        output_dir=tmp_path / "data",
+        tokenizer_name_or_path=str(tmp_path / "tokenizer"),
+        append_eod=append_eod,
     )
-    distill_output_dir = tmp_path / "distill_output"
-    distill_cmd_parts = extend_cmd_parts(
-        [
-            "torchrun",
-            f"--nproc_per_node={num_gpus}",
-            "distill.py",
-            "--use_mock_data",
-            "--document_boundaries",
-        ],
-        student_hf_path=hf_path,
-        teacher_hf_path=hf_path,
-        output_dir=distill_output_dir,
-        tp_size=num_gpus,
-        pp_size=1,
-        # Mock documents are 1-4095 tokens, so most of these windows span a document boundary.
-        seq_length=2048,
-        mbs=1,
-        gbs=4,
-        train_iters=2,
-        lr_warmup_iters=1,
-        eval_interval=2,
-        eval_iters=0,
-        log_interval=1,
+    args = argparse.Namespace(
+        student_hf_path=str(tmp_path / "tokenizer"),
+        trust_remote_code=False,
+        data_paths=["1.0", prefix],
     )
-    run_example_command(distill_cmd_parts, example_path="megatron_bridge")
-    assert (distill_output_dir / "checkpoints/iter_0000002").exists()
+    if append_eod:
+        assert _document_eod_id(args) == tokenizer.eos_token_id
+    else:
+        with pytest.raises(ValueError, match="--append_eod"):
+            _document_eod_id(args)
 
 
 def test_distill_llm_sft(tmp_path, num_gpus):

From 26d20f7a357ce571bf63d4d9cef4113548d4fc02 Mon Sep 17 00:00:00 2001
From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Date: Thu, 22 Oct 2026 00:44:45 -0700
Subject: [PATCH 3/5] Test the document-boundary cu_seqlens; fix the mock-data
 segment comment

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
---
 .../examples/megatron_bridge/test_distill.py  | 29 +++++++++++++++++--
 1 file changed, 27 insertions(+), 2 deletions(-)

diff --git a/tests/examples/megatron_bridge/test_distill.py b/tests/examples/megatron_bridge/test_distill.py
index b030cd7fd7d..98d7ebbb390 100644
--- a/tests/examples/megatron_bridge/test_distill.py
+++ b/tests/examples/megatron_bridge/test_distill.py
@@ -18,6 +18,7 @@
 import json
 import sys
 from pathlib import Path
+from types import SimpleNamespace
 
 import pytest
 import torch
@@ -29,6 +30,7 @@
     create_tiny_qwen3_moe_dir,
     get_tiny_tokenizer,
 )
+from megatron.core.datasets.gpt_dataset import GPTDataset
 from transformers import AutoModelForImageTextToText
 
 from modelopt.torch.puzzletron.anymodel import convert_model
@@ -36,7 +38,7 @@
 
 sys.path.insert(0, str(Path(__file__).resolve().parents[3] / "examples" / "megatron_bridge"))
 
-from distill import HAS_DIRECT_HF_SFT, _document_eod_id, _hf_source
+from distill import HAS_DIRECT_HF_SFT, _add_document_boundaries, _document_eod_id, _hf_source
 from export_distilled_megatron_to_hf import _get_checkpoint_export_paths
 
 
@@ -73,7 +75,8 @@ def test_distill_llm_hf_export(tmp_path, num_gpus, fake_hub_checkpoint):
         output_dir=distill_output_dir,
         tp_size=num_gpus,
         pp_size=1,
-        # Mock documents are 1-4095 tokens, so most of these windows span a document boundary.
+        # Mock token ids cycle through the 128-token vocab, so the EOD id recurs every 128 tokens
+        # and each window trains as ~17 THD segments.
         seq_length=2048,
         mbs=1,
         gbs=4,
@@ -134,6 +137,28 @@ def test_document_eod_id(tmp_path, append_eod):
             _document_eod_id(args)
 
 
+@pytest.mark.parametrize(
+    ("tokens", "cu_seqlens"),
+    [
+        ([1, 2, 9, 3, 4, 5, 9, 6], [0, 3, 7, 8]),  # EOD mid-window
+        ([1, 2, 3, 9, 4, 5, 6, 9], [0, 4, 8]),  # EOD as the last token
+        ([9, 1, 2, 3, 4, 5, 6, 7], [0, 1, 8]),  # EOD as the first token
+        ([1, 2, 3, 4, 5, 6, 7, 8], [0, 8]),  # no EOD
+    ],
+)
+def test_document_boundaries_cu_seqlens(monkeypatch, tokens, cu_seqlens):
+    """Each EOD ends a THD segment; the window's end closes the last one."""
+    monkeypatch.setattr(
+        GPTDataset, "__getitem__", lambda self, idx: {"tokens": torch.tensor(tokens)}
+    )
+    _add_document_boundaries()  # monkeypatch restores the original __getitem__ on teardown
+    dataset = SimpleNamespace(config=SimpleNamespace(tokenizer=SimpleNamespace(eod=9)))
+    sample = GPTDataset.__getitem__(dataset, 0)
+    expected = torch.tensor(cu_seqlens, dtype=torch.int32)
+    assert torch.equal(sample["cu_seqlens_q"], expected)
+    assert sample["max_seqlen_q"] == (expected[1:] - expected[:-1]).max()
+
+
 def test_distill_llm_sft(tmp_path, num_gpus):
     """--sft distills from prompt-completion jsonl instead of pre-tokenized --data_paths."""
     teacher_hf_path = create_tiny_qwen3_dir(tmp_path, with_tokenizer=True)

From 2d6c33af9365f9dbfa63dfeb0a5e3088b6e41de4 Mon Sep 17 00:00:00 2001
From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Date: Thu, 22 Oct 2026 01:21:00 -0700
Subject: [PATCH 4/5] docs: note that CP padding can push a packed pair past
 pack_max_tokens

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
---
 examples/dataset/MEGATRON_DATA_PREP.md | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/examples/dataset/MEGATRON_DATA_PREP.md b/examples/dataset/MEGATRON_DATA_PREP.md
index 357203a4264..6dbed0b9e78 100644
--- a/examples/dataset/MEGATRON_DATA_PREP.md
+++ b/examples/dataset/MEGATRON_DATA_PREP.md
@@ -234,9 +234,10 @@ num_workers: 48         # Optional; tokenizer processes (default: all CPUs).
 With `tokenizer` and `max_length`, a conversation with no trained token left after truncation is
 dropped, since Megatron-Bridge fails on it during training. `pack_max_tokens` pairs long
 conversations with short ones and writes each split pair by pair, so `--mbs 2` packs every pair
-into one micro-batch; a conversation that fits with no other is dropped. Tokenizing runs on the
-CPU: ~0.4 s per conversation per process for GLM-5.3-Flash on Nemotron-Cascade-2 (4.4K tokens
-on average).
+into one micro-batch; a conversation that fits with no other is dropped. With `--cp_size` > 1,
+Megatron-Bridge pads each conversation to a multiple of `2 * cp_size`, so a pair can exceed the cap
+by under `2 * cp_size` tokens per conversation. Tokenizing runs on the CPU: ~0.4 s per
+conversation per process for GLM-5.3-Flash on Nemotron-Cascade-2 (4.4K tokens on average).
 
 ## `reasoning_content` for Post-Training v3 Datasets
 

From 17964edf08fdfd4be5d0ca1824b5dd8feac12122 Mon Sep 17 00:00:00 2001
From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Date: Fri, 23 Oct 2026 00:46:05 -0700
Subject: [PATCH 5/5] Keep assistant-less rows for loss_mode full; raise
 chat-tokenization setup errors

With loss_mode: full, Megatron-Bridge trains rows without an assistant turn,
so the chat blend keeps them. Only Bridge's no-loss-span error marks a row as
untrainable; other ValueErrors (e.g. a tokenizer without a chat template)
would have dropped every row and are raised.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
---
 .../plugins/prepare_megatron_data_blend.py    | 27 ++++++++++++++-----
 .../test_prepare_megatron_data_blend.py       | 16 ++++++++---
 2 files changed, 33 insertions(+), 10 deletions(-)

diff --git a/modelopt/torch/utils/plugins/prepare_megatron_data_blend.py b/modelopt/torch/utils/plugins/prepare_megatron_data_blend.py
index 0a0d0ca7465..b2f6bae779e 100644
--- a/modelopt/torch/utils/plugins/prepare_megatron_data_blend.py
+++ b/modelopt/torch/utils/plugins/prepare_megatron_data_blend.py
@@ -179,7 +179,10 @@ def _prepare_sources(
 
 
 def _to_chat(
-    row: dict[str, Any], content_field: str, reasoning_content: ReasoningContentMode = "native"
+    row: dict[str, Any],
+    content_field: str,
+    reasoning_content: ReasoningContentMode = "native",
+    loss_mode: str = "assistant",
 ) -> dict[str, Any] | None:
     """Convert a source row to ``{"messages": [...]}``, or ``None`` if it has nothing to train on.
 
@@ -225,8 +228,10 @@ def _to_chat(
             if content.strip()
         ]
         chat = {"messages": turns}
-    # Nothing in a row without an assistant turn is trained under assistant-only loss.
-    if not any(turn.get("role") == "assistant" for turn in chat["messages"]):
+    # Only loss_mode full trains a row without an assistant turn.
+    if loss_mode != "full" and not any(
+        turn.get("role") == "assistant" for turn in chat["messages"]
+    ):
         return None
     return chat
 
@@ -261,6 +266,7 @@ def _take_chat_rows(
     shuffle: bool,
     seed: int,
     reasoning_content: ReasoningContentMode,
+    loss_mode: str,
 ) -> list[dict[str, Any]]:
     """Return the first ``want`` usable chat rows of a source.
 
@@ -278,7 +284,7 @@ def _take_chat_rows(
     max_rejected = 10 * want + 100
     try:
         for row in _iter_rows(source, output_dir, shuffle, seed):
-            chat = _to_chat(row, source["content_field"], reasoning_content)
+            chat = _to_chat(row, source["content_field"], reasoning_content, loss_mode)
             # Judge the source row's top-level media keys too, which the conversion does not keep.
             if chat is None or not is_text_only_chat_example({**row, **chat}):
                 rejected += 1
@@ -352,7 +358,11 @@ def _sft_length(row: dict[str, Any]) -> int:
 
     try:
         tokenized = tokenize_chat_example(row, warn_on_all_masked=False, **_SFT_TOKENIZER)
-    except ValueError:  # e.g. no assistant span matched: Bridge raises from a dataloader worker
+    except ValueError as e:
+        # No assistant span matched: Bridge raises from a dataloader worker. Other errors are setup
+        # errors (e.g. a tokenizer without a chat template) that would drop every row.
+        if "did not match any loss-contributing spans" not in str(e):
+            raise
         return 0
     return tokenized.input_ids.numel() if tokenized.assistant_mask.any() else 0
 
@@ -464,6 +474,9 @@ def prepare_chat_data_blend(config_path: Path) -> dict[str, int]:
     settings = {"shuffle": shuffle, "reasoning_content": reasoning_content}
     if shuffle:
         settings["seed"] = seed
+    loss_mode = config.get("loss_mode", "assistant")
+    if loss_mode == "full":  # keeps rows without an assistant turn
+        settings["loss_mode"] = loss_mode
 
     rows: list[dict[str, Any]] = []
     for source, want in zip(config["sources"], budgets):
@@ -473,7 +486,9 @@ def prepare_chat_data_blend(config_path: Path) -> dict[str, int]:
         part = _part_path(parts_dir, source, settings)
         got = _read_jsonl(part)
         if got is None or len(got) < want:
-            got = _take_chat_rows(source, output_dir, want, shuffle, seed, reasoning_content)
+            got = _take_chat_rows(
+                source, output_dir, want, shuffle, seed, reasoning_content, loss_mode
+            )
             _write_jsonl(part, got)
         rows.extend(got[:want])
 
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 db858f22a2d..a037f6f459d 100644
--- a/tests/gpu_megatron/torch/utils/plugins/test_prepare_megatron_data_blend.py
+++ b/tests/gpu_megatron/torch/utils/plugins/test_prepare_megatron_data_blend.py
@@ -176,8 +176,11 @@ def test_to_chat_handles_all_row_shapes():
     msgs = [{"role": "user", "content": "hi"}, {"role": "assistant", "content": "hello"}]
     assert blend_module._to_chat({"messages": msgs}, "messages") == {"messages": msgs}
     assert blend_module._to_chat({"messages": []}, "messages") is None
-    # Nothing would be trained under assistant-only loss.
+    # Nothing would be trained under assistant-only loss; loss_mode full trains the user turn.
     assert blend_module._to_chat({"messages": msgs[:1]}, "messages") is None
+    assert blend_module._to_chat({"messages": msgs[:1]}, "messages", loss_mode="full") == {
+        "messages": msgs[:1]
+    }
     # Tool definitions and a top-level system prompt reach the chat path.
     tools = [{"type": "function", "function": {"name": "add"}}]
     row = {"messages": msgs, "tools": tools, "system": "Be brief.", "id": 7}
@@ -228,13 +231,13 @@ def test_take_chat_rows_skips_non_text_rows_and_fails_without_enough(tmp_path: P
     # Media is rejected inside message content and as a top-level key the conversion drops.
     rows = [_chat(0), _image_chat(1), {**_chat(9), "images": ["a.png"]}, _chat(2), _chat(3)]
     monkeypatch.setattr(blend_module, "_iter_rows", lambda *a: iter(rows))
-    got = blend_module._take_chat_rows(source, tmp_path, 3, False, 0, "native")
+    got = blend_module._take_chat_rows(source, tmp_path, 3, False, 0, "native", "assistant")
     assert [c["messages"][0]["content"] for c in got] == ["q0", "q2", "q3"]
 
     images = [_image_chat(i) for i in range(200)]
     monkeypatch.setattr(blend_module, "_iter_rows", lambda *a: iter(images))
     with pytest.raises(RuntimeError, match=r"only 0/2 usable rows \(gave up after 121 rejected"):
-        blend_module._take_chat_rows(source, tmp_path, 2, False, 0, "native")
+        blend_module._take_chat_rows(source, tmp_path, 2, False, 0, "native", "assistant")
 
 
 def test_chat_blend_splits_weighted_rows(tmp_path: Path, chat_blend, monkeypatch):
@@ -281,8 +284,9 @@ def test_chat_blend_reuses_and_rebuilds_cached_parts(tmp_path: Path, chat_blend)
         ({}, {"shuffle": True}),
         ({"shuffle": True, "seed": 0}, {"shuffle": True, "seed": 1}),
         ({}, {"reasoning_content": "inline"}),
+        ({}, {"loss_mode": "full"}),
     ],
-    ids=["split", "files", "shuffle", "seed", "reasoning_content"],
+    ids=["split", "files", "shuffle", "seed", "reasoning_content", "loss_mode"],
 )
 def test_chat_blend_reselects_after_a_selection_change(chat_blend, first, second):
     write_config, calls = chat_blend
@@ -409,6 +413,10 @@ def iter_rows(source, output_dir, shuffle, seed):
     rows = _read_jsonl(out / "train.jsonl") + _read_jsonl(out / "validation.jsonl")
     assert 40 < len(rows) < 100  # the 50 + 50 selected rows, minus the long prompts
     assert all(blend_module._sft_length(row) for row in rows)
+    # A setup error would drop every row, so it is raised instead.
+    monkeypatch.setitem(blend_module._SFT_TOKENIZER, "loss_mode", "unknown")
+    with pytest.raises(ValueError, match="loss_mode"):
+        blend_module._sft_length(rows[0])
 
     # A cap of 2 * max_length fits any two rows, so only an odd row per split is left out.
     paired = blend_module.prepare_chat_data_blend(write_config(**config, pack_max_tokens=128))