fix(autograd/MatMul4Bit): save the packed weight via save_for_backward (#2034)

[Anai-Guo]

What

MatMul4Bit.forward saves the packed 4-bit weight as a plain ctx attribute
(ctx.tensors = (None, B)) instead of through ctx.save_for_backward(B), and
backward reads it back from ctx.tensors. This PR switches to
save_for_backward / ctx.saved_tensors. The QuantState stays on ctx.state
because it is not a tensor.

Fixes #2034.

Why

In the ordinary case the ctx-attribute approach is harmless — the weight is a
stable Parameter, so holding a reference to it is exactly as correct as saving
it, and costs nothing.

It stops being correct 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 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
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.

bf16 through the same buffer-recycling harness is unaffected, because it goes
through MmBackward0, which already uses save_for_backward. This change brings
MatMul4Bit in line with that and with PyTorch's documented mechanism for saving
tensors that must survive a checkpoint recompute.

Change

# forward
if any(ctx.needs_input_grad[:2]):
    ctx.save_for_backward(B)      # was: ctx.tensors = (None, B)
else:
    ctx.save_for_backward()       # was: ctx.tensors = (None, None)

# backward
saved = ctx.saved_tensors         # was: _, B = ctx.tensors
B = saved[0] if saved else None

No behavioural change in the ordinary training path (stable Parameter weight,
no in-place mutation between forward and backward → no version-counter interaction).

Testing

The failure only manifests with a CUDA 4-bit weight re-materialised inside a
torch.utils.checkpoint(use_reentrant=False) region, so it requires a GPU. The
issue includes a self-contained reproducer (three arms: recycled-buffer nf4,
private-buffer nf4 reference, recycled-buffer bf16 control) that I did not run
here for lack of the H100/CUDA setup it targets. Happy to add a CUDA-gated
regression test mirroring that harness if you'd like it in-tree.

🤖 Generated with Claude Code