[legacy] move gpc to legacy

colossalai/amp/naive_amp/_fp16_optimizer.py

@@ -6,9 +6,9 @@
 from torch.distributed import ProcessGroup
 from torch.optim import Optimizer
 
-from colossalai.core import global_context as gpc
 from colossalai.kernel.op_builder import FusedOptimBuilder
 from colossalai.legacy.context import ParallelMode
+from colossalai.legacy.core import global_context as gpc
 from colossalai.logging import get_dist_logger
 from colossalai.utils import clip_grad_norm_fp32, copy_tensor_parallel_attributes, multi_tensor_applier
 

colossalai/amp/naive_amp/naive_amp.py

@@ -11,8 +11,8 @@
 from torch.distributed import ReduceOp
 from torch.optim import Optimizer
 
-from colossalai.core import global_context as gpc
 from colossalai.legacy.context import ParallelMode
+from colossalai.legacy.core import global_context as gpc
 from colossalai.nn.optimizer import ColossalaiOptimizer
 
 from ._fp16_optimizer import FP16Optimizer

colossalai/amp/torch_amp/_grad_scaler.py

@@ -13,8 +13,8 @@
 from packaging import version
 from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors
 
-from colossalai.core import global_context as gpc
 from colossalai.legacy.context import ParallelMode
+from colossalai.legacy.core import global_context as gpc
 
 
 class _MultiDeviceReplicator(object):

colossalai/cli/benchmark/benchmark.py

@@ -6,9 +6,9 @@
 
 import colossalai
 from colossalai.cli.benchmark.utils import find_all_configs, get_batch_data, profile_model
-from colossalai.core import global_context as gpc
 from colossalai.legacy.context import Config
 from colossalai.legacy.context.random import reset_seeds
+from colossalai.legacy.core import global_context as gpc
 from colossalai.logging import disable_existing_loggers, get_dist_logger
 from colossalai.testing import free_port
 from colossalai.utils import MultiTimer

colossalai/communication/collective.py

@@ -6,8 +6,8 @@
 from torch import Tensor
 from torch.distributed import ReduceOp
 
-from colossalai.core import global_context as gpc
 from colossalai.legacy.context import ParallelMode
+from colossalai.legacy.core import global_context as gpc
 
 _all_gather_func = dist._all_gather_base \
     if "all_gather_into_tensor" not in dir(dist) else dist.all_gather_into_tensor

colossalai/communication/p2p.py

@@ -8,8 +8,8 @@
 import torch
 import torch.distributed as dist
 
-from colossalai.core import global_context as gpc
 from colossalai.legacy.context.parallel_mode import ParallelMode
+from colossalai.legacy.core import global_context as gpc
 from colossalai.utils import get_current_device
 
 from .utils import gather_split_1d_tensor, split_tensor_into_1d_equal_chunks

colossalai/communication/p2p_v2.py

@@ -10,8 +10,8 @@
 from torch.distributed import ProcessGroupNCCL
 from torch.distributed import distributed_c10d as c10d
 
-from colossalai.core import global_context as gpc
 from colossalai.legacy.context.parallel_mode import ParallelMode
+from colossalai.legacy.core import global_context as gpc
 
 TensorShape = Union[torch.Size, List[int], Tuple[int]]
 _pg_manager = {}

colossalai/communication/ring.py

@@ -3,8 +3,8 @@
 
 import torch
 
-from colossalai.core import global_context as gpc
 from colossalai.legacy.context.parallel_mode import ParallelMode
+from colossalai.legacy.core import global_context as gpc
 from colossalai.utils import get_current_device, synchronize
 
 

colossalai/communication/utils.py

@@ -3,8 +3,8 @@
 import torch
 import torch.distributed as dist
 
-from colossalai.core import global_context as gpc
 from colossalai.legacy.context.parallel_mode import ParallelMode
+from colossalai.legacy.core import global_context as gpc
 from colossalai.utils import get_current_device
 
 TensorShape = Union[torch.Size, List[int], Tuple[int]]

colossalai/engine/gradient_handler/_data_parallel_gradient_handler.py

@@ -1,7 +1,7 @@
-from colossalai.core import global_context as gpc
+from colossalai.legacy.context import ParallelMode
+from colossalai.legacy.core import global_context as gpc
 from colossalai.registry import GRADIENT_HANDLER
 
-from ...context.parallel_mode import ParallelMode
 from ._base_gradient_handler import BaseGradientHandler
 from .utils import bucket_allreduce
 

colossalai/engine/gradient_handler/_moe_gradient_handler.py

