From 3e4a9c1bd4b66ec95815f705536f28cdc3e61a24 Mon Sep 17 00:00:00 2001
From: OVH <kmfoda@gmail.com>
Date: Wed, 15 Jun 2022 16:05:51 +0000
Subject: [PATCH] fix import errors

---
 deepspeed/autotuning/scheduler.py                           | 2 +-
 deepspeed/comm/__init__.py                                  | 2 +-
 deepspeed/comm/comm.py                                      | 4 ++--
 deepspeed/comm/torch.py                                     | 2 +-
 deepspeed/inference/engine.py                               | 6 +++---
 deepspeed/moe/layer.py                                      | 4 ++--
 deepspeed/moe/sharded_moe.py                                | 4 ++--
 deepspeed/moe/utils.py                                      | 2 +-
 deepspeed/ops/transformer/inference/moe_inference.py        | 2 +-
 .../ops/transformer/inference/transformer_inference.py      | 2 +-
 deepspeed/runtime/activation_checkpointing/checkpointing.py | 2 +-
 deepspeed/runtime/bf16_optimizer.py                         | 2 +-
 deepspeed/runtime/comm/coalesced_collectives.py             | 2 +-
 deepspeed/runtime/comm/nccl.py                              | 2 +-
 deepspeed/runtime/config.py                                 | 2 +-
 deepspeed/runtime/engine.py                                 | 6 +++---
 deepspeed/runtime/fp16/fused_optimizer.py                   | 2 +-
 deepspeed/runtime/fp16/onebit/adam.py                       | 2 +-
 deepspeed/runtime/fp16/onebit/lamb.py                       | 2 +-
 deepspeed/runtime/fp16/onebit/zoadam.py                     | 2 +-
 deepspeed/runtime/fp16/unfused_optimizer.py                 | 2 +-
 deepspeed/runtime/pipe/engine.py                            | 2 +-
 deepspeed/runtime/pipe/module.py                            | 2 +-
 deepspeed/runtime/pipe/p2p.py                               | 2 +-
 deepspeed/runtime/pipe/topology.py                          | 2 +-
 deepspeed/runtime/swap_tensor/async_swapper.py              | 2 +-
 deepspeed/runtime/swap_tensor/optimizer_utils.py            | 2 +-
 .../runtime/swap_tensor/partitioned_optimizer_swapper.py    | 2 +-
 deepspeed/runtime/swap_tensor/partitioned_param_swapper.py  | 2 +-
 .../runtime/swap_tensor/pipelined_optimizer_swapper.py      | 2 +-
 deepspeed/runtime/swap_tensor/utils.py                      | 2 +-
 deepspeed/runtime/utils.py                                  | 2 +-
 deepspeed/runtime/zero/contiguous_memory_allocator.py       | 2 +-
 deepspeed/runtime/zero/linear.py                            | 2 +-
 deepspeed/runtime/zero/partition_parameters.py              | 2 +-
 deepspeed/runtime/zero/partitioned_param_coordinator.py     | 2 +-
 deepspeed/runtime/zero/stage3.py                            | 2 +-
 deepspeed/runtime/zero/stage_1_and_2.py                     | 2 +-
 deepspeed/runtime/zero/utils.py                             | 2 +-
 deepspeed/utils/groups.py                                   | 2 +-
 deepspeed/utils/logging.py                                  | 2 +-
 deepspeed/utils/timer.py                                    | 2 +-
 42 files changed, 49 insertions(+), 49 deletions(-)

