Release cached activations when exporting distillation models

CHANGELOG.rst

@@ -25,6 +25,7 @@ Changelog
 
 **Bug Fixes**
 
+- Release cached student and teacher activations when exporting a distillation model, including teacher inputs captured by layerwise distillation.
 - Fix ONNX INT8 entropy calibration failing or producing invalid quantization parameters for FP16 activations.
 - Fix ``--use_fsdp2`` HuggingFace checkpoint export gathering the whole model onto rank 0, which made export the dominant phase of a PTQ run and could exhaust host memory on large models. The model is now split into per-decoder-layer units dealt round-robin across ranks; each rank gathers every unit but keeps, packs, and writes only the ones it owns, so a rank buffers roughly ``model / world_size`` instead of the whole checkpoint, and rank 0 writes the combined index. Export configurations that cannot be split this way now raise instead of producing a mismatched checkpoint: FSDP2 combined with another DTensor parallelism (for example FSDP2 + tensor parallel on a 2-D mesh; HSDP is supported), models whose decoder layers cannot be discovered, a decoder layer object reused across layers, and a module that holds the decoder layers while owning parameters of its own.
 - Speed up ``mtq.quantize`` on FSDP2-sharded fused-MoE models. Promoting static-block weight quantizers gathered each expert's slice of the fused weight across ranks even though only quantizer state is read, adding a collective per expert to calibration.

modelopt/torch/distill/distillation_model.py

@@ -124,6 +124,8 @@ def export(self):
         for handle in self._hook_handles:
             handle.remove()
         self._hook_handles.clear()
+        for layer in {layer for pair in self._layers_to_loss for layer in pair}:
+            delattr(layer, "_intermediate_output")
         return super().export()
 
     @property

modelopt/torch/distill/layerwise_distillation_model.py

@@ -67,8 +67,10 @@ def _register_hooks(self):
 
     def export(self):
         """Export the distillation model."""
-        for student_layer, _ in self._layers_to_loss:
+        for student_layer in {student for student, _ in self._layers_to_loss}:
             delattr(student_layer, "_teacher_layer")
+        for teacher_layer in {teacher for _, teacher in self._layers_to_loss}:
+            delattr(teacher_layer, "_intermediate_input")
 
         if hasattr(self, "_lm_head"):
             self.lm_head = self._lm_head

tests/unit/torch/distill/test_distill.py

@@ -13,19 +13,68 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import copy
+import gc
 import inspect
 import warnings
+import weakref
 
 import pytest
 import torch
 import torch.nn as nn
+from _test_utils.torch.quantization.models import SimpleLinear
 from _test_utils.torch.vision_models import get_tiny_mobilenet_and_input
 from torchvision.models import alexnet
 
 import modelopt.torch.distill as mtd
 import modelopt.torch.opt as mto
 
 
+@pytest.mark.parametrize(
+    "pairs",
+    [
+        [("", ""), ("net", "net"), ("net", "linear1")],
+        [("", ""), ("net", "net"), ("linear1", "net")],
+    ],
+)
+def test_export_releases_captured_activations(pairs):
+    student = SimpleLinear(add_linear=True)
+    reference = copy.deepcopy(student)
+    teacher = SimpleLinear(add_linear=True)
+    hook_calls = []
+    handle = student.net.register_forward_hook(lambda *_: hook_calls.append(True))
+    model = mtd.convert(
+        student,
+        mode=[
+            (
+                "kd_loss",
+                {
+                    "teacher_model": teacher,
+                    "criterion": {pair: nn.MSELoss() for pair in pairs},
+                    "loss_balancer": mtd.StaticLossBalancer([1 / len(pairs)] * len(pairs)),
+                },
+            )
+        ],
+    )
+    inputs = student.get_input()
+    model(inputs)
+    layers = {layer for pair in model._layers_to_loss for layer in pair}
+    captures = [weakref.ref(layer._intermediate_output) for layer in layers]
+    assert all(capture() is not None for capture in captures)
+    assert model._intermediate_output.grad_fn is not None
+
+    exported = mtd.export(model)
+    gc.collect()
+
+    assert exported is student
+    assert type(exported) is SimpleLinear
+    assert all(capture() is None for capture in captures)
+    assert all(not hasattr(layer, "_intermediate_output") for layer in layers)
+    torch.testing.assert_close(exported(inputs), reference(inputs))
+    assert len(hook_calls) == 2
+    handle.remove()
+
+
 def get_input_tensor():
     """Dummy input tensor."""
     return torch.rand(2, 3, 112, 112)

tests/unit/torch/distill/test_layerwise.py

@@ -13,7 +13,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import gc
 import warnings
+import weakref
 
 import pytest
 import torch
@@ -47,6 +49,20 @@ def layerwise_distillation_model():
     return layerwise_model
 
 
+def test_layerwise_export_releases_teacher_inputs(layerwise_distillation_model):
+    model = layerwise_distillation_model
+    with model.only_teacher_forward():
+        model(get_input_tensor())
+    teacher_layers = {teacher for _, teacher in model._layers_to_loss}
+    captures = [weakref.ref(layer._intermediate_input[0]) for layer in teacher_layers]
+
+    mtd.export(model)
+    gc.collect()
+
+    assert all(capture() is None for capture in captures)
+    assert all(not hasattr(layer, "_intermediate_input") for layer in teacher_layers)
+
+
 def test_layerwise_hooks_registration(layerwise_distillation_model):
     """Test that layerwise-specific hooks are registered correctly."""
     # Check that student layers have _teacher_layer attribute