From 0870117d0efa3d16a84a930eb9288ef88f3a3fdf Mon Sep 17 00:00:00 2001
From: Michael Wyatt <michaelwyatt@microsoft.com>
Date: Tue, 6 Sep 2022 16:42:59 -0700
Subject: [PATCH] fix for pytest picking up local deepspeed dir instead of
 installed deepspeed

---
 tests/__init__.py                                         | 0
 tests/unit/autotuning/test_autotuning.py                  | 2 +-
 tests/unit/checkpoint/common.py                           | 2 +-
 tests/unit/checkpoint/test_latest_checkpoint.py           | 6 +++---
 tests/unit/checkpoint/test_lr_scheduler.py                | 6 +++---
 tests/unit/checkpoint/test_moe.py                         | 8 ++++----
 tests/unit/checkpoint/test_other_optimizer.py             | 6 +++---
 tests/unit/checkpoint/test_pipeline.py                    | 6 +++---
 tests/unit/checkpoint/test_sparse.py                      | 4 ++--
 tests/unit/checkpoint/test_tag_validation.py              | 4 ++--
 tests/unit/checkpoint/test_zero_optimizer.py              | 8 ++++----
 tests/unit/comm/test_dist.py                              | 4 ++--
 tests/unit/compression/test_compression.py                | 6 +++---
 tests/unit/elasticity/test_elastic.py                     | 4 ++--
 tests/unit/inference/test_inference.py                    | 2 +-
 tests/unit/monitor/test_monitor.py                        | 2 +-
 tests/unit/ops/adam/test_adamw.py                         | 4 ++--
 tests/unit/pipe/test_pipe_module.py                       | 2 +-
 .../unit/profiling/flops_profiler/test_flops_profiler.py  | 4 ++--
 tests/unit/runtime/fp16/onebit/test_onebit.py             | 6 +++---
 tests/unit/runtime/pipe/test_pipe.py                      | 4 ++--
 tests/unit/runtime/pipe/test_topology.py                  | 2 +-
 tests/unit/runtime/test_bf16.py                           | 6 +++---
 tests/unit/runtime/test_data.py                           | 4 ++--
 tests/unit/runtime/test_pld.py                            | 4 ++--
 tests/unit/runtime/test_runtime_utils.py                  | 2 +-
 tests/unit/test_sparse_grads.py                           | 2 +-
 tests/unit/utils/test_init_on_device.py                   | 2 +-
 28 files changed, 56 insertions(+), 56 deletions(-)
 delete mode 100644 tests/__init__.py

diff --git a/tests/__init__.py b/tests/__init__.py
deleted file mode 100644
index e69de29bb2d1..000000000000
diff --git a/tests/unit/autotuning/test_autotuning.py b/tests/unit/autotuning/test_autotuning.py
index 5ece541e1bb7..28a886c658ba 100644
--- a/tests/unit/autotuning/test_autotuning.py
+++ b/tests/unit/autotuning/test_autotuning.py
@@ -1,6 +1,6 @@
 import os
 import pytest
-from tests.unit.simple_model import create_config_from_dict
+from 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
diff --git a/tests/unit/checkpoint/common.py b/tests/unit/checkpoint/common.py
index d6422d4ba17d..d218ee7bb8d5 100644
--- a/tests/unit/checkpoint/common.py
+++ b/tests/unit/checkpoint/common.py
@@ -8,7 +8,7 @@
 from deepspeed.runtime.fp16.unfused_optimizer import FP16_UnfusedOptimizer
 from deepspeed.runtime.zero.stage3 import DeepSpeedZeroOptimizer_Stage3
 
-from tests.unit.simple_model import *
+from unit.simple_model import *
 
 
 def compare_deepspeed_states(saved_model, loaded_model):
diff --git a/tests/unit/checkpoint/test_latest_checkpoint.py b/tests/unit/checkpoint/test_latest_checkpoint.py
index f31fb269117f..1b62d232bb23 100644
--- a/tests/unit/checkpoint/test_latest_checkpoint.py
+++ b/tests/unit/checkpoint/test_latest_checkpoint.py
@@ -1,9 +1,9 @@
 import deepspeed
 
-from tests.unit.common import DistributedTest
-from tests.unit.simple_model import *
+from unit.common import DistributedTest
+from unit.simple_model import *
 
