feat(megatron-bridge): SFT-masked data support in distillation

CHANGELOG.rst

@@ -12,6 +12,7 @@ Changelog
 
 *Megatron Framework (M-LM / M-Bridge)*
 
+- Add SFT-masked data support to ``examples/megatron_bridge/distill.py``: ``--sft --sft_dataset_root <dir>`` distills on raw prompt-completion JSONL (``{"input", "output"}`` records) with the loss masked to the response tokens, using Megatron-Bridge's ``FinetuningDatasetConfig`` and the model's own HuggingFace tokenizer instead of the pretraining ``GPTDataset`` and ``NullTokenizer``.
 - Add per-expert weight quantization for Transformer Engine ``TEGroupedMLP`` (fused MoE experts): each expert now has its own ``weight_quantizer`` (a ``GroupedQuantizer`` holding one ``TensorQuantizer`` per expert) with an independent ``amax``, instead of a single shared ``amax`` across all experts. Applies to ``mtq.quantize`` calibration, HF / Megatron export, and QAD.
 - Add opt-in ``torch.compile`` execution for Transformer Engine grouped-linear per-expert weight quantizers while preserving their native checkpoint amax shapes. Set ``MODELOPT_TEGROUPED_COMPILE_WEIGHT_LOOP=1`` before quantized-module conversion; the default path remains eager.
 

examples/megatron_bridge/README.md

