Refactor Distributed Tests

.github/workflows/amd.yml

@@ -66,5 +66,5 @@ jobs:
         run: |
           if [[ -d ./torch-extensions ]]; then rm -rf ./torch-extensions; fi
           cd tests
-          TORCH_EXTENSIONS_DIR=./torch-extensions pytest --color=yes --durations=0 --verbose -n 4 unit/{comm,inference,monitor,ops,profiling,runtime}
-          #TORCH_EXTENSIONS_DIR=./torch-extensions pytest --color=yes --durations=0 --verbose -m 'sequential' unit/{comm,inference,monitor,ops,profiling,runtime}
+          TORCH_EXTENSIONS_DIR=./torch-extensions pytest --color=yes --durations=0 --verbose unit/{autotuning,checkpoint,comm,compression,elasticity,inference,launcher,monitor,ops,profiling,runtime,utils}
+          #TORCH_EXTENSIONS_DIR=./torch-extensions pytest --color=yes --durations=0 --verbose -m 'sequential' unit/{autotuning,checkpoint,comm,compression,elasticity,inference,launcher,monitor,ops,profiling,runtime,utils}

.github/workflows/nv-torch-latest-v100.yml

@@ -60,6 +60,5 @@ jobs:
           unset TORCH_CUDA_ARCH_LIST # only jit compile for current arch
           if [[ -d ./torch-extensions ]]; then rm -rf ./torch-extensions; fi
           cd tests
-          EXPECTED_TORCH=$(pip index versions torch | grep -oP -m1 "^\s*LATEST.*\s\K\d+\.\d+")
-          TORCH_EXTENSIONS_DIR=./torch-extensions pytest --color=yes --durations=0 --forked --verbose -n 4 unit/ --torch_ver=$EXPECTED_TROCH --cuda_ver="11.3"
-          TORCH_EXTENSIONS_DIR=./torch-extensions pytest --color=yes --durations=0 --forked --verbose -m 'sequential' unit/ --torch_ver=$EXPECTED_TORCH --cuda_ver="11.3"
+          TORCH_EXTENSIONS_DIR=./torch-extensions pytest --color=yes --durations=0 --forked --verbose -n 4 unit/{autotuning,checkpoint,comm,compression,elasticity,inference,launcher,monitor,ops,profiling,runtime,utils} --torch_ver="1.12" --cuda_ver="11.3"
+          TORCH_EXTENSIONS_DIR=./torch-extensions pytest --color=yes --durations=0 --forked --verbose -m 'sequential' unit/{autotuning,checkpoint,comm,compression,elasticity,inference,launcher,monitor,ops,profiling,runtime,utils} --torch_ver="1.12" --cuda_ver="11.3"

tests/unit/test_pipe.py → tests/unit/alexnet_model.py

@@ -1,30 +1,12 @@
-import copy
-
+import pytest
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-import deepspeed.comm as dist
-
-import pytest
-
 import deepspeed
+import deepspeed.comm as dist
 import deepspeed.runtime.utils as ds_utils
-
-from deepspeed.runtime.pipe.topology import PipeDataParallelTopology
-
-PipeTopo = PipeDataParallelTopology
 from deepspeed.runtime.pipe.module import PipelineModule, LayerSpec
 
-from .common import distributed_test
-
-
-def rel_diff(A, B):
-    return abs(A - B) / abs(A)
-
-
-# All models
-from .simple_model import args_from_dict
-
 
 class AlexNet(nn.Module):
     def __init__(self, num_classes=10):
@@ -103,8 +85,13 @@ def __init__(self, num_classes=10, **kwargs):
         super().__init__(layers=specs, loss_fn=nn.CrossEntropyLoss(), **kwargs)
 
 
+# Define this here because we cannot pickle local lambda functions
+def cast_to_half(x):
+    return x.half()
+
+
 def cifar_trainset(fp16=False):
-    import torchvision
+    torchvision = pytest.importorskip("torchvision", minversion="0.5.0")
     import torchvision.transforms as transforms
 
     transform_list = [
@@ -117,7 +104,7 @@ def cifar_trainset(fp16=False):
                               0.5)),
     ]
     if fp16:
-        transform_list.append(torchvision.transforms.Lambda(lambda x: x.half()))
+        transform_list.append(torchvision.transforms.Lambda(cast_to_half))
 
     transform = transforms.Compose(transform_list)
 
@@ -136,18 +123,23 @@ def cifar_trainset(fp16=False):
     return trainset
 
 
-def train_cifar(model, args, num_steps=400, average_dp_losses=True, fp16=True, seed=123):
+def train_cifar(model,
+                config,
+                num_steps=400,
+                average_dp_losses=True,
+                fp16=True,
+                seed=123):
     with torch.random.fork_rng(devices=[torch.cuda.current_device()]):
         ds_utils.set_random_seed(seed)
 
         # disable dropout
         model.eval()
 
         trainset = cifar_trainset(fp16=fp16)
-        args.local_rank = dist.get_rank()
+        config['local_rank'] = dist.get_rank()
 
         engine, _, _, _ = deepspeed.initialize(
-            args=args,
+            config=config,
             model=model,
             model_parameters=[p for p in model.parameters()],
             training_data=trainset)
