CodeGen Fix causal mask for half precision

[younesbelkada]

What does this PR do?

This PR forces the causal mask to stay in torch.uint8. An error occurs when loading a model in half precision since torch_dtype=torch.float16 casts also the buffers in fp16. Here is a minimal script to reproduce the error:

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("Salesforce/codegen-2B-mono")
model = AutoModelForCausalLM.from_pretrained("Salesforce/codegen-2B-mono", device_map="auto", torch_dtype=torch.float16)

text = "def quicksort(l):"

encoded_input = tokenizer(text, return_tensors='pt')
output_sequences = model.generate(input_ids=encoded_input['input_ids'], attention_mask=encoded_input['attention_mask'])
print(tokenizer.decode(output_sequences[0], skip_special_tokens=True))

In a future PR we could address non-casting the buffers (aka keeping them in their native dtype)

Can also confirm the slow tests pass!

cc @ydshieh

[HuggingFaceDocBuilderDev]

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

[ydshieh]

Let's have a comment here to explain why we need .to(torch.uint8) 🙏

[younesbelkada]

Added more comments on 8b81ac1
Let me know if anything is unclear!

[thomasw21]

Stupid question, why is that not

Suggested change

	causal_mask = self.causal_mask[:, :, key_length - query_length : key_length, :key_length].to(torch.uint8)
	self.causal_mask = self.causal_mask[:, :, key_length - query_length : key_length, :key_length].to(torch.uint8)

feels like something you want to do only once no?

It's a no-op when the tensor is already in the correct dtype.

[sgugger]

It won't be needed anymore as we found the root cause with Younes :-)

[sgugger]

LGTM, thanks for fixing!

[sgugger]

Let's close this one to focus on the right fix @younesbelkada :-)

[younesbelkada]

Yeah let's move the discussion to: #18471