`MatMul4Bit` keeps `B` and `quant_state` on `ctx` instead of `save_for_backward`, so gradient checkpointing cannot recompute them

[MakazhanAlpamys]

Summary

MatMul4Bit.forward stores the packed weight and the QuantState as ordinary
Python attributes on ctx (ctx.tensors = (None, B), ctx.state = quant_state)
rather than through save_for_backward. In the ordinary case this is fine and
costs nothing: the weight is a stable Parameter, and holding a reference to it
is exactly as correct as saving it. It stops being fine for a caller that
re-materialises the weight — weight offloading, layer streaming, anything that
writes a layer's bytes into a recycled device buffer just before use.
torch.utils.checkpoint discards and recomputes saved tensors, so under
checkpointing such a caller expects the recompute to hand the backward a fresh
copy of the weight. A plain ctx attribute is invisible to that mechanism: the
reference taken during the original forward survives into the backward and, by
then, points at a buffer that has been refilled with a different layer. The
result is silently wrong gradients — the forward stays bit-exact and the loss
curve looks healthy — and the only workaround available to the caller costs
O(model) memory rather than O(window), which for a streaming implementation
deletes the reason it exists. bf16 through the same buffer-recycling harness is
unaffected, because it goes through MmBackward0, which does use
save_for_backward.

Where this came from. The full measurement record is public:
benchmarks/gate-h100-validation.md.
It is the working record rather than a write-up — it contains the six hypotheses
that were tested and rejected before the ctx mechanism was found, the
discriminating controls, and the measurements that turned out invalid and why.
STEP 9 names the mechanism, STEP 6 isolates the trigger, STEP 13 covers the
gemm_4bit M-dispatch referenced below.

Environment

bitsandbytes	0.50.0
torch	2.13.0+cu130 (CUDA 13.0)
GPU	NVIDIA H100 80GB HBM3, sm90, 132 SMs
driver	590.48.01
Python	3.12.3
OS	Ubuntu 24.04.3 LTS, kernel 6.8.0-100
transformers / peft / accelerate	4.57.6 / 0.20.0 / 1.14.0
bitsandbytes install method	pip install bitsandbytes (PyPI wheel, not a source build)
older bitsandbytes versions	not tested — the ctx pattern looks unchanged in the current source, but we have only measured 0.50.0 and would rather say so than imply broader coverage

Minimal reproducer

Self-contained, no downloads, needs one CUDA device, runs in about a minute.

It builds an 8-layer stack of NF4 weights and a trainable parameter per layer
placed upstream of the 4-bit matmul, so the trainable gradient is exactly the
quantity MatMul4Bit.backward computes from B and ctx.state. Every layer's
forward is wrapped in torch.utils.checkpoint(use_reentrant=False), and the
layer's weight bytes are copied into a recycled slot inside the checkpointed
region — so the recompute does re-materialise the weight, which is the whole
premise the caller is relying on.

Three arms:

• nf4 / recycled buffers — 2 slots for 8 layers, round-robin.
• nf4 / private buffers — each layer owns its buffer, nothing is ever recycled. This is the reference.
• bf16 / recycled buffers — identical harness, native F.linear. This is the control that shows the harness itself is sound.

One note on shapes: bitsandbytes::gemm_4bit dispatches on M
(_gemm_4bit_custom_max_m = 1536 on CUDA; above it, _dequant_linear_fallback
is taken). The reproducer keeps M = 256 so it stays in the regime this report
is about; a reproducer written at a large M can end up comparing two arms that
are running the same code and concluding, incorrectly, that nothing is wrong.

"""
bitsandbytes MatMul4Bit + gradient checkpointing + a recycled weight buffer.

The trainable parameter sits upstream of the 4-bit matmul, so its gradient is
computed from `F.dequantize_4bit(B, ctx.state)` inside MatMul4Bit.backward --
i.e. from whatever `ctx` is still pointing at when the backward runs.
"""

import torch
import torch.nn.functional as F
import bitsandbytes as bnb
from bitsandbytes.functional import quantize_4bit
from torch.utils.checkpoint import checkpoint

DEV = "cuda"
DTYPE = torch.bfloat16
N_LAYERS = 8
DIM = 4096
TOKENS = 256          # keep M well below _gemm_4bit_custom_max_m = 1536
POOL_SLOTS = 2        # a streaming caller's buffer pool


class SlotPool:
    """Round-robin device buffers. Every acquire() advances the cursor, so a
    slot handed out during the ORIGINAL forward is recycled later in the step --
    which is what any streaming / offloading caller does."""

    def __init__(self, n_slots, like):
        self.slots = [torch.empty_like(like) for _ in range(n_slots)]
        self.cursor = 0

    def acquire(self, src):
        slot = self.slots[self.cursor % len(self.slots)]
        self.cursor += 1
        slot.copy_(src)
        return slot