@@ -1,9 +1,9 @@
-from colossalai.core import global_context as gpc
+from colossalai.legacy.context import ParallelMode
 from colossalai.legacy.context.moe_context import MOE_CONTEXT
+from colossalai.legacy.core import global_context as gpc
 from colossalai.registry import GRADIENT_HANDLER
 from colossalai.utils.moe import get_moe_epsize_param_dict
 
-from ...context.parallel_mode import ParallelMode
 from ._base_gradient_handler import BaseGradientHandler
 from .utils import bucket_allreduce
 

colossalai/engine/gradient_handler/_pipeline_parallel_gradient_handler.py

@@ -6,7 +6,7 @@
 import torch.distributed as dist
 from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors
 
-from colossalai.core import global_context as gpc
+from colossalai.legacy.core import global_context as gpc
 from colossalai.registry import GRADIENT_HANDLER
 
 from ._base_gradient_handler import BaseGradientHandler

colossalai/engine/gradient_handler/_sequence_parallel_gradient_handler.py

@@ -1,7 +1,7 @@
-from colossalai.core import global_context as gpc
+from colossalai.legacy.context import ParallelMode
+from colossalai.legacy.core import global_context as gpc
 from colossalai.registry import GRADIENT_HANDLER
 
-from ...context.parallel_mode import ParallelMode
 from ._base_gradient_handler import BaseGradientHandler
 from .utils import bucket_allreduce
 

colossalai/engine/schedule/_pipeline_schedule.py

@@ -8,8 +8,8 @@
 
 import colossalai.communication as comm
 from colossalai.amp.naive_amp import NaiveAMPModel
-from colossalai.core import global_context as gpc
 from colossalai.legacy.context.parallel_mode import ParallelMode
+from colossalai.legacy.core import global_context as gpc
 from colossalai.logging import get_dist_logger
 from colossalai.utils import switch_virtual_pipeline_parallel_rank
 from colossalai.utils.cuda import get_current_device

colossalai/engine/schedule/_pipeline_schedule_v2.py

@@ -7,8 +7,8 @@
 
 import colossalai.communication.p2p_v2 as comm
 from colossalai import engine
-from colossalai.core import global_context as gpc
 from colossalai.legacy.context.parallel_mode import ParallelMode
+from colossalai.legacy.core import global_context as gpc
 from colossalai.utils.cuda import get_current_device
 
 from ._pipeline_schedule import PipelineSchedule

colossalai/initialize.py

@@ -20,7 +20,6 @@
 from colossalai.amp import AMP_TYPE, convert_to_amp
 from colossalai.amp.naive_amp import NaiveAMPModel
 from colossalai.builder.builder import build_gradient_handler
-from colossalai.core import global_context as gpc
 from colossalai.engine import Engine
 from colossalai.engine.gradient_accumulation import accumulate_gradient
 from colossalai.engine.schedule import (
@@ -31,6 +30,7 @@
 )
 from colossalai.legacy.context import Config, ConfigException, ParallelMode
 from colossalai.legacy.context.moe_context import MOE_CONTEXT
