[wiederse]

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

Add --freeze_mtp to an existing examples/megatron_bridge/distill.py command.

Testing

Validated the exact signed commit d14cfcab33d5b6bc212c0da049ac91641f1e5076 on Volc A800-SXM4-80GB GPUs, using the existing NeMo 26.08 workspace image (PyTorch 2.13.0a0+8145d630e8.nv26.06, CUDA 13.3). CUDA extensions were compiled for sm_80.

Run	Result
Full Megatron distillation module, TP2	26 passed; no failures, errors or skips
MTP regressions, one GPU	11 passed; no failures, errors or skips
Same new freeze checks with the pre-fix plugin, TP2	2 expected failures; dense and MoE both report Frozen MTP must not execute forward or recomputation

CUDA_VISIBLE_DEVICES=0,1 python -m pytest \
  tests/gpu_megatron/torch/distill/plugins/test_distill_megatron.py \
  -q -o addopts='' --durations=15

CUDA_VISIBLE_DEVICES=0 python -m pytest \
  tests/gpu_megatron/torch/distill/plugins/test_distill_megatron.py \
  -k 'skip_lm_loss_with_mtp or freeze_mtp_with_quantized_mtp' \
  -q -o addopts='' --durations=15

For the negative control, only modelopt/torch/distill/plugins/megatron.py was restored from parent commit d8abfb447d; the new tests and test helper remained unchanged. The selected cases were freeze_mtp_with_quantized_mtp and True-1.0.

The extended existing training-step test covers dense/MoE models with full activation recomputation, teacher/student MTP execution guards, backbone gradients, unchanged frozen MTP weights after AdamW, MTP checkpoint keys, and config restoration after normal and exceptional forward. The quantizers in this test mark quantization state; it does not perform numerical PTQ or a full-model QAD benchmark.

JUnit reports confirm the counts above. Task t-20261020201232-kq9xt completed with exit code 0, and source hashes matched the signed commit before and after execution. The two positive runs took 273.60 s and 33.65 s, and the negative run took 258.16 s; these include fixture setup and CUDA compilation and are not performance benchmarks.

Pre-commit on all five changed files (including Ruff, mypy, Bandit and Markdown checks) and git diff --check also passed.

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

• Improvements
  • The --freeze_mtp option now bypasses MTP forward passes and loss computation for both teacher and student models, including activation recomputation.
  • Frozen MTP weights remain available in checkpoints and exports.
  • Updated option descriptions and tutorial guidance to clarify what is bypassed and which weights are retained.

