Memory leak: nn.parametrize.replace_parameter_4bit + HF gradient_checkpointing_enable() — dequant cache not released on early-stopped recompute

[aka-mnaf-zariche]

Disclosure: this bug was found during an AI-assisted debugging session and this report was drafted with AI assistance (Claude). All measurements come from real runs on our hardware; the account owner posts and vouches for this report.

System Info

• bitsandbytes 0.49.2, transformers 5.13.1, torch 2.13.0+cu126 (CUDA 12.6)
• Windows 11, RTX 3050 8 GB
• Model: allenai/OLMoE-1B-7B-0924

Context

Follow-up to #1849: transformers v5 stores fused-MoE expert weights as 3-D nn.Parameters that the 4-bit walker skips, and until Experts4bit (#1965) lands, bitsandbytes.nn.parametrize.replace_parameter_4bit is the in-tree way to quantize them. That works fine for inference and for training without gradient checkpointing. The problem is the interaction with HF's gradient_checkpointing_enable() — which is the natural companion of 4-bit training on small GPUs, so anyone following the #1849 workaround path for training is likely to hit this.

Prior-report check (why we're filing this as new): before filing we searched GitHub issues for replace_parameter_4bit (zero hits in any repository) and bitsandbytes issues around gradient checkpointing / parametrize. The closest match is #1927 (INT8 CB/SCB state machine breaking under GC recompute), which is a related family but a different mechanism — that one corrupts quantization state in autograd/_functions.py; this one leaks the 4-bit parametrization dequant cache. We found no existing report of this leak, so we're filing it.

Reproduction

import torch
from transformers import AutoModelForCausalLM, BitsAndBytesConfig
from bitsandbytes.nn import parametrize as bnb_parametrize

bnb = BitsAndBytesConfig(
    load_in_4bit=True, bnb_4bit_compute_dtype=torch.bfloat16,
    bnb_4bit_quant_type="nf4", llm_int8_skip_modules=["gate", "lm_head"],
)
model = AutoModelForCausalLM.from_pretrained(
    "allenai/OLMoE-1B-7B-0924", quantization_config=bnb, device_map={"": 0},
)
# experts are fused 3-D nn.Parameters -> quantize them via parametrize (see #1849)
for layer in model.model.layers:
    bnb_parametrize.replace_parameter_4bit(layer.mlp.experts, "gate_up_proj", quant_type="nf4")
    bnb_parametrize.replace_parameter_4bit(layer.mlp.experts, "down_proj", quant_type="nf4")

# freeze everything, train only the 16 router gates (2.1M params)
for p in model.parameters():
    p.requires_grad_(False)
for n, p in model.named_parameters():
    if ".mlp.gate." in n:
        p.requires_grad_(True)

model.gradient_checkpointing_enable()   # <-- the trigger
model.config.use_cache = False
# ... then a plain training loop: AdamW, batch 1, seq len 300, LM loss

Observed: allocated CUDA memory grows monotonically every training step until OOM on the 8 GB card.
Expected: flat memory across steps, as with any frozen-weight training.

Removing either ingredient (no replace_parameter_4bit, or no gradient_checkpointing_enable()) makes the leak disappear.

Root cause (confirmed from source)

_register_parametrization_hooks (bitsandbytes/nn/parametrize.py#L149-L153) manages the global parametrization cache with a pre-hook / post-hook pair:

def _enable_parametrization_cache(module, inputs):
    P._cache_enabled += 1

def _disable_parametrization_cache(module, inputs, output):
    P._cache_enabled -= 1
    if not P._cache_enabled:
        P._cache = {}

The post-hook is registered with plain register_forward_hook(...), and PyTorch does not run regular forward post-hooks when the module's forward raises. Meanwhile, the non-reentrant torch.utils.checkpoint implements early-stopping of recomputation by raising _StopRecomputationError from a saved-tensor pack hook (torch/utils/checkpoint.py), which unwinds straight through the module's forward.

So under HF's gradient_checkpointing_enable() (which checkpoints whole decoder layers, so the parametrized experts module is called via Module.__call__ inside the recomputed region):

1. recompute enters the experts module → pre-hook runs → P._cache_enabled += 1, dequantized weight gets cached in the global P._cache;
2. early-stop fires mid-forward → _StopRecomputationError unwinds → post-hook never runs;
3. P._cache_enabled is now permanently ≥ 1, so P._cache is never cleared again for the rest of the process — every subsequent forward pins its dequantized bf16 expert stacks (hundreds of MB each for fused-expert models) in the cache;
4. memory grows until OOM.

This also explains why our workaround below is leak-free: when the checkpointed region is the experts' bare forward function (not Module.__call__), the recompute path never touches the hook pair, so the counter stays balanced.

Workaround (validated: 900+ steps, flat memory)

Checkpoint at the experts-module boundary instead of the decoder-layer boundary, and leave HF checkpointing off:

import torch.utils.checkpoint as cpk

for layer in model.model.layers:
    ex = layer.mlp.experts
    orig = ex.forward

    def make_ckpt_fwd(orig):
        def fwd(*args, **kwargs):
            if torch.is_grad_enabled():
                return cpk.checkpoint(orig, *args, use_reentrant=False, **kwargs)
            return orig(*args, **kwargs)
        return fwd

    ex.forward = make_ckpt_fwd(orig)

With this, the same run is stable: 5.35 GB peak, ~1.0 s/step over 900 steps (router-only fine-tune of OLMoE-1B-7B on the 8 GB card).

Suggested fix

Make the cache-disable hook exception-safe:

1. Register it with register_forward_hook(_disable_parametrization_cache, always_call=True) — available since torch 2.1, designed exactly for hooks that must run when forward raises; and
2. clamp the counter (P._cache_enabled = max(0, P._cache_enabled - 1)) so an always_call invocation whose matching pre-hook never ran (exception raised before pre-hooks) cannot drive it negative.

Equivalently, the pair could be replaced with a try/finally wrapper mirroring torch.nn.utils.parametrize.cached().

Happy to share the full training script or run additional diagnostics if useful.