From 498505c1e1afd20a47746ccad5d1aea33f943ddd Mon Sep 17 00:00:00 2001 From: Anai-Guo Date: Thu, 16 Jul 2026 06:13:24 -0700 Subject: [PATCH] fix(parametrize): release the dequant cache when forward raises (#2005) _register_parametrization_hooks pairs a forward pre-hook that increments P._cache_enabled with a plain forward post-hook that decrements it. PyTorch does not run regular forward post-hooks when the module's forward raises, so any exception unwinding out of forward leaves the counter permanently >= 1 and P._cache is never cleared again -- every later forward then pins its dequantized weights for the rest of the process. This is reachable in normal use: non-reentrant torch.utils.checkpoint early-stops recomputation by raising _StopRecomputationError from a saved-tensor pack hook, straight through the module's forward. Combining replace_parameter_4bit with HF gradient_checkpointing_enable() therefore grows memory every step until OOM. Register the post-hook with always_call=True (torch >= 2.1, and the project already requires torch >= 2.4) so it still runs when forward raises, and clamp the counter at zero so an always_call invocation whose matching pre-hook never ran cannot drive it negative -- which would keep the cache pinned just the same. Signed-off-by: Anai-Guo --- bitsandbytes/nn/parametrize.py | 9 +++++++-- tests/test_parametrize.py | 28 ++++++++++++++++++++++++++++ 2 files changed, 35 insertions(+), 2 deletions(-) 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"