[CI] follow-up fixes (microsoft#3072)

CODEOWNERS

@@ -7,7 +7,7 @@
 
 
 # top-level repo folders
-/.github/ @jeffra @mrwyattii
+/.github/ @jeffra @mrwyattii @loadams
 /azure/ @jeffra @awan-10
 /benchmarks/ @jeffra @awan-10 @mrwyattii @molly-smith
 /bin/ @jeffra

tests/unit/checkpoint/test_pipeline.py

@@ -3,29 +3,18 @@
 from deepspeed.runtime.checkpoint_engine.torch_checkpoint_engine import TorchCheckpointEngine
 from unit.common import DistributedTest
 from unit.simple_model import *
-
 from unit.checkpoint.common import checkpoint_correctness_verification
+from unit.util import skip_on_arch
 
 import pytest
-import deepspeed
-import torch
-
-
-def _skip_on_older_arch(arch=7):
-    if deepspeed.accelerator.get_accelerator().device_name() == 'cuda':
-        if torch.cuda.get_device_capability()[0] < arch:
-            pytest.skip("needs higher compute capability than 7")
-    else:
-        assert deepspeed.accelerator.get_accelerator().device_name() == 'xpu'
-        return
 
 
 class TestPipelineCheckpoint(DistributedTest):
     world_size = 4
 
     @pytest.mark.parametrize("zero_stage", [0, 1])
     def test_checkpoint_pipe_engine(self, zero_stage, tmpdir):
-        _skip_on_older_arch()
+        skip_on_arch(min_arch=7)
 
         config_dict = {
             "train_batch_size": 2,

tests/unit/runtime/pipe/test_pipe.py

@@ -4,13 +4,12 @@
 import torch.nn as nn
 import pytest
 
-import torch
-import deepspeed
 import deepspeed.comm as dist
 from deepspeed.runtime.pipe.topology import PipeDataParallelTopology
 from deepspeed.runtime.pipe.module import PipelineModule
 from unit.alexnet_model import AlexNetPipe, train_cifar
 from unit.common import DistributedTest
+from unit.util import skip_on_arch
 
 PipeTopo = PipeDataParallelTopology
 
@@ -19,15 +18,6 @@ def rel_diff(A, B):
     return abs(A - B) / abs(A)
 
 
-def _skip_on_older_arch(arch=7):
-    if deepspeed.accelerator.get_accelerator().device_name() == 'cuda':
-        if torch.cuda.get_device_capability()[0] < arch:
-            pytest.skip("needs higher compute capability than 7")
-    else:
-        assert deepspeed.accelerator.get_accelerator().device_name() == 'xpu'
-        return
-
-
 @pytest.mark.parametrize('topo_config',
                          [
                              {
@@ -47,7 +37,7 @@ class TestPipeCifar10(DistributedTest):
     world_size = 4
 
     def test(self, topo_config):
-        _skip_on_older_arch()
+        skip_on_arch(min_arch=7)
 
         config_dict = {
             "train_batch_size": 16,

tests/unit/runtime/sparse_tensor/test_averaging_sparse_gradients.py

@@ -3,6 +3,7 @@
 import torch
 import deepspeed
 from unit.common import DistributedTest
+from unit.util import skip_on_arch
 
 
 class Model(torch.nn.Module):
@@ -49,6 +50,8 @@ class TestSparseAdam(DistributedTest):
     world_size = 2
 
     def test(self):
+        skip_on_arch(min_arch=7)
+
         config_dict = {
             "train_batch_size": 2,
             "steps_per_print": 1,

tests/unit/util.py

@@ -1,9 +1,20 @@
 '''Copyright The Microsoft DeepSpeed Team'''
 
+import pytest
 import torch
+import deepspeed
 from deepspeed.git_version_info import torch_info
 
 
+def skip_on_arch(min_arch=7):
+    if deepspeed.accelerator.get_accelerator().device_name() == 'cuda':
+        if torch.cuda.get_device_capability()[0] < min_arch:
+            pytest.skip(f"needs higher compute capability than {min_arch}")
+    else:
+        assert deepspeed.accelerator.get_accelerator().device_name() == 'xpu'
+        return
+
+
 def required_torch_version():
     TORCH_MAJOR = int(torch.__version__.split('.')[0])
     TORCH_MINOR = int(torch.__version__.split('.')[1])