Use torch.get_autocast_dtype instead of torch.get_autocast_gpu_dtype (#42055)

Update dtype handling for PyTorch 2.4 compatibility in flash attention models

Author: qgallouedec

src/transformers/integrations/flash_attention.py

@@ -15,7 +15,12 @@ def get_target_dtype(query: torch.Tensor, module: torch.nn.Module) -> torch.dtyp
     """If the query is in float32, return a target dtype compatible with flash attention. Return None otherwise."""
     if query.dtype == torch.float32:
         if torch.is_autocast_enabled():
-            return torch.get_autocast_gpu_dtype()
+            # NOTE: `torch.get_autocast_dtype` is there starting from PyTorch 2.4
+            return (
+                torch.get_autocast_dtype("cuda")
+                if hasattr(torch, "get_autocast_dtype")
+                else torch.get_autocast_gpu_dtype()
+            )
         # Handle the case where the model is quantized
         elif hasattr(module.config, "_pre_quantization_dtype"):
             return module.config._pre_quantization_dtype

src/transformers/kernels/falcon_mamba/selective_scan_with_ln_interface.py

@@ -250,12 +250,16 @@ def forward(
         delta_rank = delta_proj_weight.shape[1]
         d_state = A.shape[-1] * (1 if not A.is_complex() else 2)
         if torch.is_autocast_enabled():
-            x_proj_weight = x_proj_weight.to(dtype=torch.get_autocast_gpu_dtype())
-            delta_proj_weight = delta_proj_weight.to(dtype=torch.get_autocast_gpu_dtype())
-            out_proj_weight = out_proj_weight.to(dtype=torch.get_autocast_gpu_dtype())
-            out_proj_bias = (
-                out_proj_bias.to(dtype=torch.get_autocast_gpu_dtype()) if out_proj_bias is not None else None
+            # NOTE: `torch.get_autocast_dtype` is there starting from PyTorch 2.4
+            target_dtype = (
+                torch.get_autocast_dtype("cuda")
+                if hasattr(torch, "get_autocast_dtype")
+                else torch.get_autocast_gpu_dtype()
             )
+            x_proj_weight = x_proj_weight.to(dtype=target_dtype)
+            delta_proj_weight = delta_proj_weight.to(dtype=target_dtype)
+            out_proj_weight = out_proj_weight.to(dtype=target_dtype)
+            out_proj_bias = out_proj_bias.to(dtype=target_dtype) if out_proj_bias is not None else None
         if xz.stride(-1) != 1:
             xz = xz.contiguous()
         conv1d_weight = rearrange(conv1d_weight, "d 1 w -> d w")

src/transformers/models/diffllama/modeling_diffllama.py

@@ -353,6 +353,7 @@ def forward(
         device_type = query_states.device.type if query_states.device.type != "mps" else "cpu"
         if input_dtype == torch.float32:
             if torch.is_autocast_enabled():
+                # NOTE: `torch.get_autocast_dtype` is there starting from PyTorch 2.4
                 target_dtype = (
                     torch.get_autocast_dtype(device_type)
                     if hasattr(torch, "get_autocast_dtype")

src/transformers/models/diffllama/modular_diffllama.py

@@ -229,6 +229,7 @@ def forward(
         device_type = query_states.device.type if query_states.device.type != "mps" else "cpu"
         if input_dtype == torch.float32:
             if torch.is_autocast_enabled():
+                # NOTE: `torch.get_autocast_dtype` is there starting from PyTorch 2.4
                 target_dtype = (
                     torch.get_autocast_dtype(device_type)
                     if hasattr(torch, "get_autocast_dtype")

src/transformers/models/esm/modeling_esmfold.py

@@ -137,6 +137,7 @@ class EsmForProteinFoldingOutput(ModelOutput):
 
 def is_fp16_enabled(device_type):
     # Autocast world
+    # NOTE: `torch.get_autocast_dtype` is there starting from PyTorch 2.4
     autocast_dtype = (
         torch.get_autocast_dtype(device_type)
         if hasattr(torch, "get_autocast_dtype")

src/transformers/models/falcon/modeling_falcon.py

@@ -513,6 +513,7 @@ def forward(
         device_type = query_layer.device.type if query_layer.device.type != "mps" else "cpu"
         if input_dtype == torch.float32:
             if torch.is_autocast_enabled():
+                # NOTE: `torch.get_autocast_dtype` is there starting from PyTorch 2.4
                 target_dtype = (
                     torch.get_autocast_dtype(device_type)
                     if hasattr(torch, "get_autocast_dtype")

src/transformers/models/gpt_neo/modeling_gpt_neo.py

@@ -230,6 +230,7 @@ def forward(
         device_type = query.device.type if query.device.type != "mps" else "cpu"
         if query.dtype == torch.float32:
             if torch.is_autocast_enabled():
+                # NOTE: `torch.get_autocast_dtype` is there starting from PyTorch 2.4
                 target_dtype = (
                     torch.get_autocast_dtype(device_type)
                     if hasattr(torch, "get_autocast_dtype")

src/transformers/models/gptj/modeling_gptj.py

@@ -327,6 +327,7 @@ def forward(
         device_type = query.device.type if query.device.type != "mps" else "cpu"
         if input_dtype == torch.float32:
             if torch.is_autocast_enabled():
+                # NOTE: `torch.get_autocast_dtype` is there starting from PyTorch 2.4
                 target_dtype = (
                     torch.get_autocast_dtype(device_type)
                     if hasattr(torch, "get_autocast_dtype")

src/transformers/models/kyutai_speech_to_text/modeling_kyutai_speech_to_text.py

@@ -592,6 +592,7 @@ def forward(
         device_type = query_states.device.type if query_states.device.type != "mps" else "cpu"
         if input_dtype == torch.float32:
             if torch.is_autocast_enabled():
+                # NOTE: `torch.get_autocast_dtype` is there starting from PyTorch 2.4
                 target_dtype = (
                     torch.get_autocast_dtype(device_type)
                     if hasattr(torch, "get_autocast_dtype")

src/transformers/models/mimi/modeling_mimi.py

@@ -806,6 +806,7 @@ def forward(
         device_type = query_states.device.type if query_states.device.type != "mps" else "cpu"
         if input_dtype == torch.float32:
             if torch.is_autocast_enabled():
+                # NOTE: `torch.get_autocast_dtype` is there starting from PyTorch 2.4
                 target_dtype = (
                     torch.get_autocast_dtype(device_type)
                     if hasattr(torch, "get_autocast_dtype")