diff --git a/deepspeed/autotuning/scheduler.py b/deepspeed/autotuning/scheduler.py
index 526223c48cb2..e3c4fbe7708b 100755
--- a/deepspeed/autotuning/scheduler.py
+++ b/deepspeed/autotuning/scheduler.py
@@ -25,7 +25,7 @@
 thread-N: start each experiment in its own thread
 """
 
-import deepspeed.comm as dist
+from deepspeed import comm as dist
 
 from datetime import datetime
 
diff --git a/deepspeed/comm/__init__.py b/deepspeed/comm/__init__.py
index f6146f164715..1975aab8582b 100644
--- a/deepspeed/comm/__init__.py
+++ b/deepspeed/comm/__init__.py
@@ -1,6 +1,6 @@
 import torch
 from .utils import *
-import deepspeed.utils as utils
+from deepspeed import utils
 
 supported_torch_version = False
 
diff --git a/deepspeed/comm/comm.py b/deepspeed/comm/comm.py
index e79f63596a82..ca328ee28dd9 100644
--- a/deepspeed/comm/comm.py
+++ b/deepspeed/comm/comm.py
@@ -14,7 +14,7 @@
     deepspeed.comm API
         -- must be kept fully compatible (same signatures) as torch.dist API to ensure backward/cross-framework compatibility.
         -- e.g. if a client code used
-            import deepspeed.comm as dist
+            from deepspeed import comm as dist
 
             instead of
             import torch.distributed as dist
@@ -49,7 +49,7 @@ class ReduceOp(Enum):
 from deepspeed.comm.backend import Backend
 from deepspeed.comm.torch import TorchBackend
 
-import deepspeed.utils as utils
+from deepspeed import utils
 from datetime import timedelta
 
 # Current deepspeed.comm backend (cdb) global object for simple access by client code
diff --git a/deepspeed/comm/torch.py b/deepspeed/comm/torch.py
index 95ef511a2c0d..dcfae9161213 100644
--- a/deepspeed/comm/torch.py
+++ b/deepspeed/comm/torch.py
@@ -4,7 +4,7 @@
 import os
 import torch
 
-import deepspeed.utils as utils
+from deepspeed import utils
 from ..constants import TORCH_DISTRIBUTED_DEFAULT_PORT, default_pg_timeout
 from datetime import timedelta
 
diff --git a/deepspeed/inference/engine.py b/deepspeed/inference/engine.py
index 17999a9b3a38..eb98948130c6 100755
--- a/deepspeed/inference/engine.py
+++ b/deepspeed/inference/engine.py
@@ -4,7 +4,7 @@
 import torch
 import os
 from torch.nn.modules import Module
-import deepspeed.comm as dist
+from deepspeed import comm as dist
 from ..runtime.state_dict_factory import SDLoaderFactory
 from ..runtime.weight_quantizer import WeightQuantization
 from ..module_inject.replace_module import replace_transformer_layer
@@ -15,8 +15,8 @@
 from ..moe.utils import has_moe_layers
 from ..moe.layer import MoE
 
-import deepspeed.comm as dist
-import deepspeed.utils.groups as groups
+from deepspeed import comm as dist
+from deepspeed.utils import groups
 
 DS_INFERENCE_ENABLED = False
 
diff --git a/deepspeed/moe/layer.py b/deepspeed/moe/layer.py
index 0dc1cbce7849..7dcf4144c0e6 100644
--- a/deepspeed/moe/layer.py
+++ b/deepspeed/moe/layer.py
@@ -4,11 +4,11 @@
 
 import torch.nn.init as init
 import torch
-import deepspeed.comm as dist
+from deepspeed import comm as dist
 
 from deepspeed.utils import logger, log_dist
 
-import deepspeed.utils.groups as groups
+from deepspeed.utils import groups
 from .sharded_moe import MOELayer, TopKGate
 from .experts import Experts
 import copy
diff --git a/deepspeed/moe/sharded_moe.py b/deepspeed/moe/sharded_moe.py
index e6cd7e55ac25..d87d22dc7ff7 100644
--- a/deepspeed/moe/sharded_moe.py
+++ b/deepspeed/moe/sharded_moe.py
@@ -20,7 +20,7 @@
 from time import perf_counter
 import torch
 from torch import Tensor
-import deepspeed.comm as dist
+from deepspeed import comm as dist
 from torch.nn import Module, ModuleList
 import torch.nn.functional as F
 
@@ -80,7 +80,7 @@ def gumbel_rsample(shape: Tuple, device: torch.device) -> Tensor:
     return gumbel(shape)
 
 
-import deepspeed.comm as dist
+from deepspeed import comm as dist
 
 # einsum dimensions: (g)roup, (s)equence, (e)xpert, (m)odel, (c)apacity
 # See https://arxiv.org/pdf/2006.16668.pdf for details.
diff --git a/deepspeed/moe/utils.py b/deepspeed/moe/utils.py
index 09922937e6d8..2b62a66eba09 100644
--- a/deepspeed/moe/utils.py
+++ b/deepspeed/moe/utils.py
@@ -1,6 +1,6 @@
 from typing import List, Tuple, Dict
 import torch
-import deepspeed.utils.groups as groups
+from deepspeed.utils import groups
 from .layer import MoE
 
 
diff --git a/deepspeed/ops/transformer/inference/moe_inference.py b/deepspeed/ops/transformer/inference/moe_inference.py
index ffc32744161b..830110cc5f05 100644
--- a/deepspeed/ops/transformer/inference/moe_inference.py
+++ b/deepspeed/ops/transformer/inference/moe_inference.py
@@ -16,7 +16,7 @@
 import torch.nn as nn
 from .transformer_inference import DeepSpeedSelfAttention, DeepSpeedInferenceConfig
 from ....moe.sharded_moe import TopKGate
-import deepspeed.comm as dist
+from deepspeed import comm as dist
 
 import torch.nn.functional as F
 
diff --git a/deepspeed/ops/transformer/inference/transformer_inference.py b/deepspeed/ops/transformer/inference/transformer_inference.py
index 12de114e55bd..7334d38bfd16 100755
--- a/deepspeed/ops/transformer/inference/transformer_inference.py
+++ b/deepspeed/ops/transformer/inference/transformer_inference.py
@@ -10,7 +10,7 @@
 import time
 from ... import op_builder
 import torch.nn as nn
-import deepspeed.comm as dist
+from deepspeed import comm as dist
 # Cuda modules will be imported if needed
 inference_cuda_module = None
 
diff --git a/deepspeed/runtime/activation_checkpointing/checkpointing.py b/deepspeed/runtime/activation_checkpointing/checkpointing.py
index c6a44fd97105..e3a945f2eff6 100644
--- a/deepspeed/runtime/activation_checkpointing/checkpointing.py
+++ b/deepspeed/runtime/activation_checkpointing/checkpointing.py
@@ -16,7 +16,7 @@
 import copy
 import torch
 import contextlib
-import deepspeed.comm as dist
+from deepspeed import comm as dist
 
 import mmap
 from torch import _C
diff --git a/deepspeed/runtime/bf16_optimizer.py b/deepspeed/runtime/bf16_optimizer.py
index 9676b3804f8d..9dd3bdd4e329 100644
--- a/deepspeed/runtime/bf16_optimizer.py
+++ b/deepspeed/runtime/bf16_optimizer.py
@@ -1,5 +1,5 @@
 import torch
-import deepspeed.comm as dist
+from deepspeed import comm as dist
 from deepspeed.runtime.constants import PIPE_REPLICATED
 from deepspeed.ops.op_builder import UtilsBuilder
 from deepspeed.runtime import ZeROOptimizer
diff --git a/deepspeed/runtime/comm/coalesced_collectives.py b/deepspeed/runtime/comm/coalesced_collectives.py
index 90f1a6383ae0..e92af044f53e 100644
--- a/deepspeed/runtime/comm/coalesced_collectives.py
+++ b/deepspeed/runtime/comm/coalesced_collectives.py
@@ -6,7 +6,7 @@
 
 import torch
 from torch import Tensor
-import deepspeed.comm as dist
+from deepspeed import comm as dist
 # NOTE: Use torch.distributed's ProcessGroup class until we have our own.
 from torch.distributed import ProcessGroup
 import torch.nn.functional
diff --git a/deepspeed/runtime/comm/nccl.py b/deepspeed/runtime/comm/nccl.py
index 5bdd58342087..ed80059a9067 100644
--- a/deepspeed/runtime/comm/nccl.py
+++ b/deepspeed/runtime/comm/nccl.py
@@ -3,7 +3,7 @@
 '''
 
 import torch
-import deepspeed.comm as dist
+from deepspeed import comm as dist
 import time
 import cupy
 import numpy as np
diff --git a/deepspeed/runtime/config.py b/deepspeed/runtime/config.py
index 76da68bc4190..8e694c3c9d34 100755
--- a/deepspeed/runtime/config.py
+++ b/deepspeed/runtime/config.py
@@ -25,7 +25,7 @@
 from .zero.constants import *
 from .activation_checkpointing.config import DeepSpeedActivationCheckpointingConfig
 
-import deepspeed.comm as dist
+from deepspeed import comm as dist
 
 from ..git_version_info import version as __version__
 from ..utils import logger
diff --git a/deepspeed/runtime/engine.py b/deepspeed/runtime/engine.py
index 60d1a6140ba7..4b4ebc7e636a 100644
--- a/deepspeed/runtime/engine.py
+++ b/deepspeed/runtime/engine.py
@@ -45,8 +45,8 @@
 from deepspeed.checkpoint.constants import OPTIMIZER_STATE_DICT
 from deepspeed.runtime.sparse_tensor import SparseTensor
 
-import deepspeed.runtime.lr_schedules as lr_schedules
-import deepspeed.utils.groups as groups
+from deepspeed.runtime import lr_schedules
+from deepspeed.utils import groups
 from deepspeed.runtime.utils import get_grad_norm
 from deepspeed.utils import logger, log_dist, instrument_w_nvtx
 from deepspeed.comm.comm import init_distributed
@@ -218,7 +218,7 @@ def __init__(
         self.use_ds_comm = False  # False --> Use torch.dist, True --> Use ds.comm backend.
 
         global dist
-        import deepspeed.comm as dist
+        from deepspeed import comm as dist
         self._is_gradient_accumulation_boundary = None
 
         # for debug purposes - can then debug print: debug_get_module_name(module)
diff --git a/deepspeed/runtime/fp16/fused_optimizer.py b/deepspeed/runtime/fp16/fused_optimizer.py
index c73fee6b0c04..479a0f7a2839 100755
--- a/deepspeed/runtime/fp16/fused_optimizer.py
+++ b/deepspeed/runtime/fp16/fused_optimizer.py
@@ -12,7 +12,7 @@
 from deepspeed.runtime.utils import get_global_norm, get_grad_norm, CheckOverflow, get_weight_norm
 from deepspeed.runtime.fp16.loss_scaler import INITIAL_LOSS_SCALE, SCALE_WINDOW, MIN_LOSS_SCALE
 from deepspeed.utils import groups, logger, log_dist
-import deepspeed.comm as dist
+from deepspeed import comm as dist
 from deepspeed.checkpoint.constants import OPTIMIZER_STATE_DICT, CLIP_GRAD
 
 
diff --git a/deepspeed/runtime/fp16/onebit/adam.py b/deepspeed/runtime/fp16/onebit/adam.py
index 52779e75e3cb..71805176ae41 100644
--- a/deepspeed/runtime/fp16/onebit/adam.py
+++ b/deepspeed/runtime/fp16/onebit/adam.py
@@ -6,7 +6,7 @@
 import importlib
 import numpy as np
 import time
-import deepspeed.comm as dist
+from deepspeed import comm as dist
 
 from deepspeed.utils.logging import logger
 
diff --git a/deepspeed/runtime/fp16/onebit/lamb.py b/deepspeed/runtime/fp16/onebit/lamb.py
index af524b01114b..aeff08b9861b 100644
--- a/deepspeed/runtime/fp16/onebit/lamb.py
+++ b/deepspeed/runtime/fp16/onebit/lamb.py
@@ -4,7 +4,7 @@
 import types
 import torch
 import numpy as np
-import deepspeed.comm as dist
+from deepspeed import comm as dist
 from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors
 
 
diff --git a/deepspeed/runtime/fp16/onebit/zoadam.py b/deepspeed/runtime/fp16/onebit/zoadam.py
index 04145d258645..11ccb404a9ff 100644
--- a/deepspeed/runtime/fp16/onebit/zoadam.py
+++ b/deepspeed/runtime/fp16/onebit/zoadam.py
@@ -6,7 +6,7 @@
 import importlib
 import numpy as np
 import time
-import deepspeed.comm as dist
+from deepspeed import comm as dist
 
 from deepspeed.utils.logging import logger
 
diff --git a/deepspeed/runtime/fp16/unfused_optimizer.py b/deepspeed/runtime/fp16/unfused_optimizer.py
index 67ad4fa9cb19..88f0df443405 100755
--- a/deepspeed/runtime/fp16/unfused_optimizer.py
+++ b/deepspeed/runtime/fp16/unfused_optimizer.py
@@ -14,7 +14,7 @@
 from deepspeed.runtime.fp16.loss_scaler import INITIAL_LOSS_SCALE, SCALE_WINDOW, MIN_LOSS_SCALE
 from deepspeed.utils import logger
 from deepspeed.checkpoint.constants import OPTIMIZER_STATE_DICT
-import deepspeed.comm as dist
+from deepspeed import comm as dist
 
 
 class FP16_UnfusedOptimizer(DeepSpeedOptimizer):
diff --git a/deepspeed/runtime/pipe/engine.py b/deepspeed/runtime/pipe/engine.py
index 2ea05d183ab1..8d8604126219 100644
--- a/deepspeed/runtime/pipe/engine.py
+++ b/deepspeed/runtime/pipe/engine.py
@@ -12,7 +12,7 @@
 import torch
 import torch.nn as nn
 import torch.optim as optim
-import deepspeed.comm as dist
+from deepspeed import comm as dist
 
 from deepspeed.utils.logging import logger
 from deepspeed.utils.timer import SynchronizedWallClockTimer, ThroughputTimer
diff --git a/deepspeed/runtime/pipe/module.py b/deepspeed/runtime/pipe/module.py
index 1469598b7700..f47806ee8673 100644
--- a/deepspeed/runtime/pipe/module.py
+++ b/deepspeed/runtime/pipe/module.py
@@ -9,7 +9,7 @@
 
 import torch
 import torch.nn as nn
-import deepspeed.comm as dist
+from deepspeed import comm as dist
 
 from deepspeed.utils import logger
 from .. import utils as ds_utils
diff --git a/deepspeed/runtime/pipe/p2p.py b/deepspeed/runtime/pipe/p2p.py
index 6b7ff6a7c786..d0c9c2f9f364 100644
--- a/deepspeed/runtime/pipe/p2p.py
+++ b/deepspeed/runtime/pipe/p2p.py
@@ -6,7 +6,7 @@
 import typing
 
 import torch
-import deepspeed.comm as dist
+from deepspeed import comm as dist
 
 # To query whether we have send/recv support
 from packaging.version import Version
diff --git a/deepspeed/runtime/pipe/topology.py b/deepspeed/runtime/pipe/topology.py
index 9475b6293cd9..954e73592943 100644
--- a/deepspeed/runtime/pipe/topology.py
+++ b/deepspeed/runtime/pipe/topology.py
@@ -2,7 +2,7 @@
 
 from deepspeed.utils import logger
 
-import deepspeed.comm as dist
+from deepspeed import comm as dist
 import sys
 
 from collections import namedtuple
diff --git a/deepspeed/runtime/swap_tensor/async_swapper.py b/deepspeed/runtime/swap_tensor/async_swapper.py
index 1e14381ab9f8..993eb0a162f2 100644
--- a/deepspeed/runtime/swap_tensor/async_swapper.py
+++ b/deepspeed/runtime/swap_tensor/async_swapper.py
@@ -6,7 +6,7 @@
 """
 import torch
 