@@ -166,102 +158,3 @@ def train_cifar(model, args, num_steps=400, average_dp_losses=True, fp16=True, s
             losses = loss_tensor.tolist()
 
     return losses
-
-
-@pytest.mark.skip(reason="been seeing nondeterministic failures, skipping for now")
-@pytest.mark.parametrize('topo',
-                         [
-                             PipeTopo(num_pp=1,
-                                      num_dp=4),
-                             PipeTopo(num_pp=2,
-                                      num_dp=2),
-                             PipeTopo(num_pp=4,
-                                      num_dp=1),
-                         ])
-def test_pipe_cifar10(topo, tmpdir):
-    config_dict = {
-        "train_batch_size": 16,
-        "train_micro_batch_size_per_gpu": 4,
-        "steps_per_print": 20,
-        "optimizer": {
-            "type": "Adam",
-            "params": {
-                "lr": 0.001,
-                "betas": [0.9,
-                          0.999],
-                "eps": 1e-8,
-                "weight_decay": 3e-7
-            }
-        },
-        "zero_optimization": {
-            "stage": 0
-        },
-        "fp16": {
-            "enabled": False
-        },
-        "pipeline": {
-            "seed_layers": True,
-            "activation_checkpoint_interval": 1
-        }
-    }
-    args = args_from_dict(tmpdir, config_dict)
-
-    # Allocate model for consistent initial weights.
-    init_net = AlexNetPipe()
-
-    @distributed_test(world_size=4)
-    def _helper(topo, tmpdir, steps=500):
-        assert steps >= 100
-
-        base_net = copy.deepcopy(init_net)
-        base_model = PipelineModule(layers=base_net.to_layers(),
-                                    num_stages=1,
-                                    loss_fn=nn.CrossEntropyLoss())
-
-        # Train with just data parallelism
-        base_losses = train_cifar(base_model,
-                                  args,
-                                  num_steps=steps,
-                                  fp16=config_dict['fp16']['enabled'])
-
-        test_net = copy.deepcopy(init_net)
-        test_model = PipelineModule(layers=test_net.to_layers(),
-                                    topology=topo,
-                                    loss_fn=nn.CrossEntropyLoss())
-
-        #test_model = AlexNetPipe(num_classes=10,
-        #                         topology=test_topo,
-        #                         seed_layers=config_dict['pipeline']['seed_layers'])
-        test_losses = train_cifar(test_model,
-                                  args,
-                                  num_steps=steps,
-                                  fp16=config_dict['fp16']['enabled'])
-
-        abs_diffs = [l0 - l1 for l0, l1 in zip(base_losses, test_losses)]
-        rel_diffs = [rel_diff(l0, l1) for l0, l1 in zip(base_losses, test_losses)]
-        if dist.get_rank() == 0:
-            print(
-                f'abs min={min(abs_diffs)} max={max(abs_diffs)} avg={sum(abs_diffs)/len(abs_diffs)}'
-            )
-            print(
-                f'rel min={min(rel_diffs)} max={max(rel_diffs)} avg={sum(rel_diffs)/len(rel_diffs)}'
-            )
-            print(
-                f'first: base={base_losses[0]} test={test_losses[0]} abs={abs_diffs[0]} rel={rel_diffs[0]}'
-            )
-
-            for lastX in [1, 10, 100]:
-                base_avg = sum(base_losses[-lastX:]) / lastX
-                test_avg = sum(test_losses[-lastX:]) / lastX
-                print(
-                    f'last-{lastX}: base={base_avg} test={test_avg} abs={base_avg - test_avg} rel={rel_diff(base_avg, test_avg)}'
-                )
-
-        lastX = 100
-        base = base_losses[-lastX:]
-        base_avg = sum(base) / len(base)
-        test = test_losses[-lastX:]
-        test_avg = sum(test) / len(test)
-        assert rel_diff(base_avg, test_avg) < 0.03
-
-    _helper(topo, tmpdir)

tests/unit/test_autotuning.py → tests/unit/autotuning/test_autotuning.py

@@ -1,6 +1,6 @@
 import os
 import pytest
-from .simple_model import create_config_from_dict
+from tests.unit.simple_model import create_config_from_dict
 from deepspeed.launcher import runner as dsrun
 from deepspeed.autotuning.autotuner import Autotuner
 from deepspeed.autotuning.scheduler import ResourceManager

tests/unit/test_compression.py → tests/unit/compression/test_compression.py

@@ -2,10 +2,10 @@
 import pytest
 import random
 import numpy as np
-from .megatron_model import get_gpt2_model
+from tests.unit.megatron_model import get_gpt2_model
 from deepspeed.compression.compress import init_compression
-from .modeling import BertConfig
-from .modelingpreln import BertEncoder as BertEncoderPreln
+from tests.unit.modeling import BertConfig
+from tests.unit.modelingpreln import BertEncoder as BertEncoderPreln
 from deepspeed.compression.basic_layer import LinearLayer_Compress, ColumnParallelLinear_Compress, RowParallelLinear_Compress
 from deepspeed.compression.helper import convert_conv1d_to_linear