-from tests.unit.checkpoint.common import checkpoint_correctness_verification
+from unit.checkpoint.common import checkpoint_correctness_verification
 
 
 class TestLatestCheckpoint(DistributedTest):
diff --git a/tests/unit/checkpoint/test_lr_scheduler.py b/tests/unit/checkpoint/test_lr_scheduler.py
index 24bdc8c93df3..bd950b4183b6 100644
--- a/tests/unit/checkpoint/test_lr_scheduler.py
+++ b/tests/unit/checkpoint/test_lr_scheduler.py
@@ -1,10 +1,10 @@
 import deepspeed
 from deepspeed.ops.op_builder import CPUAdamBuilder
 
-from tests.unit.common import DistributedTest
-from tests.unit.simple_model import *
+from unit.common import DistributedTest
+from unit.simple_model import *
 
-from tests.unit.checkpoint.common import checkpoint_correctness_verification
+from unit.checkpoint.common import checkpoint_correctness_verification
 
 import pytest
 
diff --git a/tests/unit/checkpoint/test_moe.py b/tests/unit/checkpoint/test_moe.py
index 8f8b951e7093..41c97be747d6 100644
--- a/tests/unit/checkpoint/test_moe.py
+++ b/tests/unit/checkpoint/test_moe.py
@@ -1,10 +1,10 @@
 from deepspeed.moe.utils import split_params_into_different_moe_groups_for_optimizer
 
-from tests.unit.common import DistributedTest
-from tests.unit.simple_model import *
-from tests.unit.util import required_torch_version
+from unit.common import DistributedTest
+from unit.simple_model import *
+from unit.util import required_torch_version
 
-from tests.unit.checkpoint.common import checkpoint_correctness_verification
+from unit.checkpoint.common import checkpoint_correctness_verification
 
 import pytest
 
diff --git a/tests/unit/checkpoint/test_other_optimizer.py b/tests/unit/checkpoint/test_other_optimizer.py
index 87d40959b3b9..74a333399587 100644
--- a/tests/unit/checkpoint/test_other_optimizer.py
+++ b/tests/unit/checkpoint/test_other_optimizer.py
@@ -1,10 +1,10 @@
 import deepspeed
 from deepspeed.ops.op_builder import FusedLambBuilder
 
-from tests.unit.common import DistributedTest
-from tests.unit.simple_model import *
+from unit.common import DistributedTest
+from unit.simple_model import *
 
-from tests.unit.checkpoint.common import checkpoint_correctness_verification
+from unit.checkpoint.common import checkpoint_correctness_verification
 
 import pytest
 
diff --git a/tests/unit/checkpoint/test_pipeline.py b/tests/unit/checkpoint/test_pipeline.py
index 46b42aaa778b..a50fd8a425c4 100644
--- a/tests/unit/checkpoint/test_pipeline.py
+++ b/tests/unit/checkpoint/test_pipeline.py
@@ -1,8 +1,8 @@
 from deepspeed.runtime.checkpoint_engine.torch_checkpoint_engine import TorchCheckpointEngine
-from tests.unit.common import DistributedTest
-from tests.unit.simple_model import *
+from unit.common import DistributedTest
+from unit.simple_model import *
 
-from tests.unit.checkpoint.common import checkpoint_correctness_verification
+from unit.checkpoint.common import checkpoint_correctness_verification
 
 import pytest
 
diff --git a/tests/unit/checkpoint/test_sparse.py b/tests/unit/checkpoint/test_sparse.py
index e3b2d14fd9a4..c8db190f60c8 100644
--- a/tests/unit/checkpoint/test_sparse.py
+++ b/tests/unit/checkpoint/test_sparse.py
@@ -1,7 +1,7 @@
 import deepspeed
 
-from tests.unit.common import DistributedTest
-from tests.unit.simple_model import *
+from unit.common import DistributedTest
+from unit.simple_model import *
 
 import pytest
 
diff --git a/tests/unit/checkpoint/test_tag_validation.py b/tests/unit/checkpoint/test_tag_validation.py
index 14c8670bc304..89d954859164 100644
--- a/tests/unit/checkpoint/test_tag_validation.py
+++ b/tests/unit/checkpoint/test_tag_validation.py
@@ -1,7 +1,7 @@
 import deepspeed
 
