From d14cfcab33d5b6bc212c0da049ac91641f1e5076 Mon Sep 17 00:00:00 2001 From: wiederse <1477618732@qq.com> Date: Tue, 20 Oct 2026 11:42:32 +0000 Subject: [PATCH 1/2] Skip MTP forward and recomputation when freeze_mtp is enabled Signed-off-by: wiederse <1477618732@qq.com> --- examples/megatron_bridge/distill.py | 2 +- .../NemotronH-Omni-W4A4-QAD/README.md | 2 +- modelopt/torch/distill/plugins/megatron.py | 44 +++++++++++++---- tests/_test_utils/torch/megatron/models.py | 2 +- .../distill/plugins/test_distill_megatron.py | 47 ++++++++++++++++--- 5 files changed, 77 insertions(+), 20 deletions(-) diff --git a/examples/megatron_bridge/distill.py b/examples/megatron_bridge/distill.py index 333ccbbcb3d..4809783f3f2 100644 --- a/examples/megatron_bridge/distill.py +++ b/examples/megatron_bridge/distill.py @@ -278,7 +278,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", diff --git a/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md b/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md index 461f2e201c9..2fe6539ace0 100644 --- a/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md +++ b/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md @@ -132,7 +132,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. diff --git a/modelopt/torch/distill/plugins/megatron.py b/modelopt/torch/distill/plugins/megatron.py index 26b3c9cc5c8..8f0379caee7 100644 --- a/modelopt/torch/distill/plugins/megatron.py +++ b/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) @@ -719,8 +721,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.") @@ -743,6 +749,23 @@ 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) + 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] @@ -762,14 +785,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) diff --git a/tests/_test_utils/torch/megatron/models.py b/tests/_test_utils/torch/megatron/models.py index e1ff8424926..68cbc1976f4 100644 --- a/tests/_test_utils/torch/megatron/models.py +++ b/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, diff --git a/tests/gpu_megatron/torch/distill/plugins/test_distill_megatron.py b/tests/gpu_megatron/torch/distill/plugins/test_distill_megatron.py index 2bbcc108d69..8b9daf9d56b 100644 --- a/tests/gpu_megatron/torch/distill/plugins/test_distill_megatron.py +++ b/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 @@ -333,17 +334,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) @@ -362,7 +363,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. @@ -394,6 +405,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 ) @@ -428,6 +449,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): @@ -568,9 +600,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(): From 8f166bd72ed89a33f46a2e60cc15eada1a94202d Mon Sep 17 00:00:00 2001 From: wiederse <1477618732@qq.com> Date: Wed, 28 Oct 2026 08:11:36 +0000 Subject: [PATCH 2/2] Clarify frozen MTP bypass on non-last pipeline stages Signed-off-by: wiederse <1477618732@qq.com> --- modelopt/torch/distill/plugins/megatron.py | 1 + 1 file changed, 1 insertion(+) diff --git a/modelopt/torch/distill/plugins/megatron.py b/modelopt/torch/distill/plugins/megatron.py index cf583fe1b7b..8ef1d1e1ce5 100644 --- a/modelopt/torch/distill/plugins/megatron.py +++ b/modelopt/torch/distill/plugins/megatron.py @@ -759,6 +759,7 @@ 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