From 7eea6dfa38ff4f12bac47f97b0b022023effbedd Mon Sep 17 00:00:00 2001 From: Tai An Date: Fri, 7 Aug 2026 12:11:26 -0700 Subject: [PATCH] fix(autograd/MatMul4Bit): save the packed weight via save_for_backward (#2034) MatMul4Bit.forward stored the packed weight as a plain ctx attribute (ctx.tensors = (None, B)) instead of through save_for_backward. In the ordinary case this is harmless -- the weight is a stable Parameter. It breaks for a caller that re-materialises the weight (weight offloading, layer streaming) into a recycled device buffer under gradient checkpointing: torch.utils.checkpoint discards and recomputes *saved* tensors, so the recompute is expected to hand backward a fresh copy of the weight. A raw ctx attribute is invisible to that mechanism -- the reference taken in the original forward survives into backward and, by then, points at a buffer refilled with a different layer, producing silently wrong gradients (forward stays bit-exact). bf16 through the same harness is unaffected because MmBackward0 already uses save_for_backward. Switch to ctx.save_for_backward(B) in forward and read it back from ctx.saved_tensors in backward. The QuantState stays on ctx.state since it is not a tensor. No behavioural change in the ordinary path. --- bitsandbytes/autograd/_functions.py | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) diff --git a/bitsandbytes/autograd/_functions.py b/bitsandbytes/autograd/_functions.py index 8a069bd10..286e1b97f 100644 --- a/bitsandbytes/autograd/_functions.py +++ b/bitsandbytes/autograd/_functions.py @@ -354,10 +354,16 @@ def forward(ctx, A, B, out=None, bias=None, quant_state: Optional[F.QuantState] ctx.state = quant_state ctx.dtype_A, ctx.dtype_B, ctx.dtype_bias = A.dtype, B.dtype, None if bias is None else bias.dtype + # Save the packed weight through save_for_backward rather than as a plain + # ctx attribute. torch.utils.checkpoint discards and recomputes *saved* + # tensors, so a caller that re-materialises the weight (weight offloading, + # layer streaming) into a recycled buffer under checkpointing gets a fresh + # copy in backward. A raw ctx attribute is invisible to that mechanism and + # would keep pointing at the original buffer -- silently wrong gradients. if any(ctx.needs_input_grad[:2]): - ctx.tensors = (None, B) + ctx.save_for_backward(B) else: - ctx.tensors = (None, None) + ctx.save_for_backward() return output @@ -368,7 +374,8 @@ def backward(ctx, grad_output): return torch.zeros_like(ctx.A), torch.zeros_like(ctx.B), None, bias_grad, None req_gradA, _, _, req_gradBias, _ = ctx.needs_input_grad - _, B = ctx.tensors + saved = ctx.saved_tensors + B = saved[0] if saved else None grad_A, grad_B, grad_bias = None, None, None @@ -379,7 +386,7 @@ def backward(ctx, grad_output): # not supported by PyTorch. TODO: create work-around # if req_gradB: grad_B = torch.matmul(grad_output.t(), A) if req_gradA: - # B in ctx.tensors is already in canonical [(N*K+1)//2, 1] form (normalized in forward). + # B from saved_tensors is already in canonical [(N*K+1)//2, 1] form (normalized in forward). # dequantize returns [N, K]; matmul(grad_output[M,N], [N,K]) = grad_A[M,K]. grad_A = torch.matmul(grad_output, F.dequantize_4bit(B, ctx.state).to(grad_output.dtype))