[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

• Source mode: Auto
• Repository: NVIDIA/Model-Optimizer
• Branch: main
• File: .coderabbit.yaml
• Commit: 580ffff
• Layers: Repository YAML
• Review profile: CHILL
• Plan: Enterprise
• Run ID: 851126d1-8d9d-4ae9-b09b-ff4d1779c790

Commits

Reviewing files that changed from the base of the PR and between a49a26d and 8f166bd.

Files selected for processing: 1

• modelopt/torch/distill/plugins/megatron.py

Files skipped from review as they are similar to previous changes: 1

• modelopt/torch/distill/plugins/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, Megatron distillation temporarily bypasses MTP forward and loss computation for teacher and student models. The change restores model state after each forward and adds tests for recomputation, retained MTP state, and optional MoE configurations.

Changes: Frozen MTP execution

Layer / File(s)	Summary
Configure and apply the MTP bypass modelopt/torch/distill/plugins/megatron.py, examples/megatron_bridge/distill.py, examples/megatron_bridge/tutorials/NVIDIA-Nemotron-3.5-Super-VL-120B-A12B-BF16/README.md	freeze_mtp documentation and warnings describe bypassing MTP forward and loss computation. A context manager temporarily disables MTP processing and sets mtp_num_layers to zero around teacher and student forwards, then restores the original state. The example help text and tutorial describe the bypass and retained checkpoint weights.
Validate frozen MTP execution tests/_test_utils/torch/megatron/models.py, tests/gpu_megatron/torch/distill/plugins/test_distill_megatron.py	Tests check loss-call counts, skipped MTP forwards during recomputation, MTP parameters in the sharded state dict, and model state after forward errors. The tests cover optional MoE and sequence-parallel configurations.

Priority: Normal

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

Change: Bug fix

Sequence Diagram(s)

sequenceDiagram
  participant DistillationForward
  participant SkipMTPForward
  participant TeacherModel
  participant StudentModel
  DistillationForward->>SkipMTPForward: Enter context for teacher pass
  SkipMTPForward->>TeacherModel: Run forward with MTP disabled
  TeacherModel-->>SkipMTPForward: Return forward result
  SkipMTPForward-->>DistillationForward: Restore teacher MTP state
  DistillationForward->>SkipMTPForward: Enter context for student pass
  SkipMTPForward->>StudentModel: Run forward with MTP disabled
  StudentModel-->>SkipMTPForward: Return forward result
  SkipMTPForward-->>DistillationForward: Restore student MTP state

Loading

Possibly related PRs

• NVIDIA/Model-Optimizer#3072: This change includes a nested only_student_forward() context fix. This PR uses only_student_forward() around the student forward.

Suggested reviewers: aanoosheh

Merge Risk: Minimal · up to 8f166

No merge-blocking issue was established. The frozen-MTP bypass appears ready to merge after normal checks.

Pre-merge checks: 5 · 1

Failed checks: 1 warning

Check name	Status	Explanation	Resolution
Docstring Coverage		Docstring coverage is 57.14% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff with cyclomatic complexity of at least 2. Documenta…	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. The changed Python code adds only context-manager/config-copy logic, MTP state handling, and help text; it does not add unsafe torch.load, numpy.load(...…

Full details: Security Anti-Patterns

Explanation

No listed security anti-pattern is introduced. The changed Python code adds only context-manager/config-copy logic, MTP state handling, and help text; it does not add unsafe torch.load, numpy.load(..., allow_pickle=True), hardcoded trust_remote_code=True, external-input eval/exec, or # nosec. The PR changes no pyproject.toml or requirements file. The only added .eval() call is self._teacher_model.eval(), a PyTorch mode-setting method, not Python eval on input.

Full details: Docstring Coverage

Explanation

Docstring coverage is 57.14% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff with cyclomatic complexity of at least 2. Documentation-only edits and deletions within surviving functions are included; unchanged diff context and fully deleted functions are excluded. Analyzed 7 functions across 4 files.

Resolution

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

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.}

[kevalmorabia97]

Thanks for the contribution! This PR now has merge conflicts with main. Could you update your branch with the latest main and fix the conflicts so we can continue the review?

We've added the waiting for feedback label while we wait for your update.

[wiederse]

Thanks for the contribution! This PR now has merge conflicts with main. Could you update your branch with the latest main and fix the conflicts so we can continue the review?

We've added the waiting for feedback label while we wait for your update.

Done.

[codecov]

Codecov Report

❌ Patch coverage is 8.69565% with 21 lines in your changes missing coverage. Please review.
✅ Project coverage is 73.55%. Comparing base (580ffff) to head (a49a26d).

Files with missing lines	Patch %	Lines
modelopt/torch/distill/plugins/megatron.py	8.69%	21 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2891      +/-   ##
==========================================
- Coverage   73.57%   73.55%   -0.03%     
==========================================
  Files         661      661              
  Lines       75682    75699      +17     
==========================================
- Hits        55685    55681       -4     
- Misses      19997    20018      +21     

Flag	Coverage Δ
unit	61.49% <8.69%> (-0.02%)	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[kevalmorabia97]

The mechanism looks right: in the Megatron-LM I checked, GPTModel and the hybrid model gate MTP only through mtp_process and config.mtp_num_layers, both of which this bypasses, and the dense and MoE tests with the forward pre-hooks and negative control cover it well.

Needs action:

• Add a short real-model run with --freeze_mtp on an MTP model, ideally hybrid and with PP>1 (the README you edit targets Nemotron-3.5-Super-VL). The tests only use GPT with PP=1, and the hybrid _forward takes a different path. Report that it trains, the MTP weights are identical in the saved checkpoint, export keeps the MTP keys, and step time and memory against a run without the flag.

No action needed:

• The shallow config copy per forward is cheap, and the original config is restored in finally.
• #2784 changes the same flag's override API, so whichever merges second needs a rebase.

[wiederse]

The mechanism looks right: in the Megatron-LM I checked, GPTModel and the hybrid model gate MTP only through mtp_process and config.mtp_num_layers, both of which this bypasses, and the dense and MoE tests with the forward pre-hooks and negative control cover it well.

Needs action:

• Add a short real-model run with --freeze_mtp on an MTP model, ideally hybrid and with PP>1 (the README you edit targets Nemotron-3.5-Super-VL). The tests only use GPT with PP=1, and the hybrid _forward takes a different path. Report that it trains, the MTP weights are identical in the saved checkpoint, export keeps the MTP keys, and step time and memory against a run without the flag.在 MTP 模型上使用 --freeze_mtp 添加一次简短的真实模型运行，最好采用混合模式且 PP>1（你编辑的 README 面向 Nemotron-3.5-Super-VL）。测试仅使用 PP=1 的 GPT，而混合 _forward 走的是不同路径。报告其能够训练、保存的检查点中 MTP 权重完全一致、导出保留 MTP 键，以及与不带该标志的运行相比的步时和内存。

No action needed:

• The shallow config copy per forward is cheap, and the original config is restored in finally.
• Make freeze_mtp an explicit override of the automatic MTP freeze #2784 changes the same flag's override API, so whichever merges second needs a rebase.

The non-last-PP-stage comment is in commit 8f166bd.

I additionally validated the real NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16 checkpoint using MCore HybridModel (actual runtime class: modelopt.torch.opt.dynamic.DistillHybridModel): 52 body layers (23 Mamba, 23 MoE, 6 Attention), with its native Attention+MoE MTP block. Both conditions used TP2/PP2/EP2/ETP1 on the same four A80080GB GPUs, BF16 weights/gradients, precision-aware optimizer with full CPU optimizer offload, seq 256/MBS 1/GBS 2, and the same seed and initial weights.

Both freeze and control completed 8 steps with finite losses and no skipped iterations. All non-MTP body parameters remained trainable and body samples changed on every rank. Freeze had zero student/teacher MTP module hooks and unchanged MTP hashes. Without the flag, MTP hooks were active and 133/139 and 135/139 MTP parameter tensors changed on the two last-stage ranks.

The saved step 8 checkpoint's 14 MTP tensor entries reconstructed all 278 rank-local parameter hashes exactly. A fresh-process 1-step model-only reload matched both saved body samples and MTP hashes. The native standalone HF export retained all 270 MTP keys with identical dtype, shape and SHA256; 16 exported body samples matched the trained checkpoint and differed from the original weights.

Measured median step time (max rank per step, steps 3–8) was 13.786s with freeze versus 12.948s without. Step peak allocated memory was 51.315 versus 52.512GiB; reserved was 68.680 versus 71.498GiB. This short CPU-offloaded pair did not demonstrate a speedup: freeze was 6.48% slower, although allocator memory was lower. Initialization/checkpoint/export are excluded from these step metrics.

Scope: synthetic tokens, one ordered short pair, model-only recovery (optimizer state was not restored), and the recorded locally patched Bridge/MCore snapshots (base revisions 2a8fcdd05181ffcef261192f71ff5910d0005bf2 / 4fe0cc3fdfb5ca0e5d0c7f25897c661632469672). This validates the training/checkpoint/export mechanics.