-import deepspeed.comm as dist
+from deepspeed import comm as dist
 from deepspeed.utils.logging import logger
 from deepspeed.runtime.swap_tensor.utils import swap_out_tensors, SwapBuffer
 
diff --git a/deepspeed/runtime/swap_tensor/optimizer_utils.py b/deepspeed/runtime/swap_tensor/optimizer_utils.py
index bfcd3d2d4834..f34ff3a457a8 100644
--- a/deepspeed/runtime/swap_tensor/optimizer_utils.py
+++ b/deepspeed/runtime/swap_tensor/optimizer_utils.py
@@ -8,7 +8,7 @@
 import os
 import torch
 
-import deepspeed.comm as dist
+from deepspeed import comm as dist
 from deepspeed.utils.logging import logger
 from deepspeed.runtime.zero.offload_constants import *
 from deepspeed.runtime.swap_tensor.constants import *
diff --git a/deepspeed/runtime/swap_tensor/partitioned_optimizer_swapper.py b/deepspeed/runtime/swap_tensor/partitioned_optimizer_swapper.py
index e45292a35687..124500de888b 100644
--- a/deepspeed/runtime/swap_tensor/partitioned_optimizer_swapper.py
+++ b/deepspeed/runtime/swap_tensor/partitioned_optimizer_swapper.py
@@ -10,7 +10,7 @@
 
 from deepspeed.utils.logging import logger
 from deepspeed.ops.aio import AsyncIOBuilder
