device agnostic deepspeed&fsdp testing (#2235)

* device agnostic deepspeed testing * device agnostic fsdp testing * fix failing deepspeed test * make style --------- Co-authored-by: Zach Mueller <[email protected]>
huggingface · Dec 20, 2023 · b565a6c · b565a6c
1 parent a03c361
commit b565a6c
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 8 deletions.
diff --git a/tests/deepspeed/test_deepspeed.py b/tests/deepspeed/test_deepspeed.py
@@ -36,9 +36,9 @@
     AccelerateTestCase,
     TempDirTestCase,
     execute_subprocess_async,
-    require_cuda,
     require_deepspeed,
-    require_multi_gpu,
+    require_multi_device,
+    require_non_cpu,
     slow,
 )
 from accelerate.test_utils.training import RegressionDataset
@@ -90,7 +90,7 @@ def parameterized_custom_name_func(func, param_num, param):
 
 
 @require_deepspeed
-@require_cuda
+@require_non_cpu
 class DeepSpeedConfigIntegration(AccelerateTestCase):
     def setUp(self):
         super().setUp()
@@ -793,7 +793,7 @@ def test_basic_run(self):
 
 
 @require_deepspeed
-@require_multi_gpu
+@require_multi_device
 @slow
 class DeepSpeedIntegrationTest(TempDirTestCase):
     def setUp(self):

diff --git a/tests/fsdp/test_fsdp.py b/tests/fsdp/test_fsdp.py
@@ -28,9 +28,9 @@
     AccelerateTestCase,
     TempDirTestCase,
     execute_subprocess_async,
-    require_cuda,
     require_fsdp,
-    require_multi_gpu,
+    require_multi_device,
+    require_non_cpu,
     slow,
 )
 from accelerate.utils.constants import (
@@ -52,7 +52,7 @@
 
 
 @require_fsdp
-@require_cuda
+@require_non_cpu
 class FSDPPluginIntegration(AccelerateTestCase):
     def setUp(self):
         super().setUp()
@@ -170,7 +170,7 @@ def test_cpu_offload(self):
 
 
 @require_fsdp
-@require_multi_gpu
+@require_multi_device
 @slow
 class FSDPIntegrationTest(TempDirTestCase):
     def setUp(self):