Fix for pytest picking up wrong deepspeed

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

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):

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):

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
 

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
 

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
 

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
 

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
 

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
 

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

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
 

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
 

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

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

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
 
 

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

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

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])

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
 

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
 

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():

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
 
 

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():

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])

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():

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
 

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