Skip MTP forward and recomputation when freeze_mtp is enabled

examples/megatron_bridge/distill.py

@@ -300,7 +300,7 @@ 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.",
+        help="Freeze the student's MTP weights and bypass teacher/student MTP forward and loss.",
     )
     parser.add_argument(
         "--moe_aux_loss_coeff",

examples/megatron_bridge/tutorials/NVIDIA-Nemotron-3.5-Super-VL-120B-A12B-BF16/README.md

@@ -135,7 +135,7 @@ python -u /opt/Model-Optimizer/examples/megatron_bridge/distill.py \
 - `--sft --sft_hf_dataset ... --sft_loss_mode assistant` renders each conversation with the model's chat template and computes the loss on assistant tokens only.
 - `--seq_length 32768 --gbs 64`: 600 iterations cover 38,400 records, about 1.07 epochs of the 35,800-record training split.
 - `--recompute_*` and `--no_async_save` are needed to fit: the async-save worker needs its own CUDA context and runs out of memory next to the resident teacher.
-- `--freeze_mtp` keeps the MTP head out of training: its loss is skipped and its parameters are frozen, so the exported MTP head keeps its PTQ weights. Without it, the quantized MTP head's cross-entropy loss trains the head and, through the hidden states, embeddings and output layer it shares, the rest of the student too.
+- `--freeze_mtp` bypasses teacher/student MTP forward, loss, and activation recomputation while preserving the frozen MTP weights in checkpoints and exports. Without it, the quantized MTP head's cross-entropy loss trains the head and, through the hidden states, embeddings and output layer it shares, the rest of the student too.
 - Each saved iteration is about 1.6 TB; size `--save_interval` and `--checkpoint_keep_last` to your storage.
 
 </details>

modelopt/torch/distill/plugins/megatron.py

@@ -22,6 +22,8 @@
 import warnings
 from abc import ABCMeta
 from collections.abc import Callable
+from contextlib import contextmanager
+from copy import copy
 from dataclasses import dataclass, field
 from types import MethodType
 from typing import TYPE_CHECKING
@@ -71,8 +73,8 @@ 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: Freeze the student's MTP parameters and bypass MTP forward/loss computation
+            for both teacher and student, retaining their checkpoint weights. Default: ``False``.
     """
 
     intermediate_layer_pairs: list[tuple[str, ...]] = field(default_factory=list)
@@ -724,8 +726,12 @@ def _sharded_state_dict(self, *args, **kwargs) -> "ShardedStateDict":
                     num_frozen += 1
         if num_frozen:
             warn_rank_0(
-                f"MTP head is frozen ({num_frozen} params) and its loss is skipped"
-                + (" (freeze_mtp)." if distill_cfg.freeze_mtp else ": it is outside quantization.")
+                f"MTP head is frozen ({num_frozen} params)"
+                + (
+                    " and its forward/loss are skipped (freeze_mtp)."
+                    if distill_cfg.freeze_mtp
+                    else " and its loss is skipped: it is outside quantization."
+                )
             )
         elif parallel_state.is_pipeline_last_stage():  # earlier stages never hold the MTP head
             warn_rank_0("freeze_mtp is set but the model has no MTP head: nothing to freeze.")
@@ -748,6 +754,24 @@ def _compute_teacher_lm_loss(self, labels, logits) -> Tensor:
         _compute_teacher_lm_loss, model.teacher_model
     )
 
+    @contextmanager
+    def _skip_mtp_forward(lm):
+        mtp_process = getattr(lm, "mtp_process", False)
+        config = lm.config
+        mtp_num_layers = getattr(config, "mtp_num_layers", None)
+        # Non-last PP stages retain mtp_num_layers; entering the bypass there is harmless.
+        if not distill_cfg.freeze_mtp or not (mtp_process or mtp_num_layers):
+            yield
+            return
+        try:
+            # GPT gates the block and loss separately. Do not alter the decoder's shared config.
+            lm.mtp_process = False
+            lm.config = copy(config)
+            lm.config.mtp_num_layers = 0
+            yield
+        finally:
+            lm.mtp_process, lm.config = mtp_process, config
+
     # HACK: Pipeline-parallel Distillation requires splitting input tensor into student and teacher parts.
     def _set_student_input_tensor_shape(self, shapes: list[tuple[int]]):
         self._tensor_split_idx = shapes[0][-1]
@@ -767,14 +791,15 @@ def _set_input_tensor(self, input_tensors: list[Tensor]):
 
     # HACK: Concatenate output tensors when PP>1 so they can be passed between ranks.
     def _forward(self, *args, **kwargs):
-        with torch.no_grad():
-            self._teacher_model.eval()
-            teacher_output = self._teacher_model(*args, **kwargs)
+        with _skip_mtp_forward(self), _skip_mtp_forward(self._teacher_model):
+            with torch.no_grad():
+                self._teacher_model.eval()
+                teacher_output = self._teacher_model(*args, **kwargs)
 
-        self._lm_loss_call_count = 0  # reset loss-fn call counter for MTP tracking
+            self._lm_loss_call_count = 0  # reset loss-fn call counter for MTP tracking
 
-        with self.only_student_forward():
-            student_output = type(self).forward(self, *args, **kwargs)
+            with self.only_student_forward():
+                student_output = type(self).forward(self, *args, **kwargs)
 
         if not parallel_state.is_pipeline_last_stage():
             return torch.cat([student_output, teacher_output], dim=-1)

tests/_test_utils/torch/megatron/models.py

@@ -215,7 +215,7 @@ def squared_relu(x):
         pipeline_model_parallel_size=pipeline_model_parallel_size,
         expert_model_parallel_size=expert_model_parallel_size,
         expert_tensor_parallel_size=expert_tensor_parallel_size,
-        sequence_parallel=False,
+        sequence_parallel=config_kwargs.pop("sequence_parallel", False),
         num_layers=num_layers,
         num_layers_in_first_pipeline_stage=num_layers_in_first_pipeline_stage,
         num_layers_in_last_pipeline_stage=num_layers_in_last_pipeline_stage,

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

@@ -17,6 +17,7 @@
 import warnings
 from functools import partial
 from types import SimpleNamespace
+from unittest.mock import patch
 
 import pytest
 import torch
@@ -334,17 +335,17 @@ def _recording_loss(labels, logits):
 
     distillation_model(prompt_tokens, position_ids, attention_mask, labels=labels)
 
-    # Expect mtp_num_layers + 1 total calls: first mtp_num_layers are MTP heads,
-    # the last one is the main LM head.
-    assert len(recorded_losses) == mtp_num_layers + 1, (
-        f"Expected {mtp_num_layers + 1} loss calls, got {len(recorded_losses)}"
+    # Frozen MTP does not enter the forward/loss path; otherwise its heads precede the main LM.
+    expected_calls = 1 if freeze_mtp else mtp_num_layers + 1
+    assert len(recorded_losses) == expected_calls, (
+        f"Expected {expected_calls} loss calls, got {len(recorded_losses)}"
     )
     for i, loss in enumerate(recorded_losses[:-1]):
         assert loss.any() != mtp_skipped, f"MTP head {i} loss should be zero iff MTP is skipped"
     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):
+def _test_freeze_mtp_with_quantized_mtp(freeze_mtp, kd_loss_alpha, moe, rank, size):
     """With the MTP head quantized too, only ``freeze_mtp`` keeps it out of a training step."""
     set_seed(SEED)
 
@@ -363,7 +364,17 @@ def _test_freeze_mtp_with_quantized_mtp(freeze_mtp, kd_loss_alpha, rank, size):
         "vocab_size": vocab_size,
         "activation_func": "squared_relu",
         "mtp_num_layers": 1,
+        "recompute_granularity": "full",
+        "recompute_method": "uniform",
+        "recompute_num_layers": 1,
     }
+    if moe:
+        model_kwargs.update(
+            sequence_parallel=size > 1,
+            num_moe_experts=4,
+            moe_router_load_balancing_type="global_aux_loss",
+            moe_aux_loss_coeff=1e-3,
+        )
     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.
@@ -395,6 +406,16 @@ def _test_freeze_mtp_with_quantized_mtp(freeze_mtp, kd_loss_alpha, rank, size):
     body_params = {n: p for n, p in named_params.items() if n.startswith("decoder.")}
     assert mtp_params and body_params
     mtp_before = {n: p.detach().clone() for n, p in mtp_params.items()}
+
+    def reject_frozen_mtp(module, args):
+        assert not freeze_mtp, "Frozen MTP must not execute forward or recomputation"
+
+    # Checkpoint recomputation can call the inner transformer layer directly.
+    handles = [
+        module.register_forward_pre_hook(reject_frozen_mtp)
+        for lm in (distillation_model, distillation_model.teacher_model)
+        for module in lm.mtp.modules()
+    ]
     optimizer = torch.optim.AdamW(
         [p for p in named_params.values() if p.requires_grad], lr=1e-2, weight_decay=0.1
     )
@@ -429,6 +450,17 @@ def _test_freeze_mtp_with_quantized_mtp(freeze_mtp, kd_loss_alpha, rank, size):
     assert mtp_trained != freeze_mtp
     if freeze_mtp:
         assert all(torch.equal(p, mtp_before[n]) for n, p in mtp_params.items())
+        assert any(name.startswith("mtp.") for name in distillation_model.sharded_state_dict())
+        with (
+            patch.object(distillation_model.decoder, "forward", side_effect=RuntimeError("probe")),
+            pytest.raises(RuntimeError, match="probe"),
+        ):
+            distillation_model(prompt_tokens, position_ids, attention_mask, labels=labels)
+    for lm in (distillation_model, distillation_model.teacher_model):
+        assert lm.mtp_process and lm.config.mtp_num_layers == 1
+        assert lm.config is lm.decoder.config
+    for handle in handles:
+        handle.remove()
 
 
 def test_logits_kl_loss(dist_workers):
@@ -607,9 +639,10 @@ def test_skip_lm_loss_with_mtp(dist_workers, 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):
+@pytest.mark.parametrize("moe", [False, True])
+def test_freeze_mtp_with_quantized_mtp(dist_workers, freeze_mtp, kd_loss_alpha, moe):
     """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))
+    dist_workers.run(partial(_test_freeze_mtp_with_quantized_mtp, freeze_mtp, kd_loss_alpha, moe))
 
 
 def test_mtp_excluded_from_quantization():