-import deepspeed.comm as dist
+from deepspeed import comm as dist
 
 from deepspeed.runtime.swap_tensor.constants import *
 from deepspeed.runtime.swap_tensor.utils import swap_in_tensors, swap_out_tensors, print_object, \
diff --git a/deepspeed/runtime/swap_tensor/partitioned_param_swapper.py b/deepspeed/runtime/swap_tensor/partitioned_param_swapper.py
index c944f30c0846..6d13bca87ac8 100644
--- a/deepspeed/runtime/swap_tensor/partitioned_param_swapper.py
+++ b/deepspeed/runtime/swap_tensor/partitioned_param_swapper.py
@@ -9,7 +9,7 @@
 import shutil
 from enum import Enum
 import torch
-import deepspeed.comm as dist
+from deepspeed import comm as dist
 
 from deepspeed.utils.logging import logger
 from deepspeed.ops.aio import AsyncIOBuilder
diff --git a/deepspeed/runtime/swap_tensor/pipelined_optimizer_swapper.py b/deepspeed/runtime/swap_tensor/pipelined_optimizer_swapper.py
index a98b3e76756d..598585078632 100644
--- a/deepspeed/runtime/swap_tensor/pipelined_optimizer_swapper.py
+++ b/deepspeed/runtime/swap_tensor/pipelined_optimizer_swapper.py
@@ -10,7 +10,7 @@
 
 from deepspeed.utils.logging import logger
 from deepspeed.ops.aio import AsyncIOBuilder
