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

bitsandbytes/nn/parametrize.py

@@ -127,7 +127,9 @@ def replace_parameter_4bit(
 
 
 def _disable_parametrization_cache(module: nn.Module, inputs: tuple[Any, ...], output: Any):
-    P._cache_enabled -= 1
+    # Clamp at zero: with always_call=True this hook can fire without its matching
+    # pre-hook having run, and a negative count would never clear the cache again.
+    P._cache_enabled = max(0, P._cache_enabled - 1)
     if not P._cache_enabled:
         P._cache = {}
 
@@ -149,8 +151,11 @@ def _register_parametrization_hooks(module: nn.Module, param_name: str):
     # Register hooks to enable caching for the dequantization parametrization.
     # This helps preserve time and memory when the same quantized parameter
     # is accessed multiple times in the forward computation.
+    # always_call so that an exception escaping forward -- notably the
+    # _StopRecomputationError non-reentrant checkpointing raises to early-stop
+    # recomputation -- cannot leave the cache enabled and its tensors pinned.
     module.register_forward_pre_hook(_enable_parametrization_cache)
-    module.register_forward_hook(_disable_parametrization_cache)
+    module.register_forward_hook(_disable_parametrization_cache, always_call=True)
 
 
 def _parametrized_state_dict_post_hook(

tests/test_parametrize.py

@@ -1,6 +1,7 @@
 import pytest
 import torch
 import torch.nn as nn
+import torch.nn.utils.parametrize as P
 
 from bitsandbytes import functional as F
 from bitsandbytes.nn.parametrize import (
@@ -431,3 +432,30 @@ def test_gradient_behavior(device, dtype):
     # The dequantized output should also not require gradients
     reconstructed = module.weight_2d
     assert not reconstructed.requires_grad, "Dequantized parameter should not require gradients"
+
+
+def test_cache_released_when_forward_raises():
+    """The cache must be released even when the module's forward raises.
+
+    Non-reentrant activation checkpointing early-stops recomputation by raising
+    through the module's forward, which would otherwise leave the cache enabled
+    and pin every dequantized weight for the rest of the process.
+    """
+
+    class RaisingModule(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.weight = nn.Parameter(torch.randn(64, 64, dtype=torch.float32))
+
+        def forward(self, x):
+            _ = self.weight  # populate the cache, then unwind before the post-hook
+            raise RuntimeError("boom")
+
+    module = RaisingModule()
+    replace_parameter_4bit(module, "weight", quant_type="nf4")
+
+    with pytest.raises(RuntimeError, match="boom"):
+        module(torch.randn(1, 64))
+
+    assert P._cache_enabled == 0, "Cache should be disabled again after forward raises"
+    assert not P._cache, "Dequantized weights should not stay pinned after forward raises"