Add XPU case to is_torch_bf16_gpu_available (#37132)

* Add xpu case to is_torch_bf16_gpu_available

Signed-off-by: cyy <cyyever@outlook.com>

* Refine error messages

Signed-off-by: cyy <cyyever@outlook.com>

---------

Signed-off-by: cyy <cyyever@outlook.com>
Co-authored-by: Marc Sun <57196510+SunMarc@users.noreply.github.com>

Author: cyyever

src/transformers/training_args.py

@@ -45,6 +45,7 @@
     is_sagemaker_mp_enabled,
     is_torch_available,
     is_torch_bf16_gpu_available,
+    is_torch_cuda_available,
     is_torch_hpu_available,
     is_torch_mlu_available,
     is_torch_mps_available,
@@ -1683,11 +1684,12 @@ def __post_init__(self):
                     # cpu
                     raise ValueError("Your setup doesn't support bf16/(cpu, tpu, neuroncore). You need torch>=1.10")
                 elif not self.use_cpu:
-                    if torch.cuda.is_available() and not is_torch_bf16_gpu_available():
+                    if not is_torch_bf16_gpu_available():
+                        error_message = "Your setup doesn't support bf16/gpu."
+                        if is_torch_cuda_available():
+                            error_message += " You need Ampere+ GPU with cuda>=11.0"
                         # gpu
-                        raise ValueError(
-                            "Your setup doesn't support bf16/gpu. You need torch>=1.10, using Ampere GPU with cuda>=11.0"
-                        )
+                        raise ValueError(error_message)
 
         if self.fp16 and self.bf16:
             raise ValueError("At most one of fp16 and bf16 can be True, but not both")

src/transformers/utils/import_utils.py

@@ -512,13 +512,17 @@ def is_torch_mps_available(min_version: Optional[str] = None):
     return False
 
 
-def is_torch_bf16_gpu_available():
+def is_torch_bf16_gpu_available() -> bool:
     if not is_torch_available():
         return False
 
     import torch
 
-    return torch.cuda.is_available() and torch.cuda.is_bf16_supported()
+    if torch.cuda.is_available():
+        return torch.cuda.is_bf16_supported()
+    if torch.xpu.is_available():
+        return torch.xpu.is_bf16_supported()
+    return False
 
 
 def is_torch_bf16_cpu_available() -> bool: