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

[pjordanandrsn]

The bug

replace_parameter_4bit / replace_parameter_4bit_prequantized register a
forward_pre_hook/forward_hook pair that gates torch's global parametrization
cache by counter: pre-hook does parametrize._cache_enabled += 1, post-hook
decrements and clears parametrize._cache when it reaches 0.

Activation checkpointing with use_reentrant=False (the recommended mode, and
what HF gradient_checkpointing_enable uses) aborts its backward recompute
mid-forward by design — early stop raises an internal exception the moment
the last needed activation has been rematerialized. A plain forward_hook is
skipped when the forward is interrupted, so for the module holding the last
recomputed save the pair fires pre-only: the global counter leaks +1 per
checkpointed region per step and never returns to zero. From then on the cache
is permanently enabled and never cleared — every dequantized parameter the
model touches stays resident for the rest of training, i.e. a memory leak of
the full dequantized model size (4× the packed 4-bit bytes).

Observed in the wild training MoE models whose experts are quantized via this
path with gradient checkpointing on:

• OLMoE-1B-7B: ~12.9 GiB retained (steady ~13 GiB max_active regardless of store); plain resident training OOMs a 16 GB card at step 0 with the leak and fits with the fix applied
• Qwen3-30B-A3B: ~53 GiB retained — training peaks at 57–60 GiB on a 96 GB card
  and OOMs any 24 GB card at step 1 (allocated 22.37 GiB = base + ~14 blocks
  of retained dequants, Tried to allocate 384 MiB inside
  bitsandbytes/nn/parametrize.py::forward → dequantize_4bit).

Minimal repro of the torch-level mechanism (no bnb needed): a 3-Linear chain
under checkpoint(..., use_reentrant=False), pre/post hooks counting on the
tail module → pre=2, post=1 after one fwd/bwd.

The fix

• Register the disable hook with always_call=True (torch ≥ 2.0), so it also
  runs when the recompute is aborted — verified to balance the pair (2/2) under
  early stop.
• Clamp the decrement at zero. With always_call=True the hook can also fire
  when an earlier pre-hook raised before _enable ran; a negative counter is
  truthy, so if not P._cache_enabled would never clear the cache again.

Tests

TestParametrizationCacheCounterUnderCheckpointing:

1. test_counter_balanced_under_checkpoint_early_stop — real hook pair on the
   tail module of a checkpointed chain, 3 fwd/bwd steps, asserts counter == 0 and
   cache empty. Fails on main, passes with the fix (verified by mutation).
2. test_counter_never_goes_negative — pins the clamp.

CPU-only, no quantization kernels needed; runs in ~1.2 s.

---

Developed with Claude Code (Anthropic) — the diagnosis, fix, and regression tests were produced with AI assistance and validated by the author as described above.

[github-actions]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[matthewdouglas]

Thanks for the PR. This should fix #2005.

One thing I should maybe note is that I never really took the effort to highlight this feature in release notes, examples, demos, etc mainly for the reason that it was just a sort of experiment. Overall, it's not an ideal solution for MoEs, even if it sort of works. I think Axolotl might use it in some way, so the PR here is OK.

[pjordanandrsn]

Thanks for the quick review and merge. Hit it benchmarking Qwen3-30B QLoRA on a 24 GB 3090, where the retained dequant cache quietly erases most of 4-bit's memory win — glad it lands before anyone else meets #2005's path at scale. Agreed on the MoE fit of parametrize generally: the recompute interaction is intrinsic to caching dequantized params under checkpointing, rather than something a patch fully cures.