examples/megatron_bridge: distill on chat data straight from HuggingFace/JSONL

CHANGELOG.rst

@@ -50,6 +50,7 @@ Changelog
 
 *Megatron Framework (M-LM / M-Bridge)*
 
+- 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 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 <https://github.com/NVIDIA/Model-Optimizer/tree/main/examples/megatron_bridge/tutorials/Qwen3.6-35B-A3B/>`_ for details.

examples/megatron_bridge/README.md

@@ -175,6 +175,15 @@ token is appended after the response. A record longer than `--seq_length` is tru
 **start** of `"input"`, which drops any system prompt or opening role marker baked in there, so
 pre-filter or pre-truncate the corpus if that matters.
 
+To distill on **chat data** instead, pass `--sft --sft_hf_dataset <spec>`, where `<spec>` is a
+local `<file>.jsonl` of `{"messages": [{"role": ..., "content": ...}, ...]}` records or a Hugging
+Face Hub dataset `<hub_id>[:<split>]` (split defaults to `train`, and may be a slice such as
+`train[:1000]`). Each conversation is rendered with the model's own chat template, and
+`--sft_loss_mode` picks the trained tokens: every assistant turn (`assistant`, the default), only
+the last one (`last_turn`), or the whole conversation (`full`). With `--eval_iters > 0`, also pass
+`--sft_hf_validation` in the same format. This path needs a Megatron-Bridge with
+`DirectHFSFTDatasetConfig`.
+
 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

@@ -48,6 +48,15 @@
 from megatron.bridge.training.post_training.checkpointing import has_modelopt_state
 from megatron.bridge.training.post_training.distillation import ModelOptDistillConfig
 from megatron.bridge.utils.vocab_utils import calculate_padded_vocab_size
+
+try:
+    from megatron.bridge.data.builders import DirectHFSFTDatasetConfig
+    from megatron.bridge.data.sft_processing import ChatSFTPreprocessingConfig
+    from megatron.bridge.data.sources import HFDatasetSourceConfig
+
+    HAS_DIRECT_HF_SFT = True
+except ImportError:
+    HAS_DIRECT_HF_SFT = False
 from megatron.core.datasets.utils import get_blend_from_list
 from megatron.core.distributed import DistributedDataParallelConfig
 from megatron.core.utils import unwrap_model
@@ -117,6 +126,16 @@ def _nonnegative_int(value: str) -> int:
     return parsed
 
 
+def _hf_source(spec: str) -> "HFDatasetSourceConfig":
+    """``<file>.json[l]`` -> local chat jsonl; otherwise ``<hub_id>[:<split>]`` (split defaults to train)."""
+    if spec.endswith((".json", ".jsonl")):
+        return HFDatasetSourceConfig(
+            path_or_dataset="json", split="train", load_kwargs={"data_files": spec}
+        )
+    dataset, _, split = spec.partition(":")
+    return HFDatasetSourceConfig(path_or_dataset=dataset, split=split or "train")
+
+
 def get_args():
     """Parse command-line arguments."""
     parser = argparse.ArgumentParser(description="Distillation for Megatron-Bridge.")
@@ -181,8 +200,9 @@ def get_args():
     parser.add_argument(
         "--sft",
         action="store_true",
-        help="Distill on prompt-completion jsonl from --sft_dataset_root with the loss masked to "
-        "the completion, instead of pre-tokenized --data_paths.",
+        help="Distill on SFT data with the loss masked to the responses, instead of pre-tokenized "
+        "--data_paths: prompt-completion jsonl from --sft_dataset_root, or chat data from "
+        "--sft_hf_dataset.",
     )
     parser.add_argument(
         "--sft_dataset_root",
@@ -192,6 +212,32 @@ def get_args():
         '{"input": <prompt>, "output": <response>} records (used with --sft). See the README for '
         "how the fields are tokenized and truncated.",
     )
