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