From 10a117eb363ac20f84679455415ff4a4744bec3c Mon Sep 17 00:00:00 2001 From: Chenghao Liu Date: Thu, 10 Sep 2026 00:59:11 +0800 Subject: [PATCH] Release cached activations when exporting distillation models Signed-off-by: Chenghao Liu --- CHANGELOG.rst | 2 + modelopt/torch/distill/distillation_model.py | 2 + .../distill/layerwise_distillation_model.py | 4 +- tests/unit/torch/distill/test_distill.py | 49 +++++++++++++++++++ tests/unit/torch/distill/test_layerwise.py | 16 ++++++ 5 files changed, 72 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 335678eac5a..99a9ca28f19 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -16,6 +16,8 @@ Changelog **Bug Fixes** +- Release cached student and teacher activations when exporting a distillation model, including teacher inputs captured by layerwise distillation. + - 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. - Add FP8 and INT8 recipes that quantize timm ResNet shortcut inputs immediately before residual adds. The torch ONNX example now accepts PTQ and AutoQuantize recipes through ``--recipe`` and uses ``--qformat`` when no recipe is provided. ResNet supports only FP8 and INT8 because TensorRT has limited convolution kernel support; AutoQuantize and other quantization formats are no longer supported for ResNet. diff --git a/modelopt/torch/distill/distillation_model.py b/modelopt/torch/distill/distillation_model.py index fa344385a35..d24edf8f25b 100644 --- a/modelopt/torch/distill/distillation_model.py +++ b/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 diff --git a/modelopt/torch/distill/layerwise_distillation_model.py b/modelopt/torch/distill/layerwise_distillation_model.py index e8cbef99fe2..707b0ed225d 100644 --- a/modelopt/torch/distill/layerwise_distillation_model.py +++ b/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 diff --git a/tests/unit/torch/distill/test_distill.py b/tests/unit/torch/distill/test_distill.py index 69dec86b7f5..35aaa217845 100644 --- a/tests/unit/torch/distill/test_distill.py +++ b/tests/unit/torch/distill/test_distill.py @@ -13,12 +13,16 @@ # 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 @@ -26,6 +30,51 @@ 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) diff --git a/tests/unit/torch/distill/test_layerwise.py b/tests/unit/torch/distill/test_layerwise.py index aea1dd63509..409ef96c05f 100644 --- a/tests/unit/torch/distill/test_layerwise.py +++ b/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