+    parser.add_argument(
+        "--sft_hf_dataset",
+        type=str,
+        default=None,
+        help="Chat dataset to distill on (used with --sft): a local chat '<file>.jsonl', or "
+        "'<hub_id>[:<split>]' (split defaults to train). Only .json/.jsonl are read as local "
+        "files; anything else is taken as a Hub id. Each conversation is rendered by the "
+        "model's own chat template. Alternative to --sft_dataset_root, which expects "
+        '{"input", "output"} records already preprocessed into a dataset root.',
+    )
+    parser.add_argument(
+        "--sft_hf_validation",
+        type=str,
+        default=None,
+        help="Validation data for --sft_hf_dataset, in the same format (required when "
+        "--eval_iters > 0), e.g. '<hub_id>:test' or '<hub_id>:train[:1000]'.",
+    )
+    parser.add_argument(
+        "--sft_loss_mode",
+        type=str,
+        default="assistant",
+        choices=["assistant", "last_turn", "full"],
+        help="Which tokens --sft_hf_dataset trains on: every assistant turn, only the final one, "
+        "or the whole conversation. Multi-turn records train all of their responses under "
+        "'assistant'.",
+    )
     # Training & Eval arguments
     parser.add_argument(
         "--output_dir", type=str, required=True, help="Folder for logging and checkpoint saving"
@@ -354,19 +400,32 @@ def get_args():
     if args.validate_only and args.eval_iters == 0:
         raise ValueError("--validate_only requires --eval_iters > 0.")
 
-    if args.sft and not args.sft_dataset_root:
+    if args.sft and not (args.sft_dataset_root or args.sft_hf_dataset):
         raise ValueError(
             "--sft requires --sft_dataset_root (a directory with training.jsonl, plus "
-            "validation.jsonl when --eval_iters > 0)."
+            "validation.jsonl when --eval_iters > 0) or --sft_hf_dataset."
         )
+    if args.sft_dataset_root and args.sft_hf_dataset:
+        raise ValueError("--sft_dataset_root and --sft_hf_dataset are mutually exclusive.")
+    if args.sft_hf_dataset:
+        if not args.sft:
+            raise ValueError("--sft_hf_dataset requires --sft.")
+        if not HAS_DIRECT_HF_SFT:
+            raise ValueError(
+                "--sft_hf_dataset needs a newer Megatron-Bridge (DirectHFSFTDatasetConfig)."
+            )
+        if args.eval_iters > 0 and not args.sft_hf_validation:
+            raise ValueError("--sft_hf_dataset with --eval_iters > 0 needs --sft_hf_validation.")
+    elif args.sft_hf_validation or args.sft_loss_mode != "assistant":
+        raise ValueError("--sft_hf_validation / --sft_loss_mode require --sft_hf_dataset.")
     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 "
             "the dataset selection, so those inputs would be silently ignored."
         )
     if args.sft_dataset_root and not args.sft:
         raise ValueError("--sft_dataset_root requires --sft; without it the SFT path is not used.")
-    if args.sft:
+    if args.sft and args.sft_dataset_root:
         # Fail on a mistyped root here rather than after both checkpoints have loaded onto GPUs.
         required = ["training.jsonl"] + (["validation.jsonl"] if args.eval_iters > 0 else [])
         absent = [f for f in required if not os.path.isfile(os.path.join(args.sft_dataset_root, f))]
@@ -557,7 +616,21 @@ def _restore_student_hook(model_chunks):
         "dataloader_type": "single",
         "skip_getting_attention_mask_from_dataset": True,
     }
-    if args.sft:
+    if args.sft and args.sft_hf_dataset:
+        # Chat rows rendered by the model's own chat template, with the loss covering every
+        # assistant turn, so multi-round records train all of their responses rather than the last.
+        dataset_config = DirectHFSFTDatasetConfig(
+            seq_length=args.seq_length,
+            source=_hf_source(args.sft_hf_dataset),
+            validation_source=(
+                _hf_source(args.sft_hf_validation) if args.sft_hf_validation else None
+            ),
+            preprocessing=ChatSFTPreprocessingConfig(loss_mode=args.sft_loss_mode),
+            dataloader_type="batch",
+            do_validation=args.eval_iters > 0,
+            do_test=False,
+        )
+    elif args.sft:
         # SFT-masked distillation via Bridge's FinetuningDatasetConfig -> NeMo-style GPTSFTDataset,
         # reading {"input", "output"} jsonl. Fields are tokenized as written except that each is
         # ``.strip(" ")``-ed; see --sft_dataset_root help.

