Error with new LLaMA V2 in 4bit : mat1 and mat2 shapes cannot be multiplied

[dimaischenko]

Error with new Llama V2 model and 4bit inference:

bitsandbytes: 0.40.2
transformers: 4.31.0
torch: 2.0.1
GPU: RTX 3090 24 Gb
Cuda: 11.7

Load model:

model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-2-13b-hf",
    use_safetensors=True,
    load_in_4bit=True,
    device_map="auto",
)

Try to generate:

text = "Hello how are you?"
input_ids = tokenizer(text, return_tensors='pt').input_ids.cuda()
output = model.generate(inputs=input_ids, max_new_tokens=10)

Error:

File /usr/local/lib/python3.8/dist-packages/transformers/models/llama/modeling_llama.py:295, in <listcomp>(.0)
    292 key_slices = self.k_proj.weight.split(key_value_slicing, dim=0)
    293 value_slices = self.v_proj.weight.split(key_value_slicing, dim=0)
--> 295 query_states = [F.linear(hidden_states, query_slices[i]) for i in range(self.pretraining_tp)]
    296 query_states = torch.cat(query_states, dim=-1)
    298 key_states = [F.linear(hidden_states, key_slices[i]) for i in range(self.pretraining_tp)]

RuntimeError: mat1 and mat2 shapes cannot be multiplied (6x5120 and 1x2560)

[nazarov-yuriy]

The code is different for tp=1 and tp>1 models.
https://github.com/huggingface/transformers/blob/129cb6d5239374869beea0e1317641426e8df572/src/transformers/models/llama/modeling_llama.py#L289
7b model with tp=1 works.

[dimaischenko]

@nazarov-yuriy You are right! I manually changed the value of this parameter when loading the 13b model and everything worked.

model_id = "meta-llama/Llama-2-13b-hf"
config = AutoConfig.from_pretrained(model_id, pretraining_tp=1)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    load_in_4bit=True,
    device_map="auto",
    config=config,
)

I wonder if there is a big loss in quality with such changes, I read that the parameter is experimental:
https://huggingface.co/docs/transformers/main/model_doc/llama2#transformers.LlamaConfig.pretraining_tp

[TimDettmers]

Thank you, I will need to look into this. This is likely a changed needed in the bitsandbytes integration of the transformers stack.

[faizanahemad]

If it helps then using a different fp16 weights from TheBloke/Llama-2-70B-Chat-fp16 works. As below.

from transformers import AutoTokenizer, AutoModelForCausalLM, AutoConfig
from transformers import pipeline
# https://github.com/TimDettmers/bitsandbytes/issues/610
# repo_id = "meta-llama/Llama-2-70b-chat-hf"
repo_id = "TheBloke/Llama-2-70B-Chat-fp16"
config = AutoConfig.from_pretrained(repo_id, pretraining_tp=1)
model = AutoModelForCausalLM.from_pretrained(repo_id, device_map="auto", load_in_4bit=True, use_auth_token=True) # load_in_8bit=True, rope_scaling={"type": "dynamic", "factor": 2.0}, config=config
for p in model.parameters():
    p.requires_grad = False
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-70b-chat-hf", use_fast=True)
generate_kwargs = {"max_new_tokens":40, "do_sample":False}

question = "User:\nHow to kill a python process using terminal when the process refuses to shut down itself?\nAssistant:\n"
model_inputs = tokenizer(question, return_tensors="pt").to("cuda:0")
gen_out = model.generate(**model_inputs, **generate_kwargs)
print(tokenizer.decode(gen_out[0], skip_special_tokens=True))

@TimDettmers Not sure why these weights work but official llama ones don't but it seems there was some issue with how meta stored the weights.

[nazarov-yuriy]

Not sure why these weights work but official llama ones don't but it seems there was some issue with how meta stored the weights.

Official 70B weights doesn't work with transformers==4.31.0, but they work after huggingface/transformers#24891

[dimaischenko]

@TimDettmers @nazarov-yuriy Meta fixed pretraining_tp parameter to 1 for 13b model in their repo:

https://huggingface.co/meta-llama/Llama-2-13b-hf/blob/main/config.json

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.