From 37bc76d782cc0b4ddfc35225dff85c9e7b0bd1a1 Mon Sep 17 00:00:00 2001 From: younesbelkada Date: Thu, 4 Aug 2022 10:15:41 +0200 Subject: [PATCH 1/2] hotfix causal mask - Small hotfix for causal mask for half-precision models - Explicitly cast the causal mask to uint8 for compatibiliy with `torch.where` --- src/transformers/models/codegen/modeling_codegen.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/transformers/models/codegen/modeling_codegen.py b/src/transformers/models/codegen/modeling_codegen.py index 06581e732cdc..40e877f1bb6e 100644 --- a/src/transformers/models/codegen/modeling_codegen.py +++ b/src/transformers/models/codegen/modeling_codegen.py @@ -151,7 +151,7 @@ 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] + causal_mask = self.causal_mask[:, :, key_length - query_length : key_length, :key_length].to(torch.uint8) # Keep the attention weights computation in fp32 to avoid overflow issues query = query.to(torch.float32) From 8b81ac12339692bb3bbddc363ddcb9438eeb3e68 Mon Sep 17 00:00:00 2001 From: younesbelkada Date: Thu, 4 Aug 2022 11:46:54 +0200 Subject: [PATCH 2/2] add more comments on casting --- src/transformers/models/codegen/modeling_codegen.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/src/transformers/models/codegen/modeling_codegen.py b/src/transformers/models/codegen/modeling_codegen.py index 40e877f1bb6e..d6f8693a83b2 100644 --- a/src/transformers/models/codegen/modeling_codegen.py +++ b/src/transformers/models/codegen/modeling_codegen.py @@ -151,6 +151,10 @@ def _attn( # compute causal mask from causal mask buffer query_length, key_length = query.size(-2), key.size(-2) + + # 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) # Keep the attention weights computation in fp32 to avoid overflow issues