tests/examples/megatron_bridge/test_distill.py

@@ -34,6 +34,7 @@
 
 sys.path.insert(0, str(Path(__file__).resolve().parents[3] / "examples" / "megatron_bridge"))
 
+from distill import HAS_DIRECT_HF_SFT, _hf_source
 from export_distilled_megatron_to_hf import _get_checkpoint_export_paths
 
 
@@ -120,6 +121,75 @@ def test_distill_llm_sft(tmp_path, num_gpus):
     assert (distill_output_dir / f"checkpoints/iter_{train_iters:07d}").exists()
 
 
+@pytest.mark.skipif(not HAS_DIRECT_HF_SFT, reason="Megatron-Bridge lacks DirectHFSFTDatasetConfig")
+def test_distill_llm_sft_hf_chat(tmp_path, num_gpus):
+    """--sft_hf_dataset distills on chat jsonl rendered by the model's chat template."""
+    teacher_hf_path = create_tiny_qwen3_dir(tmp_path, with_tokenizer=True)
+    train_iters = 2
+    gbs = 4
+    # More records than train_iters * gbs so the sampler does not run dry; every 4th is multi-turn.
+    records = []
+    for i in range(64):
+        messages = [
+            {"role": "user", "content": f"what follows {i}?"},
+            {"role": "assistant", "content": f"{i + 1}"},
+        ]
+        if i % 4 == 0:
+            messages += [
+                {"role": "user", "content": "and then?"},
+                {"role": "assistant", "content": f"{i + 2}"},
+            ]
+        records.append({"messages": messages})
+    data_files = {}
+    for split in ("train", "validation"):
+        data_files[split] = tmp_path / f"chat_{split}.jsonl"
+        data_files[split].write_text("\n".join(json.dumps(r) for r in records) + "\n")
+
+    distill_output_dir = tmp_path / "distill_output"
+    distill_cmd_parts = extend_cmd_parts(
+        ["torchrun", f"--nproc_per_node={num_gpus}", "distill.py", "--sft"],
+        student_hf_path=teacher_hf_path,
+        teacher_hf_path=teacher_hf_path,
+        sft_hf_dataset=data_files["train"],
+        sft_hf_validation=data_files["validation"],
+        output_dir=distill_output_dir,
+        tp_size=num_gpus,
+        pp_size=1,
+        seq_length=64,
+        mbs=1,
+        gbs=gbs,
+        train_iters=train_iters,
+        lr_warmup_iters=1,
+        eval_interval=train_iters,
+        eval_iters=1,
+        log_interval=1,
+    )
+    run_example_command(distill_cmd_parts, example_path="megatron_bridge")
+
+    assert (distill_output_dir / f"checkpoints/iter_{train_iters:07d}").exists()
+
+
+@pytest.mark.skipif(not HAS_DIRECT_HF_SFT, reason="Megatron-Bridge lacks DirectHFSFTDatasetConfig")
+@pytest.mark.parametrize(
+    ("spec", "path_or_dataset", "split", "load_kwargs"),
+    [
+        ("data/foo.jsonl", "json", "train", {"data_files": "data/foo.jsonl"}),
+        ("data/foo.json", "json", "train", {"data_files": "data/foo.json"}),
+        ("org/ds", "org/ds", "train", None),
+        ("org/ds:test", "org/ds", "test", None),
+        # Only the first ':' separates the split, so slice syntax survives.
+        ("org/ds:train[:1000]", "org/ds", "train[:1000]", None),
+    ],
+)
+def test_hf_source_spec(spec, path_or_dataset, split, load_kwargs):
+    source = _hf_source(spec)
+    assert (source.path_or_dataset, source.split, source.load_kwargs) == (
+        path_or_dataset,
+        split,
+        load_kwargs,
+    )
+
+
 def test_distill_validate_only(tmp_path, num_gpus):
     teacher_hf_path = create_tiny_qwen3_dir(tmp_path, with_tokenizer=True)
     train_iters = 2