Release cached activations when exporting distillation models

[MrCapricornLiu]

What does this PR do?

Type of change: Bug fix

After a distillation forward pass, mtd.export() removes the capture hooks but leaves their last outputs attached to the student and teacher layers. Student outputs can retain an autograd graph even though distillation has finished. Layerwise distillation can also retain teacher inputs after a teacher-only forward.

Remove these caches when exporting, deduplicating shared layers so a layer used by multiple loss pairs is cleaned up once. Existing unrelated user hooks remain installed.

Usage

The existing mtd.export(model) call releases the distillation caches; no API change is required.

Testing

• Three new cases fail on unmodified production code because captured tensors remain reachable after export.
• All 32 tests in tests/unit/torch/distill/test_distill.py and test_layerwise.py pass.
• Tests check weak references, repeated student/teacher layers, root-layer capture, restored model class and forward behavior, user-hook preservation, and a layerwise teacher-only forward.
• A small GPT-2 CUDA reproduction on an H800 retains both captured tensors before the fix and neither after it. With the teacher model retained separately in both runs, live allocated memory drops by 32,539,648 bytes on export after the fix, versus zero before it. This measures that specific reproduction, not a general memory reduction estimate.
• Changed-file pre-commit checks pass, including Ruff, mypy, licenses, and Bandit.

Distributed execution, checkpoint serialization, and full-model training performance were not tested.

Before your PR is "Ready for review"

• Backward compatible: yes; only distillation-owned caches are removed after export.
• Copied code or new PIP dependencies: none.
• Necessary tests: added to the existing distillation modules.
• Changelog: updated.
• Upstream review: pending.

Additional Information

Related to the activation-retention report in item 6 of #1926 and #1923. Other reports in #1926 are outside this change.

Summary by CodeRabbit

• Bug Fixes

  • Fixed distillation-model export to release cached student and teacher activations.
  • Fixed layerwise distillation export to clear captured teacher inputs and temporary layer data.
  • Ensured exported models preserve their expected outputs and forward behavior.

• Tests

  • Added regression coverage verifying temporary activation data is released during export.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 6da9ac67-64c5-4e45-91aa-5012112da3a5

📥 Commits

Reviewing files that changed from the base of the PR and between 079078d and 10a117e.

📒 Files selected for processing (5)

• CHANGELOG.rst
• modelopt/torch/distill/distillation_model.py
• modelopt/torch/distill/layerwise_distillation_model.py
• tests/unit/torch/distill/test_distill.py
• tests/unit/torch/distill/test_layerwise.py

Included review availability: Your plan provides up to 12 included reviews per hour; 10 remain after this review.

---

📝 Walkthrough

Walkthrough

Distillation export now clears cached intermediate outputs and inputs from student and teacher layers. Regression tests verify garbage collection, attribute removal, preserved outputs, and hook behavior.

Changes

Distillation export cleanup

Layer / File(s)	Summary
Standard distillation export cleanup modelopt/torch/distill/distillation_model.py, tests/unit/torch/distill/test_distill.py	Export removes _intermediate_output attributes after hook cleanup. Parametrized tests verify activation release, output preservation, and hook behavior.
Layerwise export cleanup modelopt/torch/distill/layerwise_distillation_model.py, tests/unit/torch/distill/test_layerwise.py, CHANGELOG.rst	Layerwise export removes temporary teacher references and _intermediate_input attributes. Tests verify teacher input release. The changelog records the fix.

Estimated code review effort: 2 (Simple) | ~10 minutes

Suggested reviewers: aanoosheh, kevalmorabia97

Merge Risk: ⚪ Minimal · up to 10a11

Distillation export now releases cached student and teacher activations while preserving model outputs and user hooks. Current tests cover standard and layerwise cleanup, so the change is merge-ready.

🚥 Pre-merge checks | ✅ 5 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 57.14% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 7 functions across 4 files. (1 skipped: 1…	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (5 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly and concisely describes the main change: releasing cached activations during distillation-model export.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Security Anti-Patterns	✅ Passed	No listed security anti-pattern was introduced. The authoritative diff changes only two modelopt distillation files; additions are delattr cleanup operations. The changed production files contain …

Full details: Docstring Coverage

Explanation

Docstring coverage is 57.14% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 7 functions across 4 files. (1 skipped: 1 unsupported.)

• Fix all pre-merge checks with AI

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands.}