-import deepspeed.comm as dist
+from deepspeed import comm as dist
 
 from deepspeed.runtime.zero.offload_constants import *
 from deepspeed.runtime.swap_tensor.constants import *
diff --git a/deepspeed/runtime/swap_tensor/utils.py b/deepspeed/runtime/swap_tensor/utils.py
index df34fd6b9841..9d544e3d9b11 100644
--- a/deepspeed/runtime/swap_tensor/utils.py
+++ b/deepspeed/runtime/swap_tensor/utils.py
@@ -9,7 +9,7 @@
 import torch
 from deepspeed.utils.logging import logger
 
-import deepspeed.comm as dist
+from deepspeed import comm as dist
 
 from deepspeed.runtime.swap_tensor.constants import AIO_BLOCK_SIZE, AIO_QUEUE_DEPTH, \
     AIO_THREAD_COUNT, AIO_SINGLE_SUBMIT, AIO_OVERLAP_EVENTS
diff --git a/deepspeed/runtime/utils.py b/deepspeed/runtime/utils.py
index 654d31cf40ec..69660ac3c272 100755
--- a/deepspeed/runtime/utils.py
+++ b/deepspeed/runtime/utils.py
@@ -17,7 +17,7 @@
 
 import torch
 from torch._six import inf
-import deepspeed.comm as dist
+from deepspeed import comm as dist
 
 from deepspeed.utils import groups, logger
 from deepspeed.runtime.constants import PIPE_REPLICATED
