CI Enable 5 test cases on XPU (#2442)

Signed-off-by: Yao, Matrix <[email protected]>

Author: yao-matrix

scripts/launch_notebook_mp.py

@@ -22,6 +22,7 @@
 from accelerate import notebook_launcher
 
 import peft
+from peft.utils import infer_device
 
 
 def init():
@@ -33,7 +34,8 @@ def __init__(self):
         def forward(self, x):
             return self.linear(x)
 
-    model = MyModule().to("cuda")
+    device = infer_device()
+    model = MyModule().to(device)
     peft.get_peft_model(model, peft.LoraConfig(target_modules=["linear"]))
 
 

tests/test_common_gpu.py

@@ -62,7 +62,6 @@
     require_bitsandbytes,
     require_multi_accelerator,
     require_non_cpu,
-    require_torch_gpu,
 )
 
 
@@ -1252,7 +1251,7 @@ def test_4bit_dora_merging(self):
         assert torch.allclose(out_dora, out_unmerged, atol=atol, rtol=rtol)
         assert torch.allclose(out_dora, out_unloaded, atol=atol, rtol=rtol)
 
-    @require_torch_gpu
+    @require_non_cpu
     @pytest.mark.single_gpu_tests
     @require_bitsandbytes
     def test_8bit_dora_merging(self):

tests/test_gpu_examples.py

@@ -79,6 +79,7 @@
 from peft.utils.other import fsdp_auto_wrap_policy
 
 from .testing_utils import (
+    device_count,
     require_aqlm,
     require_auto_awq,
     require_auto_gptq,
@@ -302,7 +303,7 @@ def test_causal_lm_training_multi_gpu_4bit(self):
                 quantization_config=BitsAndBytesConfig(load_in_4bit=True),
             )
 
-            assert set(model.hf_device_map.values()) == set(range(torch.cuda.device_count()))
+            assert set(model.hf_device_map.values()) == set(range(device_count))
 
             model = prepare_model_for_kbit_training(model)
 
@@ -424,7 +425,7 @@ def on_optimizer_step(self, args, state, control, **kwargs):
             assert trainer.state.log_history[-1]["train_loss"] is not None
 
     @pytest.mark.single_gpu_tests
-    @require_torch_gpu
+    @require_non_cpu
     def test_8bit_adalora_causalLM(self):
         r"""
         Tests the 8bit training with adalora
@@ -497,7 +498,7 @@ def on_optimizer_step(self, args, state, control, **kwargs):
             assert trainer.state.log_history[-1]["train_loss"] is not None
 
     @pytest.mark.multi_gpu_tests
-    @require_torch_multi_gpu
+    @require_multi_accelerator
     def test_causal_lm_training_multi_gpu(self):
         r"""
         Test the CausalLM training on a multi-GPU device. This test is a converted version of
@@ -511,8 +512,8 @@ def test_causal_lm_training_multi_gpu(self):
                 quantization_config=BitsAndBytesConfig(load_in_8bit=True),
                 device_map="auto",
             )
-
-            assert set(model.hf_device_map.values()) == set(range(torch.cuda.device_count()))
+            print(f"device map: {model.hf_device_map}")
+            assert set(model.hf_device_map.values()) == set(range(device_count))
 
             tokenizer = AutoTokenizer.from_pretrained(self.causal_lm_model_id)
             model = prepare_model_for_kbit_training(model)
@@ -621,7 +622,7 @@ def test_seq2seq_lm_training_single_gpu(self):
             assert trainer.state.log_history[-1]["train_loss"] is not None
 
     @pytest.mark.multi_gpu_tests
