k-bit linear bias gradient is invalid for 1-D inputs (and wrong rank for >2-D) #2110

System Info

• bitsandbytes source at 8336490 (0.50.3.dev0)
• Debian Linux x86_64, Python 3.12.3, PyTorch 2.8.0+cu128
• Reproduced on CPU and on an NVIDIA RTX A5000 (SM86)

Summary

Linear4bit and Linear8bitLt fail during backward for 1-D input tensors when the bias requires a gradient. All three k-bit matmul autograd functions compute the bias gradient as grad_output.sum(0). For a 1-D output this is a scalar, so autograd rejects the returned gradient:

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

nn.Linear supports 1-D inputs, and the bitsandbytes forward pass already returns the correct 1-D output, so this only fails in backward. For ranks above 2, sum(0) also returns a tensor with extra leading dimensions instead of a vector; PyTorch happens to reduce it further during accumulation, but the returned gradient has the wrong documented shape.

Reproduction

import torch
import bitsandbytes as bnb

for layer in (
    bnb.nn.LinearNF4(32, 64),
    bnb.nn.Linear8bitLt(32, 64, has_fp16_weights=False),
):
    layer = layer.to("cpu")
    x = torch.randn(32, requires_grad=True)
    layer(x).sum().backward()

Both layers raise MatMul4BitBackward/MatMul8bitFpBackward returned an invalid gradient ... got [] but expected shape [64]. The same 4-bit failure occurs on CUDA.

Expected behavior

The bias gradient should be the reduction of grad_output over every leading dimension, i.e. grad_output.reshape(-1, grad_output.shape[-1]).sum(0), which is a (out_features,) vector for 1-D, 2-D, 3-D, and 4-D inputs.

Suggested labels

Bug, CUDA, x64 CPU.