Repository: huggingface/transformers
Pull request: enable flex attention ut cases on XPU

enable flex_attention test cases on XPU, it's supported by XPU torch now.

@ydshieh , pls help review, thx very much.

Changed files:
filename: tests/extended/test_trainer_ext.py
patch: @@ -259,8 +259,8 @@ def train_and_return_metrics(optim: str) -> tuple[int, float]:
             f" gpu_total_mem_bnb={gpu_total_mem_bnb}MB",
         )
 
-        self.assertEqual(
-            loss_orig, loss_bnb, f"loss should be the same, but got loss_orig={loss_orig}, loss_bnb={loss_bnb}"
+        self.assertAlmostEqual(
+            loss_orig, loss_bnb, 5, f"loss should be the same, but got loss_orig={loss_orig}, loss_bnb={loss_bnb}"
         )
 
     def run_trainer(

filename: tests/models/deepseek_v3/test_modeling_deepseek_v3.py
patch: @@ -25,7 +25,6 @@
     require_read_token,
     require_torch,
     require_torch_accelerator,
-    require_torch_gpu,
     require_torch_large_accelerator,
     slow,
     torch_device,
@@ -449,7 +448,7 @@ def test_eager_matches_sdpa_generate(self):
                     msg=f"\n{tokenizer.batch_decode(res_eager)} \nvs\n{tokenizer.batch_decode(res_sdpa)}",
                 )
 
-    @require_torch_gpu
+    @require_torch_accelerator
     def test_flex_attention_with_grads(self):
         """
         Overwriting as the namings/functionality on the attention part are different; for now it's more of a unique model.

filename: tests/models/t5gemma/test_modeling_t5gemma.py
patch: @@ -24,7 +24,6 @@
 from transformers.testing_utils import (
     require_torch,
     require_torch_accelerator,
-    require_torch_gpu,
     torch_device,
 )
 
@@ -1231,7 +1230,7 @@ def test_custom_4d_attention_mask(self):
 
     # Based on tests.test_modeling_common.ModelTesterMixin.test_flex_attention_with_grads
     # Update hidden size for encoder and decoder
-    @require_torch_gpu
+    @require_torch_accelerator
     def test_flex_attention_with_grads(self):
         for model_class in self.all_model_classes:
             # TODO: raushan, fix for composite models after making VLMs support new attn API
@@ -1516,7 +1515,7 @@ def test_training_gradient_checkpointing_use_reentrant_false(self):
 
     # Based on tests.test_modeling_common.ModelTesterMixin.test_flex_attention_with_grads
     # Update hidden size for encoder
-    @require_torch_gpu
+    @require_torch_accelerator
     def test_flex_attention_with_grads(self):
         for model_class in self.all_model_classes:
             # TODO: raushan, fix for composite models after making VLMs support new attn API

filename: tests/models/zamba2/test_modeling_zamba2.py
patch: @@ -26,6 +26,7 @@
     require_bitsandbytes,
     require_flash_attn,
     require_torch,
+    require_torch_accelerator,
     require_torch_gpu,
     slow,
     torch_device,
@@ -534,7 +535,7 @@ def test_flash_attn_2_fp32_ln(self):
     def test_new_cache_format(self, num_beams, do_sample):
         pass
 
-    @require_torch_gpu
+    @require_torch_accelerator
     def test_flex_attention_with_grads(self):
         """
         Overwriting as the base hidden size is big enough for compile.

filename: tests/test_modeling_common.py
patch: @@ -4207,7 +4207,7 @@ def update_config_headdim(config, requested_dim):
 
         return config
 
-    @require_torch_gpu
+    @require_torch_accelerator
     def test_flex_attention_with_grads(self):
         for model_class in self.all_model_classes:
             config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()