Fix MGDLoss teacher gradient leak and enforce abstract loss balancer

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):

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
 

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 = {