diff --git a/deepspeed/runtime/zero/contiguous_memory_allocator.py b/deepspeed/runtime/zero/contiguous_memory_allocator.py
index 8fbe016acea1..883e58301e04 100644
--- a/deepspeed/runtime/zero/contiguous_memory_allocator.py
+++ b/deepspeed/runtime/zero/contiguous_memory_allocator.py
@@ -1,6 +1,6 @@
 import torch
 
-import deepspeed.comm as dist
+from deepspeed import comm as dist
 
 
 def print_rank_0(message):
diff --git a/deepspeed/runtime/zero/linear.py b/deepspeed/runtime/zero/linear.py
index 7803a86b49a2..02a585b02b01 100644
--- a/deepspeed/runtime/zero/linear.py
+++ b/deepspeed/runtime/zero/linear.py
@@ -18,7 +18,7 @@
 from torch.nn import init
 from torch.nn.modules.module import Module
 from deepspeed.runtime.utils import noop_decorator
-import deepspeed.comm as dist
+from deepspeed import comm as dist
 
 tensor_map = {}
 
diff --git a/deepspeed/runtime/zero/partition_parameters.py b/deepspeed/runtime/zero/partition_parameters.py
index 131d25faedc0..ab9ff507bfa0 100755
--- a/deepspeed/runtime/zero/partition_parameters.py
+++ b/deepspeed/runtime/zero/partition_parameters.py
@@ -15,7 +15,7 @@
 
 import torch
 from torch import Tensor
