Repository navigation
Make freeze_mtp an explicit override of the automatic MTP freeze - #2784
Conversation
freeze_mtp=False could not keep the MTP head training: it was frozen anyway when the LM loss was skipped and only the backbone was quantized. freeze_mtp is now None by default, which keeps that automatic rule; True and False are authoritative (--freeze_mtp / --no-freeze_mtp). The rule's helper is renamed to _has_quantized_backbone_with_unquantized_mtp, after what it checks, and inspects only the student. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Yue <yueshen@nvidia.com>
|
No actionable comments were generated in the recent review.
 |
| Layer / File(s) | Summary |
|---|---|
Tri-state freeze option CHANGELOG.rst, examples/megatron_bridge/distill.py, modelopt/torch/distill/plugins/megatron.py |
The CLI supports --freeze_mtp and --no-freeze_mtp. DistillationConfig.freeze_mtp now defaults to None; the changelog and option help describe the explicit choices and conditional default. |
MTP freeze resolution and tests modelopt/torch/distill/plugins/megatron.py, tests/gpu_megatron/torch/distill/plugins/test_distill_megatron.py |
The quantization check requires a quantized backbone and an unquantized MTP module. Explicit freeze values control MTP loss skipping; when unset, skipping depends on LM-loss skipping and that quantization check. Tests cover loss behavior, parameter gradients, and parameter updates. |
Estimated code review effort: 3 (Moderate) | ~20 minutes
Change: Bug fix
Sequence Diagram(s)
sequenceDiagram
participant Config as DistillationConfig
participant Predicate as Quantization predicate
participant Loss as MTP loss logic
Config->>Loss: Provide freeze_mtp and kd_loss_alpha
Loss->>Predicate: Check backbone and MTP quantization
Predicate-->>Loss: Return quantization state
Loss->>Loss: Resolve whether to skip MTP loss
Possibly related PRs
- NVIDIA/Model-Optimizer#2891: Uses
freeze_mtpto skip teacher and student MTP forward and recomputation while retaining MTP weights. - NVIDIA/Model-Optimizer#2911: Changes PTQ exclusion patterns that affect quantization state used by automatic MTP freezing.
- NVIDIA/Model-Optimizer#3072: Cherry-picks the PTQ exclusion change that affects quantization state used by automatic MTP freezing.
Suggested reviewers: kevalmorabia97
Merge Risk: 
c2328
-
The three-state freeze control is ready to merge with normal checks.
Pre-merge checks: 
-
Check name Status Explanation Description Check Check skipped - CodeRabbit’s high-level summary is enabled. Title check The title states the main change: freeze_mtp becomes an explicit override of the automatic MTP freeze. The raw summary confirms this. freeze_mtp changes from a bool defaulting to False into a three-st… Docstring Coverage Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff with cyclomatic complexity of at least 2. Documentat… 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 PASS. The pull request changes only examples/megatron_bridge/distill.pyandmodelopt/torch/distill/plugins/megatron.pyamong the covered production Python files. The added code contains no `torch.…Full details: Title check
-
Explanation
The title states the main change: freeze_mtp becomes an explicit override of the automatic MTP freeze. The raw summary confirms this. freeze_mtp changes from a bool defaulting to False into a three-state setting (None, True, False). When unset, the automatic rule still applies. Explicit True and False override it. The CHANGELOG, the CLI flags (--freeze_mtp and --no-freeze_mtp), the DistillationConfig field, the helper rename to _has_quantized_backbone_with_unquantized_mtp, and the tests all match this title.
Full details: Docstring Coverage
-
Explanation
Docstring coverage is 100.00% which is sufficient. 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 5 functions across 3 files. (1 skipped: 1 unsupported.)
Full details: Security Anti-Patterns
-
Explanation
PASS. The pull request changes only
examples/megatron_bridge/distill.pyandmodelopt/torch/distill/plugins/megatron.pyamong the covered production Python files. The added code contains notorch.load(..., weights_only=False),numpy.load(..., allow_pickle=True), hardcodedtrust_remote_code=True,eval()/exec(), or# nosec. It also adds no changes topyproject.tomlorrequirements.txt. The changed test file is excluded bySECURITY.md.
-
Finishing Touches
-
Generate docstrings
-
- Commit to this branch
- Create a new PR
Generate unit tests (beta)
-
- Commit to this branch
- Create a new PR
-
- Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts
cjluo-nv
left a comment
cjluo-nv
left a comment
There was a problem hiding this comment.
Bot review (gpt-6.1-sol) — DM the bot to share feedback.
LGTM: the tri-state override is correct and well covered; the justified edits to existing tests require human sign-off.
Needs action:
- Confirm the existing-test updates in
tests/gpu_megatron/torch/distill/plugins/test_distill_megatron.py: default-behavior cases now useNone, and the renamed training-step test preserves prior coverage while adding override cases.
No action needed:
- Automatic detection now excludes the teacher, and the CLI passes all three states through unchanged.
- Tests were inspected, not executed; the PR reports its GB300 run is pending.
|
/claude review |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
|
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #2784 +/- ##
==========================================
- Coverage 80.29% 76.16% -4.14%
==========================================
Files 675 675
Lines 77278 85297 +8019
==========================================
+ Hits 62054 64964 +2910
- Misses 15224 20333 +5109
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
danielkorzekwa
left a comment
danielkorzekwa
left a comment
There was a problem hiding this comment.
good work, much clearer now with freeze-mtp=None default state
…istate Signed-off-by: Yue <yueshen@nvidia.com> # Conflicts: # CHANGELOG.rst
What does this PR do?
Type of change: Bug fix (follow-up to #2770)
Addresses @danielkorzekwa's review of #2770:
freeze_mtp=Falsecould not keep the MTP head training (comment). The MTP head was still frozen whenever the LM loss was skipped and the student was quantized except for its MTP head, because the setting was OR-ed with that automatic rule.freeze_mtpis now three-state:freeze_mtpNone(default, omitted)kd_loss_alpha == 1.0and the student is quantized except for its MTP headTrueFalseOn the CLI,
examples/megatron_bridge/distill.pytakes--freeze_mtp/--no-freeze_mtp(argparse.BooleanOptionalAction), and omitting both keeps the automatic rule. YAML configs pass throughDistillationConfigunchanged.The helper's name hid what it checks (comment).
_mtp_excluded_from_quantizationis renamed_has_quantized_backbone_with_unquantized_mtp: it is true when the model has MTP submodules and enabled quantizers, none of them inside MTP. It now inspects only the student, so a teacher's MTP head can no longer make it true. The reason for the automatic freeze moved next to where the setting is resolved.freeze_mtp=Falsechanges meaning: it used to mean "automatic", and now it keeps the MTP head training. Omitting the field keeps the previous behavior.freeze_mtpwas added in #2770 and is not in a release yet."Freeze" covers the MTP-specific parameters only. The embedding and output layer that the MTP head shares with the backbone keep training through the main loss.
Usage
Testing
tests/gpu_megatron/torch/distill/plugins/test_distill_megatron.pypasses on GB300 (28 passed):test_skip_lm_loss_with_mtpcovers omitted,TrueandFalse. It includes the regression case: a quantized backbone with an unquantized MTP head,kd_loss_alpha=1.0andfreeze_mtp=False, where the MTP loss stays active and the MTP parameters stay trainable.test_freeze_mtp_training_step, which replacestest_freeze_mtp_with_quantized_mtp, runs forward, KD backward and an AdamW step across quantized and unquantized MTP heads and all three settings. It checks that the MTP head gets gradients exactly when it is not frozen, that a frozen head is bit-identical after the step, and that the backbone always trains. The attachedTensorQuantizers are quantization-state fixtures; the test does not quantize numerics.test_has_quantized_backbone_with_unquantized_mtpcovers the renamed helper.Before your PR is "Ready for review"
freeze_mtpbehaves as before. An explicitfreeze_mtp=Falsenow keeps the MTP head training; that value exists only since Add freeze_mtp to Megatron distillation to keep the MTP head fixed #2770, which is unreleased.🤖 Generated with Claude Code
Summary by CodeRabbit
--freeze_mtpand--no-freeze_mtpto explicitly freeze or keep the student’s MTP parameters trainable. Freezing also disables MTP loss.