-from tests.unit.common import DistributedTest
-from tests.unit.simple_model import *
+from unit.common import DistributedTest
+from unit.simple_model import *
 
 import pytest
 
diff --git a/tests/unit/checkpoint/test_zero_optimizer.py b/tests/unit/checkpoint/test_zero_optimizer.py
index 0b632a202666..e8c8c12a5495 100644
--- a/tests/unit/checkpoint/test_zero_optimizer.py
+++ b/tests/unit/checkpoint/test_zero_optimizer.py
@@ -1,11 +1,11 @@
 import deepspeed
 from deepspeed.ops.op_builder import CPUAdamBuilder
 
-from tests.unit.common import DistributedTest
-from tests.unit.simple_model import *
-from tests.unit.util import required_minimum_torch_version
+from unit.common import DistributedTest
+from unit.simple_model import *
+from unit.util import required_minimum_torch_version
 
-from tests.unit.checkpoint.common import *
+from unit.checkpoint.common import *
 
 import itertools
 import pytest
diff --git a/tests/unit/comm/test_dist.py b/tests/unit/comm/test_dist.py
index 0539446eb11e..d9697f31d7c0 100644
--- a/tests/unit/comm/test_dist.py
+++ b/tests/unit/comm/test_dist.py
@@ -2,8 +2,8 @@
 import deepspeed.comm as dist
 import deepspeed
 
-from tests.unit.common import DistributedTest, get_master_port
-from tests.unit.simple_model import SimpleModel
+from unit.common import DistributedTest, get_master_port
+from unit.simple_model import SimpleModel
 
 import pytest
 
diff --git a/tests/unit/compression/test_compression.py b/tests/unit/compression/test_compression.py
index 3f5a72a4d30c..6ff215c4ec13 100644
--- a/tests/unit/compression/test_compression.py
+++ b/tests/unit/compression/test_compression.py
@@ -2,10 +2,10 @@
 import pytest
 import random
 import numpy as np
-from tests.unit.megatron_model import get_gpt2_model
+from unit.megatron_model import get_gpt2_model
 from deepspeed.compression.compress import init_compression
-from tests.unit.modeling import BertConfig
-from tests.unit.modelingpreln import BertEncoder as BertEncoderPreln
+from unit.modeling import BertConfig
+from 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
 
diff --git a/tests/unit/elasticity/test_elastic.py b/tests/unit/elasticity/test_elastic.py
index 5702e6e4d7d8..aa4b3df76cc3 100644
--- a/tests/unit/elasticity/test_elastic.py
+++ b/tests/unit/elasticity/test_elastic.py
@@ -1,9 +1,9 @@
 import pytest
 import deepspeed
-from tests.unit.common import DistributedTest
+from unit.common import DistributedTest
 from deepspeed.git_version_info import version as ds_version
 import os
-from tests.unit.simple_model import SimpleModel
+from unit.simple_model import SimpleModel
 
 
 @pytest.fixture
diff --git a/tests/unit/inference/test_inference.py b/tests/unit/inference/test_inference.py
index 1294bb0d5895..cdc3e83232de 100644
--- a/tests/unit/inference/test_inference.py
+++ b/tests/unit/inference/test_inference.py
@@ -5,7 +5,7 @@
 import itertools
 import deepspeed
 from deepspeed.git_version_info import torch_info
-from tests.unit.common import DistributedTest
+from unit.common import DistributedTest
 from packaging import version as pkg_version
 from deepspeed.ops.op_builder import OpBuilder
 from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
diff --git a/tests/unit/monitor/test_monitor.py b/tests/unit/monitor/test_monitor.py
index 674a8d7ce841..a2fa4754b91a 100644
--- a/tests/unit/monitor/test_monitor.py
+++ b/tests/unit/monitor/test_monitor.py
@@ -4,7 +4,7 @@
 from deepspeed.monitor.wandb import WandbMonitor
 from deepspeed.monitor.csv_monitor import csvMonitor
 
-from tests.unit.common import DistributedTest
+from unit.common import DistributedTest
 from deepspeed.runtime.config import DeepSpeedConfig
 
 
