FEAT / Bitsandbytes: Add dequantize API for bitsandbytes quantized models

[younesbelkada]

What does this PR do?

Fixes #30177

This PR adds a new feature dequantize in order to de-quantize models for interesting usecases such as the one described in #30177

The API is very simple:

from transformers import AutoModelForCausalLM, BitsAndBytesConfig, AutoTokenizer

model_id = "facebook/opt-125m"

model = AutoModelForCausalLM.from_pretrained(model_id, quantization_config=BitsAndBytesConfig(load_in_4bit=True))
tokenizer = AutoTokenizer.from_pretrained(model_id)

model.dequantize()

text = tokenizer("Hello my name is", return_tensors="pt").to(0)

out = model.generate(**text)
print(tokenizer.decode(out[0]))

Users just need to make sure they have enough GPU RAM in order to store the unquantized model, otherwise they might face unexpected behaviour

Added the support for 4-bit / 8-bit models and nice tests + docs to educate users on how to use this new API.

cc @amyeroberts @SunMarc

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[SunMarc]

Thanks for adding this new method in quantizer ! This will make fine-tuning with quantized model way easier ! I left a few minor comments.

[SunMarc]

The user might want to know in which precision the model was dequantized since they don't have the possibility to control that. I think it could be great to give that information since there is no default value (as opposed to from_pretrained which loads the model in fp32).
Two ways to get that:

• just check the dtype of the weights at the end ( potentially the easiest way )
• check what happens in dequantize_4bit . In the method, you see that they get the output dtype with weight.quant_state.dtype.

We can potentially add a torch_dtype attribute in the future if it makes sense.

[amyeroberts]

+1

[younesbelkada]

Nice catch! The output dtype should be correctly inferred here: https://github.com/TimDettmers/bitsandbytes/blob/b891f80ba514833f41f0e9226983b02a9fb5c44b/bitsandbytes/functional.py#L1349 through the compute_dtype so it should be accurate - I added a warning_once staement to inform users on the dequantized dtype: 1a4a906

[amyeroberts]

Thanks for adding this! +1 on all of @SunMarc's comments.

[amyeroberts]

+1

[amyeroberts]

This is already imported at the top of the module

[younesbelkada]

Nice catch ! Should be fixed now

[amyeroberts]

Thanks for adding this feature and iterating!

[amyeroberts]

One general comment, if instead you could have a private method _dequantize_and_replace, which handles the recursion, you don't need to return has_been_replaced here. When someone calls dequantize_and_replace, I don't think has_been_replaced is ever used and could be confusing e.g.:

# This is just dequantize_and_replace from before
def _dequantize_and_replace(
    model,
    modules_to_not_convert=None,
    current_key_name=None,
    quantization_config=None,
    has_been_replaced=False,
):
    ...
    return model, has_been_replaced

def dequantize_and_replace(
    model,
    modules_to_not_convert=None,
    current_key_name=None,
    quantization_config=None,
    has_been_replaced=False,
):
    model, has_been_replaced = _dequantize_and_replace(...)
    return model 

[younesbelkada]

makes sense ! Will do !

[younesbelkada]

Done in 8b904f7 !

[amyeroberts]

Nice :)

[RonanKMcGovern]

Yeah this is great, thanks

[younesbelkada]

Great thanks @RonanKMcGovern ! let us know how it goes