@@ -130,6 +130,21 @@ The distillation script expects pre-tokenized data in Megatron's binary format (
 See the **[Dataset Preparation README](../dataset/README.md#tokenizing-for-megatron-frameworks)**
 for full instructions on tokenizing JSONL files and Hugging Face datasets and get the list of output prefixes that you can use for `--data_paths` argument.
 
+Alternatively, pass `--sft --sft_dataset_root <dir>` to distill on **raw prompt-completion JSONL**
+with the loss masked to the completion. The directory must hold `training.jsonl` (and
+`validation.jsonl` when `--eval_iters > 0`) of `{"input": <prompt>, "output": <response>}` records, which are tokenized with
+the model's own HuggingFace tokenizer. Both fields are tokenized **as written**, except that
+leading and trailing spaces on each field are stripped — no chat template is applied. So if your
+model expects role/turn markers, include them in the `"input"` field yourself, and express any
+significant separator as a newline rather than a trailing space. A BOS token is prepended
+automatically when the tokenizer prepends one at inference, so do not add it yourself; an EOS
+token is appended after the response. A record longer than `--seq_length` is truncated from the
+**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.
+
+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.
+
 ### Distillation with Real Data
 
 Example usage to distill a 4B student (HF) from an 8B teacher (HF) on 8 GPUs (TP=8, PP=1):

examples/megatron_bridge/distill.py

@@ -34,6 +34,7 @@
 from megatron.bridge.training.config import (
     CheckpointConfig,
     ConfigContainer,
+    FinetuningDatasetConfig,
     GPTDatasetConfig,
     LoggerConfig,
     MockGPTDatasetConfig,
@@ -45,10 +46,11 @@
 from megatron.bridge.training.distill import distill
 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
 from megatron.core.datasets.utils import get_blend_from_list
 from megatron.core.distributed import DistributedDataParallelConfig
 from megatron.core.utils import unwrap_model
-from transformers import AutoConfig
+from transformers import AutoConfig, AutoTokenizer
 
 import modelopt.torch.distill as mtd
 import modelopt.torch.utils.distributed as dist
@@ -125,6 +127,20 @@ def get_args():
     parser.add_argument(
         "--use_mock_data", action="store_true", help="Use mock data instead of --data_paths"
     )
+    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.",
+    )
+    parser.add_argument(
+        "--sft_dataset_root",
+        type=str,
+        default=None,
+        help="Directory holding training.jsonl (and validation.jsonl when --eval_iters > 0) of "
+        '{"input": <prompt>, "output": <response>} records (used with --sft). See the README for '
+        "how the fields are tokenized and truncated.",
+    )
     # Training & Eval arguments
     parser.add_argument(
         "--output_dir", type=str, required=True, help="Folder for logging and checkpoint saving"
@@ -246,7 +262,7 @@ def get_args():
     args = parser.parse_args()
 
     # Sanity checks
-    if not args.use_mock_data and not args.data_paths:
+    if not args.sft and not args.use_mock_data and not args.data_paths:
         raise ValueError("Must provide either --data_paths or set --use_mock_data.")
 
     if args.student_hf_model is None:
@@ -256,11 +272,59 @@ 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:
+        raise ValueError(
+            "--sft requires --sft_dataset_root (a directory with training.jsonl, plus "
+            "validation.jsonl when --eval_iters > 0)."
+        )
+    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:
+        # 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))]
+        if absent:
+            raise ValueError(f"--sft_dataset_root {args.sft_dataset_root} is missing: {absent}.")
+        # Decided once here so it reaches print_args and costs a single tokenizer load.
+        args.sft_add_bos = _tokenizer_prepends_bos(args)
+
+    _check_shared_vocabulary(args)
+
     print_args(args)
 
     return args
 
 
+def _check_shared_vocabulary(args) -> None:
+    """Raise unless teacher and student use the same tokenizer."""
+    _tok = {"trust_remote_code": args.trust_remote_code}
+    student_vocab = AutoTokenizer.from_pretrained(args.student_hf_path, **_tok).get_vocab()
+    teacher_vocab = AutoTokenizer.from_pretrained(args.teacher_hf_path, **_tok).get_vocab()
+    if student_vocab != teacher_vocab:
+        raise ValueError(
+            "Distillation scores the teacher on the student's token ids, so teacher and student "
+            "must use the same tokenizer."
+        )
+
+
+def _tokenizer_prepends_bos(args) -> bool:
+    """True when the student tokenizer prepends a BOS at inference.
+
+    Probes an encode: fast tokenizers prepend via a post-processor that exposes no attribute.
+    """
+    tokenizer = AutoTokenizer.from_pretrained(
+        args.student_hf_path, trust_remote_code=args.trust_remote_code
+    )
+    if not getattr(tokenizer, "bos_token", None):
+        return False
+    return tokenizer("x").input_ids[:1] == [tokenizer.bos_token_id]
+
+
 def main(args: argparse.Namespace):
     checkpoint_dir = os.path.join(args.output_dir, "checkpoints")
     tensorboard_dir = os.path.join(args.output_dir, "tb_logs")
@@ -279,6 +343,10 @@ def _build_model_provider(hf_path, load_weights=True):
         provider.expert_model_parallel_size = args.ep_size
         provider.expert_tensor_parallel_size = 1  # Expert tensor parallelism is not supported
         provider.seq_length = args.seq_length
+        if args.sft:
+            # A response-only loss mask needs per-token reduction to combine across CP ranks.
+            # Must stay in sync with ``average_in_collective=not args.sft`` on the DDP config.
+            provider.calculate_per_token_loss = True
         if args.recompute_granularity is not None:
             provider.recompute_granularity = args.recompute_granularity
             provider.recompute_method = args.recompute_method
@@ -302,14 +370,26 @@ def _build_model_provider(hf_path, load_weights=True):
         student_provider.gradient_accumulation_fusion = False
     teacher_provider = _build_model_provider(args.teacher_hf_path)
 
+    # The KD losses compare logits elementwise over the vocab dim, so both output layers must have
+    # the same padded width. A shared tokenizer does not imply it: the HF configs can disagree.
+    padded = {
+        name: calculate_padded_vocab_size(
+            p.vocab_size, p.make_vocab_size_divisible_by, p.tensor_model_parallel_size
+        )
+        for name, p in (("student", student_provider), ("teacher", teacher_provider))
+    }
+    if padded["student"] != padded["teacher"]:
+        raise ValueError(
+            "Distillation needs student and teacher logits of equal width, but their padded vocab "
+            f"sizes differ ({padded['student']} vs {padded['teacher']})."
+        )
+
     kd_config = ModelOptDistillConfig(
         skip_lm_loss=not args.no_skip_lm_loss, kd_loss_scale=args.kd_loss_scale
     )
 
-    # VLM detection convention: HF VLM configs expose a ``vision_config``, and Megatron-Bridge nests
-    # the text model under the ``language_model`` submodule (used as ``distill_submodule`` below). If a
-    # future model breaks either convention, the ``getattr(model, "language_model")`` in the provider
-    # will error loudly rather than silently distilling the wrong module.
+    # HF VLM configs expose ``vision_config``; Megatron-Bridge nests the text model under
+    # ``language_model`` (used as ``distill_submodule`` below).
     is_vlm = hasattr(
         AutoConfig.from_pretrained(args.student_hf_path, trust_remote_code=args.trust_remote_code),
         "vision_config",
@@ -368,7 +448,33 @@ def _restore_student_hook(model_chunks):
         "dataloader_type": "single",
         "skip_getting_attention_mask_from_dataset": True,
     }
-    if args.use_mock_data:
+    if 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.
+        dataset_config = FinetuningDatasetConfig(
+            seq_length=args.seq_length,
+            dataset_root=args.sft_dataset_root,
+            seed=args.seed,
+            dataloader_type="batch",
+            # Honour --eval_iters 0 so a training-only dataset_root does not have to carry a
+            # dummy validation.jsonl just to satisfy the builder.
+            do_validation=args.eval_iters > 0,
+            do_test=False,
+            dataset_kwargs={
+                "prompt_template": "{input}{output}",
+                "label_key": "output",
+                "truncation_field": "input",
+                # Drop the oldest context. The default "right" would cut the prompt/answer
+                # boundary and then "output" itself, the only span the loss is computed on.
+                "truncation_method": "left",
+                "answer_only_loss": True,
+                # Prepended after truncation, so it survives a record that had to be cut.
+                "add_bos": args.sft_add_bos,
+                "add_eos": True,
+            },
+        )
+    elif args.use_mock_data:
         dataset_config = MockGPTDatasetConfig(**dataset_kwargs)
     else:
         # Convert flat CLI list (e.g. ["1.0", "/path/data"]) to Megatron blend format
@@ -399,7 +505,7 @@ def _restore_student_hook(model_chunks):
             grad_reduce_in_fp32=True,
             overlap_grad_reduce=True,
             overlap_param_gather=True,
-            average_in_collective=True,
+            average_in_collective=not args.sft,  # per-token loss must not be pre-averaged
             use_distributed_optimizer=True,
         ),
         dataset=dataset_config,
@@ -412,8 +518,24 @@ def _restore_student_hook(model_chunks):
             wandb_entity=args.wandb_entity,  # optional
             wandb_exp_name=args.wandb_exp_name,
         ),
