fix(parametrize): release the dequant cache when forward raises (#2005)

[Anai-Guo]

Fixes #2005.

Problem

_register_parametrization_hooks manages the global parametrization cache with a pre-hook / post-hook pair:

module.register_forward_pre_hook(_enable_parametrization_cache)   # P._cache_enabled += 1
module.register_forward_hook(_disable_parametrization_cache)      # P._cache_enabled -= 1

PyTorch does not run regular forward post-hooks when the module's forward raises. So if an exception unwinds out of forward, P._cache_enabled stays >= 1 forever and P._cache is never cleared again — every subsequent forward pins its dequantized weights for the rest of the process.

This is reachable in ordinary use, not just on error paths: non-reentrant torch.utils.checkpoint implements early-stopping of recomputation by raising _StopRecomputationError from a saved-tensor pack hook, straight through the module's forward. So replace_parameter_4bit + HF gradient_checkpointing_enable() — the natural pairing for 4-bit training of fused-MoE experts on small GPUs (see #1849) — leaks memory every step until OOM.

Fix

1. Register the post-hook with always_call=True, so it still runs when forward raises. This is available since torch 2.1 and the project already requires torch>=2.4, so no version guard is needed.
2. Clamp the counter at zero. always_call=True hooks can fire even when the matching pre-hook never ran (e.g. an earlier pre-hook raised); a negative counter would make if not P._cache_enabled false forever and pin the cache just the same.

Verification

Confirmed the underlying hook semantics on torch 2.5.1 — when forward raises, the pre-hook runs, a plain post-hook does not, and an always_call=True post-hook does.

Added test_cache_released_when_forward_raises, which drives a parametrized module whose forward raises after touching the weight, and asserts the cache is released. It reproduces the bug:

• without this fix: AssertionError: assert 1 == 0 (P._cache_enabled stuck at 1, one tensor pinned)
• with this fix: passes

Full tests/test_parametrize.py suite passes locally (159 passed, CPU + CUDA), and ruff check / ruff format are clean with the pinned v0.14.3.

🤖 Generated with Claude Code

[matthewdouglas]

Hi @Anai-Guo,

Thanks for the contribution. Closing as a duplicate of #1999, which implemented the same fix three days before #2005 was filed. For future contributions, please have your agent search open PRs before starting; #1999 was already up and would've shown in a quick search.