diff --git a/tests/unit/ops/adam/test_adamw.py b/tests/unit/ops/adam/test_adamw.py
index c898d0c0ee1c..6ca985efe0e6 100644
--- a/tests/unit/ops/adam/test_adamw.py
+++ b/tests/unit/ops/adam/test_adamw.py
@@ -4,8 +4,8 @@
 
 from deepspeed.ops.adam import FusedAdam
 from deepspeed.ops.adam import DeepSpeedCPUAdam
-from tests.unit.common import DistributedTest
-from tests.unit.simple_model import SimpleModel
+from unit.common import DistributedTest
+from unit.simple_model import SimpleModel
 
 # yapf: disable
 #'optimizer, zero_offload, torch_adam, adam_w_mode, resulting_optimizer
diff --git a/tests/unit/pipe/test_pipe_module.py b/tests/unit/pipe/test_pipe_module.py
index 99862a96a2d3..2b8cce57d7a8 100644
--- a/tests/unit/pipe/test_pipe_module.py
+++ b/tests/unit/pipe/test_pipe_module.py
@@ -10,7 +10,7 @@
 from deepspeed.pipe import PipelineModule
 from deepspeed.utils import RepeatingLoader
 
-from tests.unit.common import DistributedTest
+from unit.common import DistributedTest
 
 HIDDEN_DIM = 32
 LAYERS = 8
diff --git a/tests/unit/profiling/flops_profiler/test_flops_profiler.py b/tests/unit/profiling/flops_profiler/test_flops_profiler.py
index 734e2996fa80..fe1e77edddc8 100644
--- a/tests/unit/profiling/flops_profiler/test_flops_profiler.py
+++ b/tests/unit/profiling/flops_profiler/test_flops_profiler.py
@@ -2,8 +2,8 @@
 import pytest
 import deepspeed
 from deepspeed.profiling.flops_profiler import get_model_profile
-from tests.unit.simple_model import SimpleModel, random_dataloader
-from tests.unit.common import DistributedTest
+from unit.simple_model import SimpleModel, random_dataloader
+from unit.common import DistributedTest
 
 TORCH_MAJOR = int(torch.__version__.split('.')[0])
 TORCH_MINOR = int(torch.__version__.split('.')[1])
diff --git a/tests/unit/runtime/fp16/onebit/test_onebit.py b/tests/unit/runtime/fp16/onebit/test_onebit.py
index 139ac79c468e..451d6abb6731 100644
--- a/tests/unit/runtime/fp16/onebit/test_onebit.py
+++ b/tests/unit/runtime/fp16/onebit/test_onebit.py
@@ -10,9 +10,9 @@
 from deepspeed.runtime.pipe.topology import PipeDataParallelTopology
 from deepspeed.ops.op_builder import OpBuilder
 from deepspeed.runtime.pipe.module import PipelineModule
-from tests.unit.common import DistributedTest
-from tests.unit.simple_model import SimpleModel, random_dataloader
-from tests.unit.alexnet_model import AlexNetPipe, train_cifar
+from unit.common import DistributedTest
+from unit.simple_model import SimpleModel, random_dataloader
+from unit.alexnet_model import AlexNetPipe, train_cifar
 
 PipeTopo = PipeDataParallelTopology
 
diff --git a/tests/unit/runtime/pipe/test_pipe.py b/tests/unit/runtime/pipe/test_pipe.py
index 7c097b1522a5..e5c4ab41172a 100644
--- a/tests/unit/runtime/pipe/test_pipe.py
+++ b/tests/unit/runtime/pipe/test_pipe.py
@@ -5,8 +5,8 @@
 import deepspeed.comm as dist
 from deepspeed.runtime.pipe.topology import PipeDataParallelTopology
 from deepspeed.runtime.pipe.module import PipelineModule
-from tests.unit.alexnet_model import AlexNetPipe, train_cifar
-from tests.unit.common import DistributedTest
+from unit.alexnet_model import AlexNetPipe, train_cifar
+from unit.common import DistributedTest
 
 PipeTopo = PipeDataParallelTopology
 