-        tokenizer=TokenizerConfig(
-            tokenizer_type="NullTokenizer", vocab_size=distill_provider.vocab_size
+        tokenizer=(
+            # SFT reads raw text, so it needs the model's real tokenizer; the pretraining path
+            # consumes pre-tokenized data and keeps NullTokenizer.
+            TokenizerConfig(
+                tokenizer_type="HuggingFaceTokenizer",
+                tokenizer_model=args.student_hf_path,
+                hf_tokenizer_kwargs={
+                    "trust_remote_code": args.trust_remote_code,
+                    # Default True would make text_to_ids inject a BOS at the answer boundary,
+                    # since "{input}" and "{output}" are tokenized separately. Consumed by Bridge
+                    # in training/tokenizers/config.py.
+                    "include_special_tokens": False,
+                },
+            )
+            if args.sft
+            else TokenizerConfig(
+                tokenizer_type="NullTokenizer", vocab_size=distill_provider.vocab_size
+            )
         ),
         checkpoint=CheckpointConfig(
             save_interval=(

tests/examples/megatron_bridge/test_distill.py

@@ -14,6 +14,7 @@
 # limitations under the License.
 """Tests for prune_minitron.py and distill.py scripts."""
 
+import json
 from pathlib import Path
 
 import pytest
@@ -58,6 +59,43 @@ def test_distill_llm(tmp_path, num_gpus):
     assert (distilled_hf_path / "config.json").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)
+    train_iters = 2
+    gbs = 4
+    dataset_root = tmp_path / "sft_data"
+    dataset_root.mkdir()
+    # More records than train_iters * gbs so the sampler does not run dry.
+    records = [{"input": f"Q: what follows {i}?\nA:", "output": f" {i + 1}"} for i in range(64)]
+    for split in ("training", "validation"):
+        (dataset_root / f"{split}.jsonl").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_dataset_root=dataset_root,
+        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()
+
+
 def test_distill_validate_only(tmp_path, num_gpus):
     teacher_hf_path = create_tiny_qwen3_dir(tmp_path, with_tokenizer=True)
     train_iters = 2