[BigModeling] Final fix for dispatch int8 and fp4 models

[younesbelkada]

What does this PR do?

Fixes a silent bug regarding buffers that are not in the state dict of the model and when using quantized models. Currently the following snippet fails:

from transformers import AutoModelForCausalLM, AutoConfig, AutoTokenizer
import torch

model_path="ybelkada/gpt2-xl-8bit"

config = AutoConfig.from_pretrained(model_path)
model = AutoModelForCausalLM.from_pretrained(model_path, load_in_8bit=True)

tokenizer = AutoTokenizer.from_pretrained(model_path)

input_text = "Describe the solar system."
input_ids = tokenizer(input_text, return_tensors="pt").input_ids.to("cuda")

outputs = model.generate(input_ids, max_length=10)
print(tokenizer.decode(outputs[0]))

With the following error:

  File "/home/younes_huggingface_co/code/transformers/src/transformers/models/gpt2/modeling_gpt2.py", line 203, in _attn
    attn_weights = torch.where(causal_mask, attn_weights.to(attn_weights.dtype), mask_value)
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!

Because in #1652 we introduced a check to do nothing in case the model is quantized. This leads to buffers that are not in the state dict being on CPU.

As the .to operation is not supported for quantized models, the fix is to always force-dispatch the model if the model is quantized.

cc @sgugger @SunMarc

Currently the corresponding test is implemented in huggingface/transformers#24543 as it relies on the main branch of transformers.

[sgugger]

Thanks for the fix!

[HuggingFaceDocBuilderDev]

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

[SunMarc]

LGTM !