Make freeze_mtp an explicit override of the automatic MTP freeze

CHANGELOG.rst

@@ -162,7 +162,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 ``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 ``freeze_mtp`` to the Megatron ``DistillationConfig`` (``--freeze_mtp`` / ``--no-freeze_mtp`` in ``examples/megatron_bridge/distill.py``): ``True`` freezes the student's MTP-specific parameters and disables its MTP loss, ``False`` keeps them training. Unset, they are frozen only when ``kd_loss_alpha`` is 1.0 and the student is quantized except for its MTP head, as before.
 - 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 <https://github.com/NVIDIA/Model-Optimizer/tree/main/examples/dataset/MEGATRON_DATA_PREP.md#chat-jsonl-blends-for-assistant-only-loss>`_.
 - 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.

examples/megatron_bridge/distill.py

@@ -302,8 +302,11 @@ def get_args():
     )
     parser.add_argument(
         "--freeze_mtp",
-        action="store_true",
-        help="Keep the student's MTP head at its input weights: skip its loss and freeze it.",
+        action=argparse.BooleanOptionalAction,
+        default=None,
+        help="Freeze the student's MTP-specific parameters and disable its MTP loss "
+        "(--no-freeze_mtp keeps them training). Unset, they are frozen only when "
+        "--kd_loss_alpha is 1.0 and the student is quantized except for its MTP head.",
     )
     parser.add_argument(
         "--moe_aux_loss_coeff",

modelopt/torch/distill/plugins/megatron.py

@@ -71,8 +71,11 @@ class DistillationConfig:
             Only the smallest prefix of the (sorted) Top-K whose cumulative teacher probability
             reaches this value contributes to the loss. Requires ``logit_kl_topk``. Must be in (0, 1].
         logit_kl_top_p_min_k: Minimum number of Top-K entries kept per token when top-P is active.
-        freeze_mtp: Keep the student's MTP head fixed: skip its loss, which otherwise also trains the
-            shared layers it reads from, and freeze its parameters. Default: ``False``.
+        freeze_mtp: Whether to freeze the student's MTP-specific parameters and disable its auxiliary
+            MTP loss, which otherwise also trains the shared layers it reads from. ``True`` freezes
+            it and ``False`` keeps it training. ``None`` (default) freezes it only when
+            ``kd_loss_alpha == 1.0`` and the student is quantized except for its MTP head, since
+            there is then no quantization error for the MTP loss to recover.
     """
 
     intermediate_layer_pairs: list[tuple[str, ...]] = field(default_factory=list)
@@ -82,7 +85,7 @@ class DistillationConfig:
     logit_kl_topk: int | None = None
     logit_kl_top_p: float | None = None
     logit_kl_top_p_min_k: int = 1
-    freeze_mtp: bool = False
+    freeze_mtp: bool | None = None
     criterion: Criterion | None = None
     loss_balancer: mtd.DistillationLossBalancer | None = None
 
@@ -743,20 +746,16 @@ def _init_weights(self, module):
 ########################################################
 
 
-def _mtp_excluded_from_quantization(model: nn.Module) -> bool:
-    """Whether the model is quantized but its MTP submodules are not.
-
-    Only then is the MTP loss pure overhead, since distillation has no quantization error to
-    recover there. Plain distillation (e.g. pruning recovery) still trains the MTP head.
-    """
-    has_mtp = model_quantized = mtp_quantized = False
+def _has_quantized_backbone_with_unquantized_mtp(model: nn.Module) -> bool:
+    """Whether ``model`` has MTP submodules and enabled quantizers, none of them inside MTP."""
+    has_mtp = backbone_quantized = mtp_quantized = False
     for name, module in model.named_modules():
         is_mtp = "mtp" in name.split(".")
         has_mtp |= is_mtp
         if isinstance(module, TensorQuantizer) and module.is_enabled:
-            model_quantized = True
             mtp_quantized |= is_mtp
-    return has_mtp and model_quantized and not mtp_quantized
+            backbone_quantized |= not is_mtp
+    return has_mtp and backbone_quantized and not mtp_quantized
 
 
 def adjust_distillation_model_for_mcore(
@@ -775,13 +774,15 @@ def _sharded_state_dict(self, *args, **kwargs) -> "ShardedStateDict":
     # Uses a per-forward call counter so that MTP head calls (which always precede the
     # main LM head call in _postprocess) still receive real CE loss even when
     # skip_lm_loss=True — only the final main-head call is zeroed.
-    # An MTP head left out of quantization is exempt from that: there is no quantization
-    # error to recover there, and its CE materialises an fp32 [seq, vocab] tensor.
-    # ``freeze_mtp`` skips the MTP loss unconditionally.
+    # ``freeze_mtp`` decides whether the MTP head trains. Unset, it is frozen only when the LM loss
+    # is skipped and the student is quantized except for its MTP head: there is then no
+    # quantization error to recover there, and its CE materialises an fp32 [seq, vocab] tensor.
     skip_lm_loss = distill_cfg.skip_lm_loss
-    skip_mtp_loss = distill_cfg.freeze_mtp or (
-        skip_lm_loss and _mtp_excluded_from_quantization(model)
-    )
+    if distill_cfg.freeze_mtp is None:
+        with model.hide_teacher_model():
+            skip_mtp_loss = skip_lm_loss and _has_quantized_backbone_with_unquantized_mtp(model)
+    else:
+        skip_mtp_loss = distill_cfg.freeze_mtp
     if skip_mtp_loss:
         # Freeze the untrained MTP head: DDP's overlapped grad reduce asserts on params with no grad.
         num_frozen = 0

tests/gpu_megatron/torch/distill/plugins/test_distill_megatron.py

@@ -36,7 +36,7 @@
     LogitsKLLoss,
     TopKLogitsKLLoss,
     TopLogitsKLLoss,
-    _mtp_excluded_from_quantization,
+    _has_quantized_backbone_with_unquantized_mtp,
     adjust_distillation_model_for_mcore,
     setup_distillation_config,
 )
@@ -240,7 +240,8 @@ def _test_topk_logits_kl_loss(kd_kwargs, rank, size):
 def _test_skip_lm_loss_with_mtp(quantized, skip_lm_loss, freeze_mtp, rank, size):
     """Test that skip_lm_loss only zeroes the main LM head, unless MTP is left out of quantization.
 
-    ``freeze_mtp`` zeroes the MTP loss and freezes the MTP head regardless.
+    An explicit ``freeze_mtp`` overrides that: ``True`` zeroes the MTP loss and freezes the MTP head,
+    ``False`` keeps both.
     """
     set_seed(SEED)
 
@@ -284,7 +285,7 @@ def _test_skip_lm_loss_with_mtp(quantized, skip_lm_loss, freeze_mtp, rank, size)
         mtp_num_layers=mtp_num_layers,
     ).cuda()
     if quantized:
-        # A quantized body with the MTP head left out of quantization.
+        # Quantization-state fixture: a quantized body with the MTP head left out of quantization.
         student_model.decoder.weight_quantizer = TensorQuantizer()
 
     distill_cfg = setup_distillation_config(
@@ -304,7 +305,7 @@ def _test_skip_lm_loss_with_mtp(quantized, skip_lm_loss, freeze_mtp, rank, size)
     adjust_distillation_model_for_mcore(distillation_model, distill_cfg)
 
     # An untrained MTP head is frozen so DDP doesn't wait for its grads.
-    mtp_skipped = freeze_mtp or (quantized and skip_lm_loss)
+    mtp_skipped = (quantized and skip_lm_loss) if freeze_mtp is None else freeze_mtp
     with distillation_model.hide_teacher_model():
         mtp_params = [p for n, p in distillation_model.named_parameters() if "mtp" in n.split(".")]
     assert mtp_params
@@ -346,8 +347,8 @@ def _recording_loss(labels, logits):
     assert recorded_losses[-1].any() != skip_lm_loss, "Main LM head loss is zero iff skip_lm_loss"
 
 
-def _test_freeze_mtp_with_quantized_mtp(freeze_mtp, kd_loss_alpha, rank, size):
-    """With the MTP head quantized too, only ``freeze_mtp`` keeps it out of a training step."""
+def _test_freeze_mtp_training_step(mtp_quantized, freeze_mtp, kd_loss_alpha, rank, size):
+    """Through a training step, the MTP head is frozen exactly when ``freeze_mtp`` resolves to it."""
     set_seed(SEED)
 
     vocab_size = 32
@@ -368,10 +369,12 @@ def _test_freeze_mtp_with_quantized_mtp(freeze_mtp, kd_loss_alpha, rank, size):
     }
     teacher_model = get_mcore_gpt_model(initialize_megatron=True, **model_kwargs).cuda()
     student_model = get_mcore_gpt_model(initialize_megatron=False, **model_kwargs).cuda()
-    # A recipe that quantizes the MTP head as well as the body.
+    # Quantization-state fixtures: a quantized body, with or without a quantized MTP head.
     student_model.decoder.weight_quantizer = TensorQuantizer()
-    student_model.mtp.weight_quantizer = TensorQuantizer()
-    assert not _mtp_excluded_from_quantization(student_model)
+    if mtp_quantized:
+        student_model.mtp.weight_quantizer = TensorQuantizer()
+    assert _has_quantized_backbone_with_unquantized_mtp(student_model) != mtp_quantized
+    frozen = (kd_loss_alpha == 1.0 and not mtp_quantized) if freeze_mtp is None else freeze_mtp
 
     distill_cfg = setup_distillation_config(
         config_or_path=DistillationConfig(kd_loss_alpha=kd_loss_alpha, freeze_mtp=freeze_mtp),
@@ -389,7 +392,7 @@ def _test_freeze_mtp_with_quantized_mtp(freeze_mtp, kd_loss_alpha, rank, size):
         adjust_distillation_model_for_mcore(distillation_model, distill_cfg)
     if rank == 0:
         frozen_msg = [w for w in caught if "MTP head is frozen" in str(w.message)]
-        assert bool(frozen_msg) == freeze_mtp
+        assert bool(frozen_msg) == frozen
 
     with distillation_model.hide_teacher_model():
         named_params = dict(distillation_model.named_parameters())
@@ -428,8 +431,8 @@ def _test_freeze_mtp_with_quantized_mtp(freeze_mtp, kd_loss_alpha, rank, size):
     assert any(p.grad is not None and p.grad.any() for p in body_params.values())
     # The MTP head trains unless frozen; frozen, it gets no gradient and the step leaves it as is.
     mtp_trained = any(p.grad is not None and p.grad.any() for p in mtp_params.values())
-    assert mtp_trained != freeze_mtp
-    if freeze_mtp:
+    assert mtp_trained != frozen
+    if frozen:
         assert all(torch.equal(p, mtp_before[n]) for n, p in mtp_params.items())
 
 
@@ -639,26 +642,40 @@ def test_distillation_config_top_p_validation():
 @pytest.mark.parametrize(
     ("quantized", "skip_lm_loss", "freeze_mtp"),
     [
-        (False, True, False),
-        (True, True, False),
-        (True, False, False),
+        (False, True, None),
+        (True, True, None),
+        (True, False, None),
         (False, True, True),
         (False, False, True),
+        (True, True, False),  # explicit False overrides the automatic freeze
     ],
 )
 def test_skip_lm_loss_with_mtp(dist_workers, quantized, skip_lm_loss, freeze_mtp):
     """Test that skip_lm_loss only zeroes the main LM head, unless MTP is left out of quantization."""
     dist_workers.run(partial(_test_skip_lm_loss_with_mtp, quantized, skip_lm_loss, freeze_mtp))
 
 
-@pytest.mark.parametrize(("freeze_mtp", "kd_loss_alpha"), [(True, 1.0), (True, 0.5), (False, 1.0)])
-def test_freeze_mtp_with_quantized_mtp(dist_workers, freeze_mtp, kd_loss_alpha):
-    """With the MTP head quantized too, only ``freeze_mtp`` keeps it out of a training step."""
-    dist_workers.run(partial(_test_freeze_mtp_with_quantized_mtp, freeze_mtp, kd_loss_alpha))
+@pytest.mark.parametrize(
+    ("mtp_quantized", "freeze_mtp", "kd_loss_alpha"),
+    [
+        (True, None, 1.0),  # automatic: a quantized MTP head keeps training
+        (False, None, 1.0),  # automatic: an unquantized MTP head is frozen
+        (False, False, 1.0),  # explicit False overrides the automatic freeze
+        (True, True, 1.0),
+        (True, True, 0.5),
+        (False, True, 0.5),
+        (True, False, 1.0),
+    ],
+)
+def test_freeze_mtp_training_step(dist_workers, mtp_quantized, freeze_mtp, kd_loss_alpha):
+    """Through a training step, the MTP head is frozen exactly when ``freeze_mtp`` resolves to it."""
+    dist_workers.run(
+        partial(_test_freeze_mtp_training_step, mtp_quantized, freeze_mtp, kd_loss_alpha)
+    )
 
 
-def test_mtp_excluded_from_quantization():
-    """MTP loss is skipped only when the model is quantized and MTP is left out of it."""
+def test_has_quantized_backbone_with_unquantized_mtp():
+    """True only for a model with MTP and enabled quantizers, none of them inside MTP."""
 
     def _quantizer(enabled: bool) -> TensorQuantizer:
         quantizer = TensorQuantizer()
@@ -676,18 +693,20 @@ def _model(*, with_mtp: bool, body_quant: bool, mtp_quant: bool) -> nn.Module:
             model.mtp.weight_quantizer = _quantizer(mtp_quant)
         return model
 
-    # Plain distillation (e.g. pruning recovery) still trains the MTP head.
-    assert not _mtp_excluded_from_quantization(
+    # Nothing quantized.
+    assert not _has_quantized_backbone_with_unquantized_mtp(
         _model(with_mtp=True, body_quant=False, mtp_quant=False)
     )
-    # QAD with MTP excluded from the recipe: no quantization error to recover there.
-    assert _mtp_excluded_from_quantization(_model(with_mtp=True, body_quant=True, mtp_quant=False))
-    # QAD with a quantized MTP head: keep its loss.
-    assert not _mtp_excluded_from_quantization(
+    # Quantized backbone, MTP left out of quantization.
+    assert _has_quantized_backbone_with_unquantized_mtp(
+        _model(with_mtp=True, body_quant=True, mtp_quant=False)
+    )
+    # MTP quantized too.
+    assert not _has_quantized_backbone_with_unquantized_mtp(
         _model(with_mtp=True, body_quant=True, mtp_quant=True)
     )
     # No MTP at all.
-    assert not _mtp_excluded_from_quantization(
+    assert not _has_quantized_backbone_with_unquantized_mtp(
         _model(with_mtp=False, body_quant=True, mtp_quant=False)
     )