From c4282e242890b2aa85c5a40dba686068f6716b09 Mon Sep 17 00:00:00 2001 From: LinCanNerd Date: Tue, 6 Oct 2026 09:58:37 +0200 Subject: [PATCH] Fix MGDLoss teacher gradient leak and enforce abstract loss balancer 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 #1926. Co-Authored-By: Claude Opus 5.5 (1M context) Signed-off-by: LinCanNerd --- modelopt/torch/distill/loss_balancers.py | 2 +- modelopt/torch/distill/losses.py | 2 +- tests/unit/torch/distill/test_distill.py | 19 +++++++++++++++++++ 3 files changed, 21 insertions(+), 2 deletions(-) diff --git a/modelopt/torch/distill/loss_balancers.py b/modelopt/torch/distill/loss_balancers.py index 0ceaca3e3c2..18ee15e8e43 100644 --- a/modelopt/torch/distill/loss_balancers.py +++ b/modelopt/torch/distill/loss_balancers.py @@ -30,7 +30,7 @@ STUDENT_LOSS_KEY = "student_loss" -class DistillationLossBalancer(nn.Module): +class DistillationLossBalancer(nn.Module, metaclass=abc.ABCMeta): """Interface for loss balancers.""" def __init__(self): diff --git a/modelopt/torch/distill/losses.py b/modelopt/torch/distill/losses.py index 832437b7e59..f8290b0b46c 100644 --- a/modelopt/torch/distill/losses.py +++ b/modelopt/torch/distill/losses.py @@ -257,7 +257,7 @@ def _loss_fn(self, out_s: torch.Tensor, out_t: torch.Tensor): masked_feats = torch.mul(out_s, mat) new_feats = self.generation(masked_feats) - kd_loss = F.mse_loss(new_feats, out_t) + kd_loss = F.mse_loss(new_feats, out_t.detach()) return kd_loss diff --git a/tests/unit/torch/distill/test_distill.py b/tests/unit/torch/distill/test_distill.py index 7c938aee90d..015973d4d83 100644 --- a/tests/unit/torch/distill/test_distill.py +++ b/tests/unit/torch/distill/test_distill.py @@ -80,6 +80,25 @@ def count_params(params): assert next(kd_loss.parameters()).dtype == torch.float16 +def test_mgd_loss_does_not_backprop_into_teacher(): + teacher_layer = nn.Conv2d(3, 48, 1) + out_s = torch.randn(2, 32, 8, 8, requires_grad=True) + out_t = teacher_layer(torch.randn(2, 3, 8, 8)) + + mtd.MGDLoss(32, 48)(out_s, out_t).backward() + + assert out_s.grad is not None + assert teacher_layer.weight.grad is None + + +def test_loss_balancer_requires_forward(): + class NoForwardBalancer(mtd.DistillationLossBalancer): + pass + + with pytest.raises(TypeError, match="abstract"): + NoForwardBalancer() + + def test_distillation_model_no_balancer(): student = tiny_mobilenet().train() config = {