Update attn_mask_utils.py

attn_mask_utils.py

@@ -1,7 +1,19 @@
 from typing import List, Optional, Tuple, Union
 import torch
+from packaging import version
+import importlib.metadata
 from transformers.modeling_attn_mask_utils import AttentionMaskConverter
 
+from transformers.utils.import_utils import _is_package_available
+
+def is_transformers_attn_greater_or_equal_4_39():
+    if not _is_package_available("transformers"):
+        return False
+
+    return version.parse(importlib.metadata.version("transformers")) >= version.parse(
+        "4.39.0"
+    )
+
 def _prepare_4d_attention_mask_for_sdpa(
     attention_mask: Optional[torch.Tensor],
     input_shape: Union[torch.Size, Tuple, List],
@@ -59,9 +71,14 @@ def _prepare_4d_attention_mask_for_sdpa(
         # From PyTorch 2.1 onwards, F.scaled_dot_product_attention with the memory-efficient attention backend
         # produces nans if sequences are completely unattended in the attention mask. Details: https://github.com/pytorch/pytorch/issues/110213
         if query_length > 1:
-            expanded_4d_mask = AttentionMaskConverter._unmask_unattended(
-                expanded_4d_mask, attention_mask, unmasked_value=0.0
-            )
+            if is_transformers_attn_greater_or_equal_4_39():
+                expanded_4d_mask = AttentionMaskConverter._unmask_unattended(
+                    expanded_4d_mask, min_dtype=torch.finfo(inputs_embeds.dtype).min
+                )
+            else:
+                expanded_4d_mask = AttentionMaskConverter._unmask_unattended(
+                    expanded_4d_mask, attention_mask, unmasked_value=0.0
+                )
 
     return expanded_4d_mask
 
@@ -195,8 +212,13 @@ def _prepare_4d_causal_attention_mask_for_sdpa(
         # controlflow that can not be captured properly.
         # TODO: _unmask_unattended does not work either with torch.compile when using fullgraph=True. We should find a way to detect this case.
         if query_length > 1 and not is_tracing:
-            expanded_4d_mask = AttentionMaskConverter._unmask_unattended(
-                expanded_4d_mask, attention_mask, unmasked_value=0.0
-            )
+            if is_transformers_attn_greater_or_equal_4_39():
+                expanded_4d_mask = AttentionMaskConverter._unmask_unattended(
+                    expanded_4d_mask, min_dtype=torch.finfo(inputs_embeds.dtype).min
+                )
+            else:
+                expanded_4d_mask = AttentionMaskConverter._unmask_unattended(
+                    expanded_4d_mask, attention_mask, unmasked_value=0.0
+                )
 
-    return expanded_4d_mask
+    return expanded_4d_mask