-    @require_torch_multi_gpu
+    @require_multi_accelerator
     def test_seq2seq_lm_training_multi_gpu(self):
         r"""
         Test the Seq2SeqLM training on a multi-GPU device. This test is a converted version of
@@ -636,7 +637,7 @@ def test_seq2seq_lm_training_multi_gpu(self):
                 device_map="balanced",
             )
 
-            assert set(model.hf_device_map.values()) == set(range(torch.cuda.device_count()))
+            assert set(model.hf_device_map.values()) == set(range(device_count))
 
             tokenizer = AutoTokenizer.from_pretrained(self.seq2seq_model_id)
             model = prepare_model_for_kbit_training(model)
@@ -920,7 +921,7 @@ def test_causal_lm_training_multi_gpu_4bit_dora(self):
                 quantization_config=BitsAndBytesConfig(load_in_4bit=True),
             )
 
-            assert set(model.hf_device_map.values()) == set(range(torch.cuda.device_count()))
+            assert set(model.hf_device_map.values()) == set(range(device_count))
 
             model = prepare_model_for_kbit_training(model)
 
@@ -1037,7 +1038,7 @@ def test_causal_lm_training_multi_gpu_8bit_dora(self):
                 quantization_config=BitsAndBytesConfig(load_in_8bit=True),
             )
 
-            assert set(model.hf_device_map.values()) == set(range(torch.cuda.device_count()))
+            assert set(model.hf_device_map.values()) == set(range(device_count))
 
             model = prepare_model_for_kbit_training(model)
 
@@ -1284,7 +1285,7 @@ def test_causal_lm_training_multi_gpu_vera(self):
                 quantization_config=BitsAndBytesConfig(load_in_8bit=True),
             )
 
-            assert set(model.hf_device_map.values()) == set(range(torch.cuda.device_count()))
+            assert set(model.hf_device_map.values()) == set(range(device_count))
 
             model = prepare_model_for_kbit_training(model)
 
@@ -1343,7 +1344,7 @@ def test_causal_lm_training_multi_gpu_4bit_vera(self):
                 quantization_config=BitsAndBytesConfig(load_in_4bit=True),
             )
 
-            assert set(model.hf_device_map.values()) == set(range(torch.cuda.device_count()))
+            assert set(model.hf_device_map.values()) == set(range(device_count))
 
             model = prepare_model_for_kbit_training(model)
 
@@ -1656,7 +1657,7 @@ def test_causal_lm_training_multi_gpu(self):
                 quantization_config=self.quantization_config,
             )
 
-            assert set(model.hf_device_map.values()) == set(range(torch.cuda.device_count()))
+            assert set(model.hf_device_map.values()) == set(range(device_count))
 
             model = prepare_model_for_kbit_training(model)
 
@@ -2552,7 +2553,7 @@ def test_config_no_loftq_config(self):
 
 
 @require_bitsandbytes
-@require_torch_gpu
+@require_non_cpu
 class MultiprocessTester(unittest.TestCase):
     def test_notebook_launcher(self):
         script_path = os.path.join("scripts", "launch_notebook_mp.py")
@@ -3187,7 +3188,7 @@ def test_causal_lm_training_multi_gpu(self):
                 device_map="auto",
             )
 
-            assert set(model.hf_device_map.values()) == set(range(torch.cuda.device_count()))
+            assert set(model.hf_device_map.values()) == set(range(device_count))
 
             model = prepare_model_for_kbit_training(model)
 
@@ -3335,7 +3336,7 @@ def test_causal_lm_training_multi_gpu_eetq(self):
                 quantization_config=quantization_config,
             )
 
-            assert set(model.hf_device_map.values()) == set(range(torch.cuda.device_count()))
+            assert set(model.hf_device_map.values()) == set(range(device_count))
 
             model = prepare_model_for_kbit_training(model)
 
@@ -3586,7 +3587,7 @@ def test_causal_lm_training_multi_gpu_torchao(self, quant_type):
                 torch_dtype=torch.bfloat16,
             )
 
-            assert set(model.hf_device_map.values()) == set(range(torch.cuda.device_count()))
+            assert set(model.hf_device_map.values()) == set(range(device_count))
 
             model = prepare_model_for_kbit_training(model)
             model.model_parallel = True
@@ -3646,7 +3647,7 @@ def test_causal_lm_training_multi_gpu_torchao_int4_raises(self):
             torch_dtype=torch.bfloat16,
         )
 
-        assert set(model.hf_device_map.values()) == set(range(torch.cuda.device_count()))
+        assert set(model.hf_device_map.values()) == set(range(device_count))
 
         model = prepare_model_for_kbit_training(model)
         model.model_parallel = True