Skip to content

CodeGen Fix causal mask for half precision #18467

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 2 commits into from
Closed

CodeGen Fix causal mask for half precision #18467

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
37bc76d
hotfix causal mask
younesbelkada Aug 4, 2022
8b81ac1
add more comments on casting
younesbelkada Aug 4, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 5 additions & 1 deletion src/transformers/models/codegen/modeling_codegen.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -151,7 +151,11 @@ def _attn(

# compute causal mask from causal mask buffer
query_length, key_length = query.size(-2), key.size(-2)
causal_mask = self.causal_mask[:, :, key_length - query_length : key_length, :key_length]

# Here we force cast the causal mask to uint8 to avoid errors related to torch.where
# combined with torch_dtype="auto" where it casts all variables including buffers to
# fp16. See the related issue here: https://github.com/huggingface/transformers/pull/18467
causal_mask = self.causal_mask[:, :, key_length - query_length : key_length, :key_length].to(torch.uint8)
Copy link
Collaborator

@ydshieh ydshieh Aug 4, 2022

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Contributor Author

@younesbelkada younesbelkada Aug 4, 2022

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Contributor

@thomasw21 thomasw21 Aug 4, 2022 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Collaborator

@sgugger sgugger Aug 4, 2022 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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


# Keep the attention weights computation in fp32 to avoid overflow issues
query = query.to(torch.float32)
Expand Down