Skip 1Bit Compression and sparsegrad tests for HPU. (microsoft#5270)

HPU accelrator does not support 1bit and sparsegrad tests.
dbyoung18 · Jun 11, 2024 · 816eec2 · 816eec2
1 parent 6c8d334
commit 816eec2
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 1 deletion.
diff --git a/tests/unit/runtime/half_precision/onebit/test_onebit.py b/tests/unit/runtime/half_precision/onebit/test_onebit.py
@@ -33,6 +33,9 @@
     pytest.skip("NCCL-based 1-bit compression is not yet supported w. ROCm 5 until cupy supports ROCm 5",
                 allow_module_level=True)
 
+if get_accelerator().device_name() == 'hpu':
+    pytest.skip("1-bit compression is not supported by HPU.", allow_module_level=True)
+
 
 @pytest.mark.parametrize("dtype", [torch.float32, torch.float16], ids=["fp32", "fp16"])
 class TestOneBitAdamBasic(DistributedTest):

diff --git a/tests/unit/runtime/sparse_tensor/test_averaging_sparse_gradients.py b/tests/unit/runtime/sparse_tensor/test_averaging_sparse_gradients.py
@@ -4,9 +4,14 @@
 # DeepSpeed Team
 
 import torch
+import pytest
 import deepspeed
 from unit.common import DistributedTest
 from unit.util import skip_on_arch
+from deepspeed.accelerator import get_accelerator
+
+if get_accelerator().device_name() == 'hpu':
+    pytest.skip("sparse_gradients not supported by HPU.", allow_module_level=True)
 
 
 class Model(torch.nn.Module):

diff --git a/tests/unit/runtime/sparse_tensor/test_sparse_grads.py b/tests/unit/runtime/sparse_tensor/test_sparse_grads.py
@@ -4,11 +4,15 @@
 # DeepSpeed Team
 
 import torch
+import pytest
 import deepspeed
 from unit.common import DistributedTest
-
+from deepspeed.accelerator import get_accelerator
 import deepspeed.utils.groups as groups
 
+if get_accelerator().device_name() == 'hpu':
+    pytest.skip("sparse_gradients not supported by HPU.", allow_module_level=True)
+
 
 class Model(torch.nn.Module):