Repository navigation
Fix MGDLoss teacher gradient leak and enforce abstract loss balancer - #2668
Merged
AAnoosheh merged 1 commit intoOct 9, 2026
Merged
Conversation
Contributor
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configuration
📒 Files selected for processing (3)
Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 11 remain after this review. 📝 WalkthroughWalkthrough
ChangesDistillation loss behavior
Priority: ⬇️ Low Estimated code review effort: 2 (Simple) | ~10 minutes Change: Bug fix Merge Risk: ⚪ Minimal · up to The change prevents teacher-side gradient flow and rejects balancers missing 🚥 Pre-merge checks | ✅ 5 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #2668 +/- ##
==========================================
+ Coverage 71.69% 78.10% +6.41%
==========================================
Files 642 642
Lines 71285 71285
==========================================
+ Hits 51109 55679 +4570
+ Misses 20176 15606 -4570
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:
|
Contributor
|
/ok to test b8f1a0d |
LinCanNerd
force-pushed
the
fix/distill-mgd-detach-abc-balancer
branch
from
October 8, 2026 08:57
b8f1a0d to
61496a2
Compare
Contributor
Author
|
@AAnoosheh Thanks for kicking off CI. The two failed required checks weren't caused by this PR; the branch was just based on an old
I rebased onto the latest |
Contributor
|
/ok to test 61496a2 |
MGDLoss now detaches the teacher feature map like the other distillation losses, so standalone use no longer backpropagates into the teacher. DistillationLossBalancer now uses ABCMeta, so subclasses missing forward() fail at instantiation instead of at the first loss call. Addresses the MGDLoss and DistillationLossBalancer items in NVIDIA#1926. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Signed-off-by: LinCanNerd <lincanecdl@gmail.com>
LinCanNerd
force-pushed
the
fix/distill-mgd-detach-abc-balancer
branch
from
October 9, 2026 08:50
61496a2 to
c4282e2
Compare
Contributor
|
/ok to test c4282e2 |
AAnoosheh
approved these changes
Oct 9, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do?
Type of change: Bug fix
Fixes two verified defects listed under "Smaller items" in #1926:
MGDLossbackpropagates into the teacher. UnlikeLogitsDistillationLossandMFTLoss,MGDLosspassed the teacher feature map toF.mse_losswithout detaching it, so standalone use (outsideDistillationModel, which runs the teacher underno_grad) populated teacher gradients, wasting memory and compute and updating the teacher if its params are in the optimizer. The teacher target is now detached.DistillationLossBalancerdid not enforce its abstractforward. It used@abc.abstractmethodwithoutABCMeta, so a subclass missingforward()instantiated fine and only failed with a bareNotImplementedErrorat the first loss call. The base class now usesmetaclass=abc.ABCMeta, so such subclasses raiseTypeErrorat construction. Both in-tree subclasses (StaticLossBalancer, Megatron'sLogitsAndIntermediatesLossBalancer) implementforwardand are unaffected.Usage
No API changes.
Testing
test_mgd_loss_does_not_backprop_into_teacherandtest_loss_balancer_requires_forwardtotests/unit/torch/distill/test_distill.py; both fail onmainand pass with this change.pytest tests/unit/torch/distill(CPU, torch 2.14.1): 36 passed.pre-commit runon the changed files: all hooks pass.Before your PR is "Ready for review"
Make sure you read and follow Contributor guidelines and your commits are signed (
git commit -s -S).Make sure you read and follow the Security Best Practices (e.g. avoiding hardcoded
trust_remote_code=True,torch.load(..., weights_only=False),pickle, etc.).DistillationLossBalancersubclass withoutforward()now fails at construction instead of at first call; such a subclass was already unusable)CONTRIBUTING.md: N/AAdditional Information
Addresses the
MGDLossandDistillationLossBalanceritems in #1926 (the issue tracks other unrelated items, so it should stay open).🤖 Generated with Claude Code
Summary by CodeRabbit
Bug Fixes
forwardmethod can no longer be instantiated.Tests