[Feature request] Support bias in linear_nf4

[ebsmothers]

In Llama 3.2 vision models, we load a pretrained checkpoint where some of the linear layers have biases in them (specifically in the vision encoder MLPs, see e.g. here). We would like to be able to quantize these linears as well, but linear_nf4 currently only operates on a weight, no bias (ref). I know we had some discussions about the use of bias in NF4 linears way back when, but curious if there are any hard technical blockers for supporting this in ao. If not, can we add this?

[gau-nernst]

I think you can just add bias to the outputs? i.e. out = linear_nf4(input, weight_nf4) + bias

[ebsmothers]

@gau-nernst thanks. So in that case we just keep bias in the higher precision all the time? I guess the memory implications of keeping the bias in bf16 aren't too terrible, but was wondering about quantizing the bias as well (not sure about implications for model accuracy, but to me it seems the intuitive thing to do).

[msaroufim]

It's actually quite common to keep the biases in high precision, see this popular blog as an example https://www.maartengrootendorst.com/blog/quantization/#weights-and-biases

[ebsmothers]

@msaroufim thanks, this is a really nice blog. Let me go this route for now then.

[pbontrager]

The blog post mostly says that the biases aren’t quantized (as heavily) because they are only a small fraction of the model, but it doesn’t explain why they’re needed to be in int16 instead of int4. Is their precision more important to model quality?

Also, while the solution above is great and works for us, I’m a bit concerned with how it diverges from the core library Linear, where there is a bias by default. Are these int tensors meant to be used as a drop in replacement for nn.Linear or are they more low level where a user is expected to know the above pattern?

[gau-nernst]

@ebsmothers @pbontrager If you want to quantize the biases, which I don't think it's worth the efforts, you can do something like this

from torchao.dtypes.nf4tensor import to_nf4, linear_nf4
import torch

inputs = torch.randn(4, 2048)
weight = to_nf4(torch.randn(2048, 2048))
bias = to_nf4(torch.randn(2048), scaler_block_size=1)

# bias.to(dtype) will implicitly dequantize
out = linear_nf4(inputs, weight) + bias.to(weight.dtype)

Note:

• Since bias is too small, you would need to adjust block_size and scalar_block_size accordingly
• I assume you don't need grad_bias, thus doing like this should be fine (i.e. no need custom torch.autograd.Function)