Scope out autocast in extended atten mask gen (#287)

huggingface · Jul 10, 2023 · dbefacd · dbefacd
1 parent 32f8555
commit dbefacd
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/optimum/habana/transformers/models/albert/modeling_albert.py b/optimum/habana/transformers/models/albert/modeling_albert.py
@@ -69,7 +69,8 @@ def gaudi_albert_forward(
     # torch.finfo must take the dtype of encoder_extended_attention_mask
     extended_attention_mask = extended_attention_mask.to(dtype=self.dtype)  # bf16 compatibility
     extended_attention_mask = 1.0 - extended_attention_mask
-    extended_attention_mask = extended_attention_mask * torch.finfo(extended_attention_mask.dtype).min
+    with torch.autocast(enabled=False, device_type="hpu"):
+        extended_attention_mask = extended_attention_mask * torch.finfo(extended_attention_mask.dtype).min
     head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)
 
     embedding_output = self.embeddings(

diff --git a/optimum/habana/transformers/models/modeling_all_models.py b/optimum/habana/transformers/models/modeling_all_models.py
@@ -90,7 +90,8 @@ def gaudi_get_extended_attention_mask(
     # torch.finfo must take the dtype of encoder_extended_attention_mask
     extended_attention_mask = extended_attention_mask.to(dtype=dtype)  # bf16 compatibility
     extended_attention_mask = 1.0 - extended_attention_mask
-    extended_attention_mask = extended_attention_mask * torch.finfo(extended_attention_mask.dtype).min
+    with torch.autocast(enabled=False, device_type="hpu"):
+        extended_attention_mask = extended_attention_mask * torch.finfo(extended_attention_mask.dtype).min
 
     return extended_attention_mask