def run(quant, recycle, params_init):
    torch.manual_seed(0)
    weights = [torch.randn(DIM, DIM, device=DEV, dtype=DTYPE) / 32
               for _ in range(N_LAYERS)]

    if quant == "nf4":
        sources, states = [], []
        for w in weights:
            packed, st = quantize_4bit(w, blocksize=64, quant_type="nf4",
                                       compress_statistics=True)
            sources.append(packed)
            states.append(st)
    else:
        sources, states = weights, [None] * N_LAYERS

    pool = SlotPool(POOL_SLOTS, sources[0]) if recycle else None
    params = [p.clone().requires_grad_(True) for p in params_init]

    def body(idx):
        def fn(x, scale):
            h = x * scale                                  # trainable, upstream
            # The fill is INSIDE the checkpointed region on purpose: the
            # recompute is supposed to re-materialise the weight.
            w = pool.acquire(sources[idx]) if recycle else sources[idx]
            if quant == "nf4":
                return bnb.matmul_4bit(h, w.t(), quant_state=states[idx])
            return F.linear(h, w)
        return fn

    x = torch.randn(TOKENS, DIM, device=DEV, dtype=DTYPE, requires_grad=True)
    for i in range(N_LAYERS):
        x = checkpoint(body(i), x, params[i], use_reentrant=False)

    x.float().pow(2).mean().backward()
    return [p.grad.detach().float().clone() for p in params]


def main():
    torch.manual_seed(1234)
    init = [torch.ones(DIM, device=DEV, dtype=DTYPE)
            + 0.01 * torch.randn(DIM, device=DEV, dtype=DTYPE)
            for _ in range(N_LAYERS)]

    ref = run("nf4", recycle=False, params_init=init)    # nothing recycled
    got = run("nf4", recycle=True, params_init=init)     # 2 slots, 8 layers
    bf16 = run("bf16", recycle=True, params_init=init)   # control
    bf16_ref = run("bf16", recycle=False, params_init=init)

    print(f"{'layer':>5}  {'nf4 recycled vs private':>24}  {'bf16 recycled vs private':>25}")
    for i in range(N_LAYERS):
        d4 = (got[i] - ref[i]).abs().max().item()
        d16 = (bf16[i] - bf16_ref[i]).abs().max().item()
        print(f"{i:>5}  {d4:>16.6e} {'MISMATCH' if d4 else '  ok    '}"
              f"  {d16:>16.6e} {'MISMATCH' if d16 else '  ok    '}")


if __name__ == "__main__":
    main()

What it prints

The slot a layer's forward captured is recycled by a later acquire(), so
MatMul4Bit.backward dequantises another layer's bytes. Only the layers whose
slot happens never to be reclaimed before their backward runs come out right —
with POOL_SLOTS = 2 that is the last one or two.

Broken (bitsandbytes as it is today). This is the verbatim output of the script
above on the environment in the table, exit code 0:

layer   nf4 recycled vs private   bf16 recycled vs private
    0      6.762500e+01 MISMATCH      0.000000e+00   ok
    1      5.903125e+01 MISMATCH      0.000000e+00   ok
    2      5.750000e+01 MISMATCH      0.000000e+00   ok
    3      5.987500e+01 MISMATCH      0.000000e+00   ok
    4      5.750000e+01 MISMATCH      0.000000e+00   ok
    5      5.587500e+01 MISMATCH      0.000000e+00   ok
    6      0.000000e+00   ok          0.000000e+00   ok
    7      0.000000e+00   ok          0.000000e+00   ok

Two details worth reading off it. The surviving layers are the LAST
POOL_SLOTS — 6 and 7 with two slots — which is what you would predict if the
captured reference is simply overwritten by later acquire() calls. And the
magnitudes are not small: these gradients are wrong by order 60 against their own
scale, not by a rounding difference, while the forward stays bit-exact.

Correct (what save_for_backward would give, and what the bf16 column already
shows):

layer   nf4 recycled vs private   bf16 recycled vs private
    0     0.000000e+00   ok              0.000000e+00   ok
   ...
    7     0.000000e+00   ok              0.000000e+00   ok

The bf16 column is the point of the control: the identical recycling harness,
identical checkpointing, and a native op that saves its operand properly, gives
exact gradients. Nothing about buffer recycling is inherently unsound here.

Root cause

bitsandbytes/autograd/_functions.py, MatMul4Bit (line numbers as of 0.50.0):

55:  ctx.state = quant_state          # plain attribute
59:  ctx.tensors = (None, B)          # plain attribute -- not save_for_backward
...
85:  grad_A = torch.matmul(grad_output, F.dequantize_4bit(B, ctx.state).to(...))

B is the packed 4-bit weight and ctx.state carries absmax, offset and
code. Both are captured during the forward and read back verbatim in the
backward.

torch.utils.checkpoint (non-reentrant) works through
torch.autograd.graph.saved_tensors_hooks: it packs a placeholder for each
saved tensor during the original forward and re-materialises it by re-running
the forward at backward time. Attributes hung on ctx never enter that path, so
they are neither dropped nor recomputed. For a caller whose weight tensor is a
recycled buffer, the recompute produces a correct, fresh weight that the node
then ignores in favour of the reference it took on the first pass.