diff --git a/tests/unit/runtime/pipe/test_topology.py b/tests/unit/runtime/pipe/test_topology.py
index ae6fb7ecfeee..35860c5f5167 100644
--- a/tests/unit/runtime/pipe/test_topology.py
+++ b/tests/unit/runtime/pipe/test_topology.py
@@ -7,7 +7,7 @@
 from deepspeed.runtime.pipe.topology import ProcessTopology as Topo
 from deepspeed.runtime.pipe.topology import _prime_factors
 
-from tests.unit.common import DistributedTest
+from unit.common import DistributedTest
 
 
 def test_topology_2d():
diff --git a/tests/unit/runtime/test_bf16.py b/tests/unit/runtime/test_bf16.py
index d75b92441117..de15a0868df4 100644
--- a/tests/unit/runtime/test_bf16.py
+++ b/tests/unit/runtime/test_bf16.py
@@ -2,10 +2,10 @@
 import deepspeed
 import pytest
 from deepspeed.ops.adam import FusedAdam
-from tests.unit.common import DistributedTest
+from unit.common import DistributedTest
 from deepspeed.ops.op_builder import CPUAdamBuilder
-from tests.unit.simple_model import SimpleModel, SimpleOptimizer, random_dataloader
-from tests.unit.util import bf16_required_version_check
+from unit.simple_model import SimpleModel, SimpleOptimizer, random_dataloader
+from unit.util import bf16_required_version_check
 from deepspeed import comm as dist
 
 
diff --git a/tests/unit/runtime/test_data.py b/tests/unit/runtime/test_data.py
index c15d2f48eeef..e87f6c5e96f0 100644
--- a/tests/unit/runtime/test_data.py
+++ b/tests/unit/runtime/test_data.py
@@ -2,8 +2,8 @@
 import torch
 import pytest
 import deepspeed
-from tests.unit.common import DistributedTest
-from tests.unit.simple_model import SimpleModel, random_dataset
+from unit.common import DistributedTest
+from unit.simple_model import SimpleModel, random_dataset
 
 
 def test_repeating_loader():
diff --git a/tests/unit/runtime/test_pld.py b/tests/unit/runtime/test_pld.py
index 736fb3459438..c7766a8be7e1 100644
--- a/tests/unit/runtime/test_pld.py
+++ b/tests/unit/runtime/test_pld.py
@@ -3,8 +3,8 @@
 import pytest
 from deepspeed.runtime.progressive_layer_drop import ProgressiveLayerDrop
 
-from tests.unit.common import DistributedTest
-from tests.unit.simple_model import SimpleModel, PLD_SimpleModel, random_dataloader
+from unit.common import DistributedTest
+from unit.simple_model import SimpleModel, PLD_SimpleModel, random_dataloader
 
 
 @pytest.mark.parametrize('theta', [0, 0.1, 0.9, 1.0])
diff --git a/tests/unit/runtime/test_runtime_utils.py b/tests/unit/runtime/test_runtime_utils.py
index e59ede26c8df..33f40ad30a06 100644
--- a/tests/unit/runtime/test_runtime_utils.py
+++ b/tests/unit/runtime/test_runtime_utils.py
@@ -6,7 +6,7 @@
 import deepspeed.runtime.utils as ds_utils
 import deepspeed.utils.groups as groups
 
-from tests.unit.common import DistributedTest
+from unit.common import DistributedTest
 
 
 def test_call_to_str():
diff --git a/tests/unit/test_sparse_grads.py b/tests/unit/test_sparse_grads.py
index 94cc389db7b5..d63a47347988 100644
--- a/tests/unit/test_sparse_grads.py
+++ b/tests/unit/test_sparse_grads.py
@@ -1,6 +1,6 @@
 import torch
 import deepspeed
-from tests.unit.common import DistributedTest
+from unit.common import DistributedTest
 
 import deepspeed.utils.groups as groups
 
diff --git a/tests/unit/utils/test_init_on_device.py b/tests/unit/utils/test_init_on_device.py
index 8bbd70c7f7fc..46d179b439cf 100644
--- a/tests/unit/utils/test_init_on_device.py
+++ b/tests/unit/utils/test_init_on_device.py
@@ -1,6 +1,6 @@
 import torch
 import pytest
-from tests.unit.simple_model import SimpleModel
+from unit.simple_model import SimpleModel
 from deepspeed import OnDevice
 from packaging import version as pkg_version