Megatron-Bridge distillation: packed chat-SFT data prep and per-document attention

CHANGELOG.rst

@@ -21,6 +21,7 @@ Changelog
 *Megatron Framework (M-LM / M-Bridge)*
 
 - Add ``--sft_in_batch_packing``, ``--moe_aux_loss_coeff`` and ``--dsa_kernel_backend`` to ``examples/megatron_bridge/distill.py`` (``--dsa_kernel_backend`` also to ``quantize.py``): pack ``--sft_hf_dataset`` conversations per micro-batch instead of padding each to ``--seq_length``, override the MoE load-balancing loss weight (e.g. 0 for padded SFT data), and select fused DeepSeek sparse-attention kernels for long sequences.
+- Add ``--document_boundaries`` to ``examples/megatron_bridge/distill.py`` to attend within each document of a pre-tokenized window, and ``tokenizer`` / ``max_length`` / ``pack_max_tokens`` to ``chat_jsonl`` blends to drop conversations with nothing to train and pair the rest for ``--sft_in_batch_packing``. Packed distillation of Gated DeltaNet models needs ``flash-linear-attention`` >= 0.5.2 (earlier versions compute it wrongly); install ``examples/megatron_bridge/requirements.txt`` in ``nemo:26.08``.
 - Add ``--num_layers_in_first_pipeline_stage``, ``--num_layers_in_last_pipeline_stage`` and ``--recompute_num_layers_per_stage`` to ``examples/megatron_bridge/distill.py`` to balance pipeline-parallel distillation: give the last stage, which also runs the output layers and KD loss, fewer layers, and recompute fewer layers on later stages.
 - Add ``--vllm_fakequant_export`` to ``examples/megatron_bridge/export_quantized_megatron_to_hf.py`` to evaluate formats vLLM has no kernels for (e.g. GGML IQ) as BF16 fake-quant checkpoints, and allow ``--ep_size`` > 1 for checkpoints built with ``--no_moe_grouped_gemm``.
 - GGML GPTQ checkpoints from ``examples/megatron_bridge`` keep their calibrated codes when the Megatron checkpoint is reloaded for QAD or export, and layerwise calibration now runs on Megatron vision-language models.

examples/dataset/MEGATRON_DATA_PREP.md

@@ -221,6 +221,24 @@ rows, or runs out of them stops the run. Each source's selection is cached under
 the settings that choose its rows, so an interrupted run resumes and a changed configuration selects afresh.
 Nothing prunes `parts`; delete it once `train.jsonl` is written to reclaim the space.
 
