Fix nested distillation forward contexts

[Michael-RDev]

What does this PR do?

Type of change: Bug fix

I found that nesting only_student_forward() or only_teacher_forward() causes the outer context to stop working when the inner one exits. This also happens when the inner context uses enable=False.

Both context managers reset their flags to False on exit. This change saves the previous value and restores it in finally, so the enclosing context keeps its execution mode, including after an exception.

Public APIs, checkpoint handling, and loss calculations are unchanged.

Testing

Added eight regression cases using small local models and real forward passes. They cover:

• Both teacher-only and student-only contexts.
• Enabled and disabled inner contexts.
• Normal exits and exceptions.
• Normal execution after leaving a context.
• Overlapping teacher/student contexts.

All eight cases fail before the fix and pass afterward.

Verified with CPU PyTorch:

• Distillation and layerwise test modules — 39 passed.
• All applicable pre-commit checks — passed, including Ruff, mypy, Bandit, and license checks.

GPU and distributed integration tests were not run.

Before this PR is ready for review

• Backward compatible: Yes.
• Copied code or new dependencies: None.
• Necessary tests added: Yes.
• Changelog updated: Yes.
• Claude approval: Not obtained, as I don't use Claude.

Summary by CodeRabbit

• Bug Fixes
  • Nested student-only and teacher-only forwarding contexts now preserve the surrounding execution mode when they exit, including when an inner context is disabled or raises an exception.
  • Exiting a forwarding context restores the expected behavior for subsequent model calls, including when contexts are nested or exceptions occur inside or outside them.

[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]

Navigate logical layers of code changes, visualize relationships, and explore their blast radius.

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

• Configuration used: Repository: NVIDIA/Model-Optimizer/.coderabbit.yaml
• Review profile: CHILL
• Plan: Enterprise
• Run ID: 8192cbb8-48c2-400b-a0cc-66e4d9f33525

📥 Commits

Reviewing files that changed from the base of the PR and between c25fd72 and 5c33d27.

📒 Files selected for processing (2)

• modelopt/torch/distill/distillation_model.py
• tests/unit/torch/distill/test_distill.py

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

---

📝 Walkthrough

Walkthrough

The student-only and teacher-only forward contexts now restore their prior mode flags when they exit. Tests cover nested contexts, disabled contexts, and exceptions.

Changes

Nested forward contexts

Layer / File(s)	Summary
Restore forward mode after nested contexts modelopt/torch/distill/distillation_model.py, tests/unit/torch/distill/test_distill.py, CHANGELOG.rst	The context managers save and restore their prior mode flags. Tests check nested context behavior, outputs, and call counts, including exception cases. The changelog records the fix.

Priority: ⬇️ Low

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

Change: Bug fix

Merge Risk: ⚪ Minimal · up to 5c33d

Nested distillation contexts now restore the enclosing forward mode, preventing inner contexts from affecting later forwards. The implementation and regression coverage align; no actionable merge-blocking risk remains.

🚥 Pre-merge checks | ✅ 5 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 36.36% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 11 functions across 2 files.	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: fixing nested distillation forward contexts.
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	The changed package code only saves and restores the two forward-mode flags. The Python diff adds no prohibited deserialization options, hardcoded trust_remote_code=True, external-input eval() or …

• Fix all pre-merge checks with AI

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create a new PR

• Autopilot · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

---

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