[bnb] Add fp4 support for dispatch

[younesbelkada]

What does this PR do?

Fixes #1504

This PR applies the similar enhancement as #1228 for FP4 layers

Now the script below outputs the desired dtype:

import torch
from transformers import AutoConfig, AutoModel, BitsAndBytesConfig
from transformers.utils.bitsandbytes import replace_8bit_linear

from accelerate import init_empty_weights
from huggingface_hub import hf_hub_download

from accelerate import load_checkpoint_and_dispatch


with init_empty_weights():
    model = AutoModel.from_config(AutoConfig.from_pretrained("bigscience/bloom-560m"))

quantization_config = BitsAndBytesConfig(
    load_in_4bit=True
)

model = replace_8bit_linear(model, quantization_config=quantization_config)

# For some reason replace_8bit_linear creates parameters with requires_grad=True but it's irrelevant rn
for p in model.parameters():
    p.requires_grad = False

model_path = hf_hub_download("bigscience/bloom-560m", "pytorch_model.bin")

model = load_checkpoint_and_dispatch(
    model,
    checkpoint=model_path,
    device_map="auto",
)

print(f"{model.h[0].self_attention.query_key_value.weight.device}\n{model.h[0].self_attention.query_key_value.weight.dtype}")
>>> torch.uint8

cc @sgugger @BlackSamorez

[sgugger]

The .cuda function is very very deprecated. You should use to.

[younesbelkada]

Hmm I think the way bitsandbytes has designed its Linear4bit layers we need to call cuda: https://github.com/TimDettmers/bitsandbytes/blob/ac5550a0238286377ee3f58a85aeba1c40493e17/bitsandbytes/nn/modules.py#L152 it seems to be the only way to quantize the weights :/ I tried it with to and it didn't worked. (note that at that point module.weight is a bnb.nn.Params4bit module)

[sgugger]

Oh ok. Not very PyTorch-ic then.

[younesbelkada]

yeah ! :/

[HuggingFaceDocBuilderDev]

The documentation is not available anymore as the PR was closed or merged.

[sgugger]

Isn't this function renamed to something else?

[younesbelkada]

ah yes, let me modify that

[sgugger]

To adapt.