diff --git a/bitsandbytes/nn/parametrize.py b/bitsandbytes/nn/parametrize.py index 4a956c7fa..b6de3c289 100644 --- a/bitsandbytes/nn/parametrize.py +++ b/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( diff --git a/tests/test_parametrize.py b/tests/test_parametrize.py index 50260123b..5293d5d29 100644 --- a/tests/test_parametrize.py +++ b/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"