[anh-nguyenn]

Fixes #2108

Summary

MatMul8bitFp.backward() computes the weight gradient as torch.matmul(A.t(), grad_output).t(). A.t() only accepts tensors with rank <= 2, so CPU/XPU training with Linear8bitLt crashes for 1-D, 3-D, or 4-D inputs even though the forward pass (via F.linear) accepts them. grad_output was also only flattened for exactly 3-D inputs, so 1-D activations reached the matmul with mismatched ranks.

This flattens any non-2-D grad_output and reshapes A to 2-D before computing the weight gradient. The input-gradient path was already rank-independent and is unchanged.

Verification

Added tests/test_modules.py::test_linear8bitlt_cpu_backprop_non_2d, covering (32,), (2, 3, 32), and (2, 2, 2, 32) and comparing both input and weight gradients against torch.nn.Linear.

Before the fix:

BNB_TEST_DEVICE=cpu pytest -q tests/test_modules.py::test_linear8bitlt_cpu_backprop_non_2d
# 3 failed
RuntimeError: t() expects a tensor with <= 2 dimensions, but self is 3D  (and 4D)
RuntimeError: inconsistent tensor size, expected tensor [32] and src [64]  (1D)

After the fix:

BNB_TEST_DEVICE=cpu pytest -q tests/test_modules.py::test_linear8bitlt_cpu_backprop_non_2d
# 3 passed

BNB_TEST_DEVICE=cpu pytest -q tests/test_modules.py \
  -k "linear8bitlt_cpu_backprop_non_2d or kbit_backprop or linear8bitlt_inference or linear8bitlt_no_fp16_weights"
# 25 passed, 67 deselected

ruff check bitsandbytes/autograd/_functions.py tests/test_modules.py
ruff format --check bitsandbytes/autograd/_functions.py tests/test_modules.py
# All checks passed; 2 files already formatted

GPU validation

This change only touches the CPU/XPU MatMul8bitFp path; CUDA uses MatMul8bitLt and is unaffected. The regression test runs on CPU, where the bug reproduces. The full CPU test run above was performed on an NVIDIA RTX A5000 host (SM86) using the CPU device.

AI assistance

I used an AI assistant (opencode) to help investigate, reproduce, and implement this change. I reviewed the reproducer, the patch, and the test results before submitting.

Suggested labels: Bug, CPU.