-import deepspeed.comm as dist
+from deepspeed import comm as dist
 from torch.nn import Module
 from torch.nn import Parameter
 
diff --git a/deepspeed/runtime/zero/partitioned_param_coordinator.py b/deepspeed/runtime/zero/partitioned_param_coordinator.py
index 4d9f263365d0..a9a02ac56679 100644
--- a/deepspeed/runtime/zero/partitioned_param_coordinator.py
+++ b/deepspeed/runtime/zero/partitioned_param_coordinator.py
@@ -12,7 +12,7 @@
 from torch.cuda import Event, Stream
 from torch.nn import Module, Parameter
 
-import deepspeed.comm as dist
+from deepspeed import comm as dist
 from deepspeed.utils.logging import logger
 from deepspeed.runtime.zero.partition_parameters import *
 from deepspeed.runtime.zero.offload_constants import *
diff --git a/deepspeed/runtime/zero/stage3.py b/deepspeed/runtime/zero/stage3.py
index 6b87cbc13e02..f1b0012b35fd 100755
--- a/deepspeed/runtime/zero/stage3.py
+++ b/deepspeed/runtime/zero/stage3.py
@@ -14,7 +14,7 @@
 import torch
 from torch.cuda import Event, Stream
 from torch.nn import Module, Parameter
-import deepspeed.comm as dist
+from deepspeed import comm as dist
 import math
 from torch._six import inf
 from torch.nn import Module
diff --git a/deepspeed/runtime/zero/stage_1_and_2.py b/deepspeed/runtime/zero/stage_1_and_2.py
index 49a154b80b55..7b683d12eeb4 100755
--- a/deepspeed/runtime/zero/stage_1_and_2.py
+++ b/deepspeed/runtime/zero/stage_1_and_2.py
@@ -3,7 +3,7 @@
 '''
 
 import torch
-import deepspeed.comm as dist
+from deepspeed import comm as dist
 from torch._six import inf
 from packaging import version as pkg_version
 
diff --git a/deepspeed/runtime/zero/utils.py b/deepspeed/runtime/zero/utils.py
index ff01aa79ef0a..665802ce8e93 100755
--- a/deepspeed/runtime/zero/utils.py
+++ b/deepspeed/runtime/zero/utils.py
@@ -2,7 +2,7 @@
 from typing import List
 
 import torch
-import deepspeed.comm as dist
+from deepspeed import comm as dist
 from deepspeed.utils import logger
 from deepspeed.ops.adam import DeepSpeedCPUAdam
 from deepspeed.ops.adam import FusedAdam
diff --git a/deepspeed/utils/groups.py b/deepspeed/utils/groups.py
index cdc94f8aab3d..f45b6681c77b 100644
--- a/deepspeed/utils/groups.py
+++ b/deepspeed/utils/groups.py
@@ -25,7 +25,7 @@
 """
 
 import torch
-import deepspeed.comm as dist
+from deepspeed import comm as dist
 
 from deepspeed.utils import logger, log_dist
 from deepspeed.utils.exceptions import DeprecatedException
diff --git a/deepspeed/utils/logging.py b/deepspeed/utils/logging.py
index 9f3a28f572ee..55636173b0a4 100644
--- a/deepspeed/utils/logging.py
+++ b/deepspeed/utils/logging.py
@@ -2,7 +2,7 @@
 import sys
 import os
 
-import deepspeed.comm as dist
+from deepspeed import comm as dist
 
 log_levels = {
     "debug": logging.DEBUG,
diff --git a/deepspeed/utils/timer.py b/deepspeed/utils/timer.py
index 91203778d10b..19671add031f 100755
--- a/deepspeed/utils/timer.py
+++ b/deepspeed/utils/timer.py
@@ -8,7 +8,7 @@
 import torch
 from numpy import mean
 from deepspeed.utils.logging import log_dist
-import deepspeed.comm as dist
+from deepspeed import comm as dist
 
 from deepspeed.utils import logger