Linear8bitLt CPU/XPU backward crashes for non-2D inputs in MatMul8bitFp #2108

System Info

• bitsandbytes source at 8336490 (0.50.3.dev0)
• Debian Linux x86_64, Python 3.12.3, PyTorch 2.8.0+cu128
• CPU repro also relevant to XPU, which uses the same MatMul8bitFp path

Summary

Training bnb.nn.Linear8bitLt on CPU (and XPU) crashes for input tensors that are not 2-D. MatMul8bitFp.backward() computes the weight gradient with torch.matmul(A.t(), grad_output).t(), but A.t() only works for <= 2 dimensions. The forward pass supports arbitrary ranks because it uses F.linear, so 1-D, 3-D and 4-D inputs fail only during backward.

The same code path is also wrong for 1-D inputs: grad_output is only flattened for exactly 3-D, so a 1-D activation reaches a 2-D matmul with mismatched ranks.

Reproduction

import torch
import bitsandbytes as bnb

layer = bnb.nn.Linear8bitLt(32, 64, has_fp16_weights=True, bias=False).to("cpu")
x = torch.randn(2, 3, 32, requires_grad=True)  # also fails for (32,) and (2, 2, 2, 32)
layer(x).sum().backward()

RuntimeError: t() expects a tensor with <= 2 dimensions, but self is 3D
  File "bitsandbytes/autograd/_functions.py", line 288, in backward
    grad_B = torch.matmul(A.t(), grad_output).t()

For (32,) the failure is:

RuntimeError: inconsistent tensor size, expected tensor [32] and src [64]

Expected behavior

Backward should work for the same input ranks that nn.Linear and the 8-bit forward pass accept, and produce gradients matching a float reference layer (within int8 quantization error).

Notes

• CUDA uses MatMul8bitLt, not MatMul8bitFp, so this is CPU/XPU-specific.
• has_fp16_weights=False with frozen weights does not hit req_gradB, so it can hide the issue.
• Suggested labels: Bug, CPU.