+To feed `distill.py --sft_in_batch_packing`, also count tokens the way Megatron-Bridge trains them:
+
+```yaml
+tokenizer: Qwen/Qwen3.6-35B-A3B  # The student's tokenizer.
+max_length: 32768       # distill.py --seq_length; longer conversations are truncated.
+loss_mode: assistant    # Optional; as distill.py --sft_loss_mode.
+pack_max_tokens: 33280  # Optional; order each split as pairs of at most this many tokens.
+num_workers: 48         # Optional; tokenizer processes (default: all CPUs).
+```
+
+With `tokenizer` and `max_length`, a conversation with no trained token left after truncation is
+dropped, since Megatron-Bridge fails on it during training. `pack_max_tokens` pairs long
+conversations with short ones and writes each split pair by pair, so `--mbs 2` packs every pair
+into one micro-batch; a conversation that fits with no other is dropped. With `--cp_size` > 1,
+Megatron-Bridge pads each conversation to a multiple of `2 * cp_size`, so a pair can exceed the cap
+by under `2 * cp_size` tokens per conversation. Tokenizing runs on the CPU: ~0.4 s per
+conversation per process for GLM-5.3-Flash on Nemotron-Cascade-2 (4.4K tokens on average).
+
 ## `reasoning_content` for Post-Training v3 Datasets
 
 v3 datasets include a `reasoning_content` field in assistant messages (chain-of-thought separate from

examples/megatron_bridge/README.md

@@ -184,6 +184,25 @@ the last one (`last_turn`), or the whole conversation (`full`). With `--eval_ite
 `--sft_hf_validation` in the same format. This path needs a Megatron-Bridge with
 `DirectHFSFTDatasetConfig`.
 
+Megatron-Bridge pads every conversation to `--seq_length` when `--ep_size` > 1, so for MoE models
+add `--sft_in_batch_packing --mbs 2 --moe_aux_loss_coeff 0`: each micro-batch packs two
+conversations into one unpadded THD sequence, and the MoE load-balancing loss, which would also
+run over padding, is dropped. Order the rows as pairs that fit together with the chat-blend
+`pack_max_tokens` option ([Dataset Preparation README](../dataset/MEGATRON_DATA_PREP.md#chat-jsonl-blends-for-assistant-only-loss)).
+Megatron-Bridge holds every row in host memory on each rank, so keep `train.jsonl` to the rows
+the run needs.
+
+With pre-tokenized data, `--document_boundaries` (`--mbs 1`, `--cp_size 1`) restricts attention to
+each EOD-separated document of a window instead of the whole window. Tokenize the data with
+`--append_eod`: the student tokenizer's EOS marks each document's end. Chat-template data has no
+EOD between conversations, so distill.py rejects it.
+
+> [!IMPORTANT]
+> Packed training (`--sft_in_batch_packing` or `--document_boundaries`) of Gated DeltaNet models
+> (e.g. Qwen3.5 / Qwen3.6) needs `flash-linear-attention` >= 0.5.2: earlier versions compute the
+> packed linear attention wrongly. `nemo:26.08` ships 0.4.2, so run
+> `python -m pip install -r requirements.txt` in this directory first.
+
 Teacher and student must share a tokenizer — distillation scores the teacher on the student's
 token ids, and the KD losses compare the two models' logits elementwise over the vocab dimension.
 

examples/megatron_bridge/distill.py

@@ -58,6 +58,8 @@
 except ImportError:
     HAS_DIRECT_HF_SFT = False
 from megatron.core import parallel_state
+from megatron.core.datasets.gpt_dataset import GPTDataset
+from megatron.core.datasets.indexed_dataset import IndexedDataset
 from megatron.core.datasets.utils import get_blend_from_list
 from megatron.core.distributed import DistributedDataParallelConfig
 from megatron.core.utils import unwrap_model
@@ -214,6 +216,12 @@ def get_args():
         nargs="+",
         help="List of tokenized data paths to load from (weight1 path1 weight2 path2 ...)",
     )
+    parser.add_argument(
+        "--document_boundaries",
+        action="store_true",
+        help="Pre-tokenized data only: attend within each EOD-separated document of a sequence (THD "
+        "segments) instead of across the whole packed window. Needs --mbs 1 and --cp_size 1.",
+    )
     parser.add_argument(
         "--data_path_to_cache", type=str, default=None, help="Path to cache the dataset indices"
     )
@@ -492,6 +500,9 @@ def get_args():
         raise ValueError(
             "--sft_hf_validation / --sft_loss_mode / --sft_in_batch_packing require --sft_hf_dataset."
         )
+    if args.document_boundaries and (args.sft or args.mbs != 1 or args.cp_size != 1):
+        # Each window carries its own number of documents, so windows cannot be stacked.
+        raise ValueError("--document_boundaries needs pre-tokenized data, --mbs 1 and --cp_size 1.")
     if args.sft and (args.data_paths or args.use_mock_data):
         raise ValueError(
             "--sft is mutually exclusive with --data_paths / --use_mock_data: the SFT branch wins "
@@ -540,13 +551,56 @@ def _tokenizer_prepends_bos(args) -> bool:
     return tokenizer("x").input_ids[:1] == [tokenizer.bos_token_id]
 
 
+def _document_eod_id(args) -> int:
+    """The token ending each document: the student tokenizer's EOS, which ``--append_eod`` appends.
+
+    Fails if a data file's first document does not end with it, e.g. chat-template data, which has
+    no EOD between conversations, rather than silently attending across documents.
+    """
+    eod = AutoTokenizer.from_pretrained(
+        args.student_hf_path, trust_remote_code=args.trust_remote_code
+    ).eos_token_id
+    if eod is None:
+        raise ValueError("--document_boundaries needs a student tokenizer with an EOS token.")
+    for prefix in get_blend_from_list(args.data_paths)[0] if args.data_paths else []:
+        # Bound to a name: a document is a view of the dataset's memory map, closed with it.
+        dataset = IndexedDataset(prefix)
+        if dataset[0][-1] != eod:
+            raise ValueError(
+                f"--document_boundaries: the first document of {prefix} does not end with the "
+                f"student tokenizer's EOS ({eod}). Tokenize text data with --append_eod."
+            )
+    return eod
+
+
+def _add_document_boundaries() -> None:
+    """Give each EOD-separated document of a pre-tokenized window its own THD segment."""
+    getitem = GPTDataset.__getitem__
+
+    def getitem_with_boundaries(self, idx):
+        sample = getitem(self, idx)
+        tokens = sample["tokens"]
+        n = tokens.numel()
+        ends = (tokens == self.config.tokenizer.eod).nonzero(as_tuple=True)[0] + 1
+        # On the tokens' (CPU) device: CUDA tensors from a dataloader worker are shared over IPC.
+        bounds = [tokens.new_zeros(1), ends[ends < n], tokens.new_tensor([n])]
+        cu_seqlens = torch.cat(bounds).to(torch.int32)
+        sample["cu_seqlens_q"] = cu_seqlens
+        sample["max_seqlen_q"] = (cu_seqlens[1:] - cu_seqlens[:-1]).max()
+        return sample
+
+    GPTDataset.__getitem__ = getitem_with_boundaries
+
+
 def main(args: argparse.Namespace, owns_the_run: bool = True):
     # --student_hf_path stays as given. The student is built and exported from the local copy of
     # the whole checkpoint in student_local_checkpoint_path, since the HF export reads the source's
     # files from local disk only; student_hub_model_id is the Hub ID, or None for a local path.
     args.student_hub_model_id, args.student_local_checkpoint_path = ensure_local_checkpoint(
         args.student_hf_path
     )
+    # NullTokenizer's EOD defaults to vocab_size - 1, which pre-tokenized data never contains.
+    eod_id = _document_eod_id(args) if args.document_boundaries else None
     student_has_modelopt_state = args.student_megatron_path is not None and has_modelopt_state(
         args.student_megatron_path
     )
@@ -683,6 +737,8 @@ def _restore_student_hook(model_chunks):
             return model_chunks
 
         distill_provider.register_pre_wrap_hook(_restore_student_hook, prepend=True)
+    if args.document_boundaries:
+        _add_document_boundaries()
 
     if args.recompute_num_layers_per_stage is not None:
         # Set this pipeline stage's student recompute depth before the model is wrapped.
@@ -714,7 +770,7 @@ def _per_stage_recompute_hook(model_chunks):
         "path_to_cache": args.data_path_to_cache,
         "random_seed": args.seed,
         "reset_attention_mask": False,
-        "reset_position_ids": False,
+        "reset_position_ids": args.document_boundaries,
         "eod_mask_loss": False,
         "num_dataset_builder_threads": 1,
         "data_sharding": True,
@@ -829,7 +885,9 @@ def _per_stage_recompute_hook(model_chunks):
             )
             if args.sft
             else TokenizerConfig(
-                tokenizer_type="NullTokenizer", vocab_size=distill_provider.vocab_size
+                tokenizer_type="NullTokenizer",
+                vocab_size=distill_provider.vocab_size,
+                null_tokenizer_eod_id=eod_id,
             )
         ),
         checkpoint=CheckpointConfig(

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