Why save_for_backward fixes it, and what it changes

Saving B and the tensor fields of quant_state (absmax, and under double
quantisation nested_absmax / nested_offset, plus code) through
save_for_backward, and reassembling the QuantState in backward, makes the
node participate in the mechanism checkpointing relies on. The recompute would
then supply the operands, exactly as it does for every native op.

What this costs in the ordinary path: nothing that we can see. Without
checkpointing, a saved tensor is held alive by the graph in the same way a ctx
attribute is — same lifetime, same math, same allocation. The one behavioural
change is that saved tensors carry autograd's version counter, so an in-place
mutation of B between forward and backward becomes a raised error rather than
silently wrong numbers. In our experience that is a feature: when we disabled
checkpointing in our own harness, autograd's version counter caught exactly this
class of overwrite on the unquantised tensors in the same layer, and the error
message was the most useful diagnostic of the whole investigation.

What it gains: correctness under gradient checkpointing for any caller that
re-materialises weights, and — see below — a memory profile that stays
O(window) instead of O(model).

The awkward part is that quant_state is a Python object rather than a tensor,
so the change is not a one-line substitution; the tensor fields have to be routed
through save_for_backward and the QuantState rebuilt on the way out. The
ctx.tensors = (None, B) shape also looks like it is inherited from
MatMul8bitLt, so there may be a reason for it that is not visible from outside.

The evidence, briefly

Two discriminating experiments, each with its control in the same process, five
independent runs and three backward repetitions per run.

Forward-only vs backward-only de-aliasing. If the forward-time capture is the
consumer, handing the op a private copy on the forward must fix it and doing
the same on the backward must not.

run1 control    -> WRONG  ['60/128', '8/128', '8/128']
run1 clone_fwd  -> EXACT  ['128/128', '128/128', '128/128']
run1 clone_bwd  -> WRONG  ['8/128', '8/128', '8/128']
runs 2-5        -> identical in all three arms

5/5 runs, 15/15 repetitions; the two arms separate completely and in the
predicted direction.

Only the bitsandbytes-captured tensors matter. De-aliasing only the packed
weight and its ::absmax / ::nested_* sidecars must fix it, while de-aliasing
only the rest of the layer (layernorms and biases, which reach native ops that do
use save_for_backward) must not. Without the second arm the first is satisfied
by "copying most of the bytes helps somehow".

run1 control            -> WRONG  ['128/128', '8/128', '8/128']
run1 clone_fwd_quant    -> EXACT  ['128/128', '128/128', '128/128']
run1 clone_fwd_nonquant -> WRONG  ['8/128', '8/128', '8/128']
runs 2-5                -> identical in all three arms

Both experiments reproduce on a real Qwen2.5-32B checkpoint, not only on the
synthetic stack: 5/5 runs, 15/15 repetitions.

Some things we ruled out along the way, in case they save someone time: a full
torch.cuda.synchronize() after the buffer is ready does not fix it, so this
is aliasing rather than a timing race; and bf16 is exact through the same harness
while moving 3.9x more bytes per layer, which is what first pointed at the 4-bit
path specifically.

Workarounds we tried, and their cost

All measured with correctness asserted inside each arm, so a fast arm cannot
quietly be the wrong one.

workaround	gradients	throughput	peak VRAM
none (baseline)	wrong	1.00x	4,220 MiB
hand the op a private copy of the captured tensors on the forward	correct	1.03x	19,720 MiB
pageable instead of page-locked host memory	correct	7.41x slower	unchanged

(Real Qwen2.5-32B NF4 for the first two rows; the pinning row is the same effect
measured on a synthetic stack of the same shape.)

The de-aliasing row is the important one, and it is why this needs a fix
upstream rather than in our code. Three percent of throughput is nothing. But
19,720 MiB is approximately the entire NF4 store — because bitsandbytes holds
each layer's forward-time reference until the backward reaches that layer, every
private copy has to stay alive across the whole forward-to-backward span. Any
de-aliasing repair is therefore O(model), not O(window), by construction.
For a caller whose entire premise is that peak memory is bounded by one layer,
that is not a workaround; it is the feature deleted.

We also confirmed the other end of the same trade: making the buffer pool large
enough that nothing is ever recycled restores exact gradients, and "large enough"
turns out to be one buffer per layer — the whole model resident again.

The remaining option on our side is to not route the streamed weight through
MatMul4Bit at all — dequantise inside the checkpointed region and use a native
matmul. We checked what that would cost numerically: the gradient is bit-exact
across 423 of 423 (shape, M, seed) rows, worst max_abs exactly 0.0, which is
unsurprising since MatMul4Bit.backward is already dequantize_4bit followed by
a matmul. The forward differs only where the fused kernel genuinely runs, and
then by up to one bf16 ulp. It is a workable escape hatch for us, but it means
opting out of the fused kernel to avoid an autograd bookkeeping detail, which
seems like the wrong place for the fix to live.

Happy to test a patch on the setup above, or to provide the fuller measurement
record if it is useful.