[anh-nguyenn]

Fixes #2110

Summary

All three k-bit matmul autograd functions (MatMul8bitLt, MatMul8bitFp, MatMul4Bit) computed the bias gradient as grad_output.sum(0). For 1-D inputs this returns a scalar, so backward raised:

Function MatMul4BitBackward returned an invalid gradient at index 2
  - got [] but expected shape compatible with [64]

For inputs with rank > 2 it returned a tensor with extra leading dimensions instead of an (out_features,) vector; autograd happened to reduce it further during accumulation, but the returned gradient had the wrong shape.

This reduces over every leading dimension explicitly:

grad_bias = grad_output.reshape(-1, grad_output.shape[-1]).sum(0, dtype=ctx.dtype_bias)

This is identical for 2-D inputs, makes 1-D inputs work, and returns the documented vector shape for 3-D/4-D inputs.

Verification

Added tests/test_modules.py::test_kbit_bias_grad_non_2d, parametrized over Linear8bitLt(has_fp16_weights=False) and LinearNF4, with input shapes (32,), (2, 3, 32), and (2, 2, 2, 32). It compares the bias gradient against torch.nn.Linear using a random upstream gradient.

Before the fix:

BNB_TEST_DEVICE=cpu  ... test_kbit_bias_grad_non_2d   # 2 failed, 4 passed
BNB_TEST_DEVICE=cuda ... test_kbit_bias_grad_non_2d   # 1 failed, 2 passed, 3 skipped
# both 1-D cases fail with "invalid gradient at index 2 - got []"

After the fix:

BNB_TEST_DEVICE=cpu  ... test_kbit_bias_grad_non_2d   # 6 passed
BNB_TEST_DEVICE=cuda ... test_kbit_bias_grad_non_2d   # 3 passed, 3 skipped
# (CUDA int8 1-D forward is a separate pre-existing limitation; covered on CPU)

Regression runs:

BNB_TEST_DEVICE=cpu pytest -q tests/test_autograd.py tests/test_modules.py \
  -k "bias_grad_non_2d or kbit_backprop or matmul_4bit or matmullt"
# 792 passed, 71 deselected

BNB_TEST_DEVICE=cuda pytest -q tests/test_autograd.py -k "matmul_4bit or matmullt"
# 768 passed

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

CPU tests were run on a host with an NVIDIA RTX A5000 (SM86). The CUDA regression runs above were executed on that A5000 and cover the 4-bit bias-gradient path; the CUDA 1-D case is now fixed and the existing 2-D/3-D/4-D autograd tests pass.

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, CUDA, x64 CPU.