[zero] reorganize zero/gemini folder structure

applications/Chat/coati/trainer/strategies/colossalai.py

@@ -14,17 +14,16 @@
 import colossalai
 from colossalai.logging import get_dist_logger
 from colossalai.nn.optimizer import CPUAdam, HybridAdam
-from colossalai.nn.parallel import ZeroDDP, zero_model_wrapper, zero_optim_wrapper
-from colossalai.nn.parallel.utils import get_static_torch_model
 from colossalai.tensor import ProcessGroup, ShardSpec
 from colossalai.utils import get_current_device
-from colossalai.utils.model.colo_init_context import ColoInitContext
-
-logger = get_dist_logger(__name__)
+from colossalai.zero import ColoInitContext, ZeroDDP, zero_model_wrapper, zero_optim_wrapper
+from colossalai.zero.gemini.utils import get_static_torch_model
 
 from .base import Strategy
 from .ddp import DDPStrategy
 
+logger = get_dist_logger(__name__)
+
 
 class ColossalAIStrategy(DDPStrategy):
     """

colossalai/auto_parallel/offload/base_offload_module.py

@@ -4,8 +4,8 @@
 import torch
 import torch.nn as nn
 
-from colossalai.gemini.tensor_utils import free_storage
 from colossalai.nn.parallel.data_parallel import _cast_float
+from colossalai.zero.legacy.gemini.tensor_utils import free_storage
 
 from .region_manager import RegionManager
 from .util import GlobalRuntimeInfo

colossalai/auto_parallel/offload/region.py

@@ -1,7 +1,10 @@
-from typing import List, Dict, Tuple
+from typing import Dict, List, Tuple
+
 import torch
 from torch.fx import Node
-from colossalai.gemini.tensor_utils import alloc_storage, free_storage
+
+from colossalai.zero.legacy.gemini.tensor_utils import alloc_storage, free_storage
+
 
 class Region:
     """
@@ -52,15 +55,13 @@ def init_param_data(self, pre_alloc_tensor: torch.Tensor = None):
         Map the parameters in the region to a contiguous memory space.
         """
 
-        self.fp16_data = torch.zeros(
-            self.param_num, dtype=torch.half, device='cuda')
+        self.fp16_data = torch.zeros(self.param_num, dtype=torch.half, device='cuda')
         offset = 0
         for param in self.fp16_params:
             param.data = param.data.cuda()
             p_num = param.data.numel()
             self.fp16_data[offset:offset + p_num].copy_(param.data.flatten())
-            param.data = self.fp16_data[offset:offset +
-                                               p_num].view(param.data.shape)
+            param.data = self.fp16_data[offset:offset + p_num].view(param.data.shape)
             self.param_to_range[param] = (offset, offset + p_num)
             offset += p_num
 
@@ -141,4 +142,4 @@ def split(self, cut_node_idx: int, cut_param_idx: int):
     def __update_params_ptr(self) -> None:
         for param in self.fp16_params:
             begin, end = self.param_to_range[param]
-            param.data = self.fp16_data[begin:end].view(param.data.shape)
+            param.data = self.fp16_data[begin:end].view(param.data.shape)

colossalai/booster/plugin/gemini_plugin.py

@@ -14,12 +14,12 @@
 
 from colossalai.checkpoint_io import CheckpointIO, GeneralCheckpointIO
 from colossalai.cluster import DistCoordinator
-from colossalai.gemini.memory_tracer import MemStats
 from colossalai.interface import ModelWrapper, OptimizerWrapper
-from colossalai.nn.parallel import GeminiDDP, zero_model_wrapper, zero_optim_wrapper
 from colossalai.tensor.colo_parameter import ColoParameter
 from colossalai.utils import get_current_device
-from colossalai.utils.model.colo_init_context import _convert_to_coloparam
+from colossalai.zero import GeminiDDP, zero_model_wrapper, zero_optim_wrapper
+from colossalai.zero.gemini.colo_init_context import _convert_to_coloparam
+from colossalai.zero.gemini.memory_tracer import MemStats
 
 from .plugin_base import Plugin
 

colossalai/engine/_base_engine.py

@@ -10,8 +10,8 @@
 
 from colossalai.engine.gradient_handler import BaseGradientHandler
 from colossalai.engine.schedule import BaseSchedule, InterleavedPipelineSchedule, NonPipelineSchedule, PipelineSchedule
-from colossalai.gemini.ophooks import BaseOpHook, register_ophooks_recursively
 from colossalai.logging import get_dist_logger
+from colossalai.zero.legacy.gemini import BaseOpHook, register_ophooks_recursively
 
 
 class Engine:

colossalai/engine/schedule/_pipeline_schedule.py

@@ -157,7 +157,7 @@ def load_micro_batch(self):
         return self._move_to_device(mciro_batch_data)
 
     def pre_processing(self, engine):
-        from colossalai.zero.sharded_model.sharded_model_v2 import ShardedModelV2
+        from colossalai.zero.legacy import ShardedModelV2
 
         # TODO: remove this after testing new zero with pipeline parallelism
         model = engine.model

colossalai/initialize.py

@@ -29,13 +29,12 @@
     PipelineSchedule,
     get_tensor_shape,
 )
-from colossalai.gemini.ophooks import BaseOpHook
 from colossalai.logging import get_dist_logger
 from colossalai.nn.optimizer.colossalai_optimizer import ColossalaiOptimizer
 from colossalai.utils import get_current_device, is_using_ddp, is_using_pp, is_using_sequence, sync_model_param
 from colossalai.utils.moe import sync_moe_model_param
-from colossalai.zero import convert_to_zero_v2
-from colossalai.zero.sharded_optim.sharded_optim_v2 import ShardedOptimizerV2
+from colossalai.zero.legacy import ShardedOptimizerV2, convert_to_zero_v2
+from colossalai.zero.legacy.gemini.ophooks import BaseOpHook
 
 
 def get_default_parser():

colossalai/nn/layer/moe/experts.py

@@ -9,7 +9,7 @@
 from colossalai.context import ParallelMode, seed
 from colossalai.context.moe_context import MOE_CONTEXT
 from colossalai.utils import get_current_device
-from colossalai.zero.init_ctx import no_shard_zero_decrator
+from colossalai.zero.legacy.init_ctx import no_shard_zero_decrator
 
 
 class MoeExperts(nn.Module):

colossalai/nn/layer/moe/layers.py

@@ -18,7 +18,7 @@
 from colossalai.nn.layer.moe.routers import MoeRouter, Top1Router, Top2Router
 from colossalai.nn.layer.moe.utils import NormalNoiseGenerator, UniformNoiseGenerator
 from colossalai.utils import get_current_device
-from colossalai.zero.init_ctx import no_shard_zero_context, no_shard_zero_decrator
+from colossalai.zero.legacy.init_ctx import no_shard_zero_context, no_shard_zero_decrator
 
 
 @no_shard_zero_decrator(is_replicated=True)

colossalai/nn/parallel/__init__.py

@@ -1,5 +1,5 @@
-from .data_parallel import ColoDDP, ZeroDDP
-from .gemini_parallel import GeminiDDP
-from .zero_wrapper import zero_model_wrapper, zero_optim_wrapper
+from .data_parallel import ColoDDP
 
-__all__ = ['ColoDDP', 'ZeroDDP', 'GeminiDDP', 'zero_model_wrapper', 'zero_optim_wrapper']
+__all__ = [
+    'ColoDDP',
+]