fix(parametrize): cache-gate hooks leak the global parametrize cache under activation checkpointing (all dequantized params retained)

bitsandbytes/nn/parametrize.py

@@ -127,7 +127,11 @@ def replace_parameter_4bit(
 
 
 def _disable_parametrization_cache(module: nn.Module, inputs: tuple[Any, ...], output: Any):
-    P._cache_enabled -= 1
+    # Clamp instead of a bare decrement: with ``always_call=True`` this hook also runs
+    # when the forward raised before the pre-hook incremented (e.g. an earlier pre-hook
+    # failed), and the counter must never go negative — a negative value is truthy, so
+    # ``if not P._cache_enabled`` would stop clearing the cache forever.
+    P._cache_enabled = max(0, P._cache_enabled - 1)
     if not P._cache_enabled:
         P._cache = {}
 
@@ -149,8 +153,18 @@ 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=True`` is load-bearing: activation checkpointing with
+    # ``use_reentrant=False`` aborts its backward recompute mid-forward by design
+    # (early stop, via an internal exception) once the last needed activation has
+    # been rematerialized. A plain forward hook is skipped in that case, so the
+    # global ``parametrize._cache_enabled`` counter leaks upward once per
+    # checkpointed region per step, after which the cache is enabled (and never
+    # cleared) for the remainder of training — every dequantized parameter this
+    # module produces stays resident, i.e. a memory leak of the full dequantized
+    # model size (4x the packed 4-bit bytes).
     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

@@ -431,3 +431,69 @@ 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"
+
+
+class TestParametrizationCacheCounterUnderCheckpointing:
+    """The cache-gate hooks must stay balanced under activation checkpointing.
+
+    ``use_reentrant=False`` checkpointing aborts its backward recompute mid-forward by
+    design (early stop) once the last needed activation is rematerialized. A plain
+    forward hook is skipped for the module holding that last save, so the global
+    ``parametrize._cache_enabled`` counter leaked +1 per checkpointed region per step,
+    the cache was never cleared again, and every dequantized parameter stayed resident
+    (4x the packed model bytes). ``always_call=True`` keeps the pair balanced.
+    """
+
+    def _reset(self):
+        import torch.nn.utils.parametrize as P
+
+        P._cache_enabled = 0
+        P._cache = {}
+
+    def test_counter_balanced_under_checkpoint_early_stop(self):
+        import torch.nn.utils.parametrize as P
+        from torch.utils.checkpoint import checkpoint
+
+        from bitsandbytes.nn.parametrize import _register_parametrization_hooks
+
+        self._reset()
+
+        class Chain(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.a = nn.Linear(8, 8)
+                self.tail = nn.Linear(8, 8)
+
+            def forward(self, x):
+                return self.tail(self.a(x))
+
+        m = Chain()
+        # Register the real hook pair on the tail module — the one whose recompute
+        # holds the last save, where early stop aborts before the forward hook.
+        # (torch < 2.5 skips the state-dict hook inside; the cache pair is what we need.)
+        _register_parametrization_hooks(m.tail, "weight")
+
+        for _ in range(3):
+            x = torch.randn(2, 8, requires_grad=True)
+            checkpoint(m, x, use_reentrant=False).sum().backward()
+
+        assert P._cache_enabled == 0, (
+            f"parametrize cache counter leaked to {P._cache_enabled}; dequantized "
+            "parameters would be retained for the rest of training"
+        )
+        assert P._cache == {}
+        self._reset()
+
+    def test_counter_never_goes_negative(self):
+        import torch.nn.utils.parametrize as P
+
+        from bitsandbytes.nn.parametrize import _disable_parametrization_cache
+
+        self._reset()
+        P._cache[("k",)] = torch.zeros(1)
+        # always_call can fire the disable hook when an earlier pre-hook raised before
+        # enable ran; a negative counter is truthy and would stop cache clearing forever.
+        _disable_parametrization_cache(nn.Identity(), (), None)
+        assert P._cache_enabled == 0
+        assert P._cache == {}
+        self._reset()