+from colossalai.legacy.core import global_context as gpc
 from colossalai.logging import get_dist_logger
 from colossalai.nn.optimizer.colossalai_optimizer import ColossalaiOptimizer
 from colossalai.utils import (

colossalai/legacy/context/moe_context.py

@@ -9,7 +9,7 @@
 
 
 def _check_sanity():
-    from colossalai.core import global_context as gpc
+    from colossalai.legacy.core import global_context as gpc
     if gpc.tensor_parallel_size > 1 or gpc.pipeline_parallel_size > 1:
         raise NotImplementedError("Moe is not compatible with tensor or "
                                   "pipeline parallel at present.")
@@ -61,7 +61,7 @@ def setup(self, seed: int, use_kernel_optim: bool = True):
 
         self.world_size = dist.get_world_size()
 
-        from colossalai.core import global_context as gpc
+        from colossalai.legacy.core import global_context as gpc
         self.max_ep_size = gpc.config.get('max_ep_size', self.world_size)
         assert self.world_size % self.max_ep_size == 0, \
             "Maximum expert parallel size must be a factor of the number of GPUs"

colossalai/legacy/context/random/_helper.py

@@ -161,7 +161,7 @@ def wrapper(*args, **kwargs):
 
 def moe_set_seed(seed):
     if torch.cuda.is_available():
-        from colossalai.core import global_context as gpc
+        from colossalai.legacy.core import global_context as gpc
         global_rank = gpc.get_global_rank()
         diff_seed = seed + global_rank
         add_seed(ParallelMode.TENSOR, diff_seed, True)

colossalai/nn/layer/base_layer.py

@@ -5,8 +5,8 @@
 
 import torch.nn as nn
 
-from colossalai.core import global_context as gpc
 from colossalai.legacy.context import ParallelMode
+from colossalai.legacy.core import global_context as gpc
 
 
 class ParallelLayer(nn.Module):

colossalai/nn/layer/parallel_1d/_operation.py

@@ -1,7 +1,7 @@
 import torch
 import torch.distributed as dist
 
-from colossalai.core import global_context as gpc
+from colossalai.legacy.core import global_context as gpc
 
 try:
     import fused_mix_prec_layer_norm_cuda
@@ -40,10 +40,10 @@ def backward(ctx, grad_output):
         input_, weight_, bias_, mean, invvar = ctx.saved_tensors
         grad_input = grad_weight = grad_bias = None
         grad_input, grad_weight, grad_bias \
-          = fused_mix_prec_layer_norm_cuda.backward_affine(
-            grad_output.contiguous(), mean, invvar,
-            input_, ctx.normalized_shape,
-            weight_, bias_, ctx.eps)
+            = fused_mix_prec_layer_norm_cuda.backward_affine(
+                grad_output.contiguous(), mean, invvar,
+                input_, ctx.normalized_shape,
+                weight_, bias_, ctx.eps)
 
         return grad_input, grad_weight, grad_bias, None, None
 

colossalai/nn/layer/parallel_1d/_utils.py

@@ -3,8 +3,9 @@
 
 import torch
 import torch.distributed as dist
-from colossalai.core import global_context as gpc
+
 from colossalai.global_variables import tensor_parallel_env as env
+from colossalai.legacy.core import global_context as gpc
 
 from ..utils import divide
 
@@ -124,7 +125,7 @@ def backward(ctx, grad_output):
 class _SplitForwardGatherBackward(torch.autograd.Function):
     """
     Split the input and keep only the corresponding chuck to the rank.
-    
+
     Args:
         input_: input matrix.
         parallel_mode: parallel mode.

colossalai/nn/layer/parallel_1d/layers.py

@@ -11,10 +11,10 @@
 from torch.nn.parameter import Parameter
 
 from colossalai.communication import broadcast
-from colossalai.core import global_context as gpc
 from colossalai.global_variables import tensor_parallel_env as env
 from colossalai.kernel import LayerNorm
 from colossalai.legacy.context import ParallelMode, seed
+from colossalai.legacy.core import global_context as gpc
 from colossalai.nn import init as init
 from colossalai.registry import LAYERS
 from colossalai.utils.checkpointing import (

colossalai/nn/layer/parallel_2d/_operation.py

@@ -6,9 +6,9 @@
 from torch.cuda.amp import custom_bwd, custom_fwd
 
 from colossalai.communication.collective import all_gather, all_reduce, reduce, reduce_scatter
-from colossalai.core import global_context as gpc
 from colossalai.global_variables import tensor_parallel_env as env
 from colossalai.legacy.context.parallel_mode import ParallelMode
+from colossalai.legacy.core import global_context as gpc
 from colossalai.utils import get_current_device
 
 

colossalai/nn/layer/parallel_2d/_utils.py

@@ -1,6 +1,6 @@
-from colossalai.core import global_context as gpc
 from colossalai.global_variables import tensor_parallel_env as env
 from colossalai.legacy.context.parallel_mode import ParallelMode
+from colossalai.legacy.core import global_context as gpc
 
 
 def get_summa_dim_from_env() -> int:

colossalai/nn/layer/parallel_2d/layers.py

@@ -9,9 +9,9 @@
 from torch.nn import Parameter
 
 from colossalai.communication import broadcast
-from colossalai.core import global_context as gpc
 from colossalai.global_variables import tensor_parallel_env as env
 from colossalai.legacy.context import ParallelMode, seed
+from colossalai.legacy.core import global_context as gpc
 from colossalai.nn import init as init
 from colossalai.registry import LAYERS
 from colossalai.utils.checkpointing import gather_tensor_parallel_state_dict, partition_tensor_parallel_state_dict

colossalai/nn/layer/parallel_2p5d/_operation.py

@@ -6,8 +6,8 @@
 from torch.cuda.amp import custom_bwd, custom_fwd
 
 from colossalai.communication.collective import all_gather, all_reduce, reduce_scatter
-from colossalai.core import global_context as gpc
 from colossalai.legacy.context.parallel_mode import ParallelMode
+from colossalai.legacy.core import global_context as gpc
 from colossalai.utils import get_current_device
 
 

colossalai/nn/layer/parallel_2p5d/_utils.py

@@ -1,6 +1,6 @@
-from colossalai.core import global_context as gpc
 from colossalai.global_variables import tensor_parallel_env as env
 from colossalai.legacy.context.parallel_mode import ParallelMode
+from colossalai.legacy.core import global_context as gpc
 
 
 def get_tesseract_dim_dep_from_env():

colossalai/nn/layer/parallel_2p5d/layers.py

@@ -9,9 +9,9 @@
 from torch.nn import Parameter
 
 from colossalai.communication import broadcast
-from colossalai.core import global_context as gpc
 from colossalai.global_variables import tensor_parallel_env as env
 from colossalai.legacy.context import ParallelMode, seed
+from colossalai.legacy.core import global_context as gpc
 from colossalai.nn import init as init
 from colossalai.registry import LAYERS
 from colossalai.utils.checkpointing import (

colossalai/nn/layer/parallel_3d/_operation.py

@@ -9,8 +9,8 @@
 
 from colossalai.communication import all_gather, all_reduce, broadcast, reduce, reduce_scatter
 from colossalai.constants import INPUT_GROUP_3D, WEIGHT_GROUP_3D
-from colossalai.core import global_context as gpc
 from colossalai.legacy.context.parallel_mode import ParallelMode
+from colossalai.legacy.core import global_context as gpc
 
 from ._utils import get_parallel_mode_from_env, push_async_grad
 

colossalai/nn/layer/parallel_3d/_utils.py

@@ -5,8 +5,8 @@
 from torch import Tensor
 
 from colossalai.constants import INPUT_GROUP_3D, INPUT_X_WEIGHT_3D, OUTPUT_GROUP_3D, OUTPUT_X_WEIGHT_3D, WEIGHT_GROUP_3D
-from colossalai.core import global_context as gpc
 from colossalai.global_variables import tensor_parallel_env as env
+from colossalai.legacy.core import global_context as gpc
 
 
 def get_depth_from_env() -> int:

colossalai/nn/layer/parallel_3d/layers.py

@@ -10,9 +10,9 @@
 
 from colossalai.communication import all_reduce, broadcast
 from colossalai.constants import INPUT_GROUP_3D, INPUT_X_WEIGHT_3D, OUTPUT_GROUP_3D, OUTPUT_X_WEIGHT_3D, WEIGHT_GROUP_3D
-from colossalai.core import global_context as gpc
 from colossalai.global_variables import tensor_parallel_env as env
 from colossalai.legacy.context import ParallelMode, seed
+from colossalai.legacy.core import global_context as gpc
 from colossalai.nn import init as init
 from colossalai.nn.layer.base_layer import ParallelLayer
 from colossalai.registry import LAYERS

colossalai/nn/layer/parallel_sequence/_operation.py

@@ -6,8 +6,8 @@
 from torch.cuda.amp import custom_bwd, custom_fwd
 
 from colossalai.communication import ring_forward
-from colossalai.core import global_context as gpc
 from colossalai.legacy.context.parallel_mode import ParallelMode
+from colossalai.legacy.core import global_context as gpc
 from colossalai.nn.layer.parallel_sequence._utils import _calc_current_device_range, _calc_incoming_device_range
 from colossalai.utils import get_current_device
 

colossalai/nn/layer/parallel_sequence/layers.py

@@ -9,11 +9,11 @@
 from torch.nn import Parameter
 
 import colossalai
-from colossalai.core import global_context as gpc
 from colossalai.kernel import FusedScaleMaskSoftmax
 from colossalai.kernel.cuda_native.scaled_softmax import AttnMaskType
 from colossalai.legacy.context import seed
 from colossalai.legacy.context.parallel_mode import ParallelMode
+from colossalai.legacy.core import global_context as gpc
 from colossalai.nn.layer.parallel_sequence._operation import RingAV, RingQK
 from colossalai.registry import LAYERS
 

colossalai/nn/layer/wrapper/pipeline_wrapper.py

@@ -3,8 +3,8 @@
 import torch.distributed as dist
 import torch.nn as nn
 
-from colossalai.core import global_context as gpc
 from colossalai.legacy.context import ParallelMode
+from colossalai.legacy.core import global_context as gpc
 
 
 class PipelineSharedModuleWrapper:

colossalai/nn/loss/loss_1d.py

@@ -3,8 +3,8 @@
 from torch.cuda.amp import custom_bwd, custom_fwd
 from torch.nn.modules.loss import _Loss
 
-from colossalai.core import global_context as gpc
 from colossalai.legacy.context import ParallelMode
+from colossalai.legacy.core import global_context as gpc
 from colossalai.registry import LOSSES