[wiederse]

Superseded by #2891 (freeze-mtp-skip-forward) after GPU validation.

What does this PR do?

Type of change: Bug fix.

--freeze_mtp previously froze the student's MTP-specific parameters and skipped its cross-entropy loss, but both teacher and student still executed the MTP transformer. The student's MTP activations could also be recomputed during backward through shared backbone tensors.

Explicitly freezing MTP now bypasses its forward and loss paths for both models. The GPT wrapper temporarily disables both MTP gates using a shallow config copy, leaving the decoder's shared config intact for MoE auxiliary losses and activation recomputation. The original config and MTP flags are restored even if forward raises, and MTP parameters remain available in checkpoints and exports. The default behavior and the automatic freeze rule are unchanged.

Usage

python examples/megatron_bridge/distill.py <existing training arguments> --freeze_mtp

Testing

• Passed pre-commit on all five changed files, including Ruff, mypy, Bandit and Markdown checks.
• Passed git diff --check.
• Collected all 11 selected cases with:

  python -m pytest tests/gpu_megatron/torch/distill/plugins/test_distill_megatron.py --collect-only -q -o addopts='' -k 'freeze_mtp_with_quantized_mtp or skip_lm_loss_with_mtp'

• Extended the existing test_skip_lm_loss_with_mtp to expect no MTP loss calls when explicitly frozen.
• Extended the existing test_freeze_mtp_with_quantized_mtp to cover dense and MoE models with full activation recomputation; forward hooks reject teacher/student MTP execution, while existing backward/optimizer assertions verify that the backbone trains and frozen MTP parameters remain unchanged. Additional assertions check checkpoint inclusion and config restoration after normal and exceptional forward.
• GPU and multi-GPU execution, before/after regression execution, and representative GPU time/peak-memory measurements were not run: this environment has no CUDA devices and no active GPU-container relay. The new training-path assertions still require GPU validation.

Before your PR is "Ready for review"

• Is this change backward compatible?: ✅ The API, checkpoint layout, and behavior when the flag is omitted are unchanged.
• If you copied code from any other sources or added a new PIP dependency, did you follow guidance in CONTRIBUTING.md: N/A.
• Did you write any new necessary tests?: ✅ Extended two existing GPU tests; no new test functions or unit-test overhead.
• Did you update Changelog?: N/A. Follow-up fix for freeze_mtp in the same unreleased cycle.
• Did you get Claude approval on this PR?: ❌ Not requested.

Additional Information

Related: #2784 changes the freeze override API; this PR retains the current API and addresses execution when freeze_mtp=True.

Summary by CodeRabbit

• New Features

  • Freezing MTP now bypasses student and teacher MTP processing, loss calculation, and activation recomputation during distillation.
  • Frozen MTP weights remain available in checkpoints and exported models.
  • Behavior is supported across quantized MTP and MoE configurations.

• Documentation

  • Updated command-line help and QAD guidance to clarify the effects of freezing MTP.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

• Configuration used: Repository: NVIDIA/Model-Optimizer/.coderabbit.yaml
• Review profile: CHILL
• Plan: Enterprise
• Run ID: 75a910f5-8719-490f-b887-04da858c8c0c

📥 Commits

Reviewing files that changed from the base of the PR and between d8abfb4 and d14cfca.

📒 Files selected for processing: 5

• examples/megatron_bridge/distill.py
• examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md
• modelopt/torch/distill/plugins/megatron.py
• tests/_test_utils/torch/megatron/models.py
• tests/gpu_megatron/torch/distill/plugins/test_distill_megatron.py

Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 11 remain after this review.

📝 Walkthrough

When freeze_mtp is enabled, the distillation plugin now bypasses MTP processing during teacher and student forward passes. The change also updates the CLI help, tutorial documentation, and tests for frozen MTP behavior and related configurations.

Changes: Frozen MTP forward bypass

Layer / File(s)	Summary
Implement and describe frozen MTP bypass modelopt/torch/distill/plugins/megatron.py, examples/megatron_bridge/distill.py, examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md	The plugin temporarily disables MTP processing and sets a copied language-model configuration's MTP layer count to zero during teacher and student forward passes, then restores the original state. The CLI help and tutorial documentation describe the bypass.
Validate bypass and configurations tests/_test_utils/torch/megatron/models.py, tests/gpu_megatron/torch/distill/plugins/test_distill_megatron.py	Tests check MTP loss-call counts, confirm frozen MTP modules do not execute, and verify model state and decoder behavior. Coverage also includes recomputation, quantized MTP, and MoE configurations.

Priority: ⬇️ Low

Estimated code review effort: 3 (Moderate) | ~20 minutes

Change: Bug fix

Suggested reviewers: aanoosheh

Merge Risk: ⚪ Minimal · up to d14cf

No merge-blocking issue is established. The new training-path assertions still await GPU execution.

Pre-merge checks: 5 · 1

Failed checks: 1 warning

Check name	Status	Explanation	Resolution
Docstring Coverage		Docstring coverage is 58.33% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 12 functions across 4 files. (1 skipped: …	Write docstrings for the functions missing them to satisfy the coverage threshold.

Passed checks: 5 passed

Check name	Status	Explanation
Description Check		Check skipped - CodeRabbit’s high-level summary is enabled.
Title check		The title clearly and concisely describes the main change: skipping MTP forward and recomputation when freeze_mtp is enabled.
Linked Issues check		Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check		Check skipped because no linked issues were found for this pull request.
Security Anti-Patterns		No listed security anti-pattern is introduced by this pull request. The added modelopt and example code only adds context-manager/configuration logic and help text. No new torch.load(weights_only=Fals…

Full details: Security Anti-Patterns

Explanation

No listed security anti-pattern is introduced by this pull request. The added modelopt and example code only adds context-manager/configuration logic and help text. No new torch.load(weights_only=False), pickle-enabled NumPy load, eval/exec on external input, or # nosec comment appears in the changed Python additions. The existing trust_remote_code option is controlled by the caller and defaults to false via store_true. The pull request changes no pyproject.toml or requirements file.

Full details: Docstring Coverage

Explanation

Docstring coverage is 58.33% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 12 functions across 4 files. (1 skipped: 1 unsupported.)

Resolution

Write docstrings for the functions missing them to satisfy the coverage threshold.

• Fix all pre-merge checks with AI

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create a new PR

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

_{Comment @coderabbitai help to get the list of available commands.}