[zero] reorganize zero/gemini folder structure (#3424)

* [zero] refactor low-level zero folder structure

* [zero] fix legacy zero import path

* [zero] fix legacy zero import path

* [zero] remove useless import

* [zero] refactor gemini folder structure

* [zero] refactor gemini folder structure

* [zero] refactor legacy zero import path

* [zero] refactor gemini folder structure

* [zero] refactor gemini folder structure

* [zero] refactor gemini folder structure

* [zero] refactor legacy zero import path

* [zero] fix test import path

* [zero] fix test

* [zero] fix circular import

* [zero] update import

colossalai/nn/parallel/zero_wrapper.py → colossalai/zero/wrapper.py

@@ -4,7 +4,7 @@ from typing import Dict, Optional
 import torch
 import torch.nn as nn
 
-from .gemini_parallel import GeminiDDP
+from .gemini import GeminiDDP
 
 
 def zero_model_wrapper(model: nn.Module, zero_stage: int = 1, gemini_config: Optional[Dict] = None):
@@ -99,11 +99,11 @@ def zero_optim_wrapper(model: nn.Module,
     config_dict['max_scale'] = max_scale
 
     if zero_stage in [1, 2]:
-        from colossalai.zero.sharded_optim.low_level_optim import LowLevelZeroOptimizer
+        from colossalai.zero.low_level import LowLevelZeroOptimizer
         config_dict['partition_grad'] = zero_stage == 2
         config_dict['clip_grad_norm'] = max_norm
         return LowLevelZeroOptimizer(optimizer, **config_dict)
     else:
-        from colossalai.nn.optimizer.zero_optimizer import ZeroOptimizer
+        from colossalai.zero.gemini.gemini_optimizer import ZeroOptimizer
         config_dict['clipping_norm'] = max_norm
         return ZeroOptimizer(optimizer, model, **config_dict)

docs/source/en/features/nvme_offload.md

@@ -78,7 +78,7 @@ from transformers.models.gpt2.modeling_gpt2 import GPT2LMHeadModel
 
 import colossalai
 from colossalai.nn.optimizer import HybridAdam
-from colossalai.nn.parallel import zero_model_wrapper, zero_optim_wrapper
+from colossalai.zero import zero_model_wrapper, zero_optim_wrapper
 from colossalai.utils.model.colo_init_context import ColoInitContext
 ```
 

docs/source/zh-Hans/features/nvme_offload.md

@@ -77,7 +77,7 @@ from transformers.models.gpt2.configuration_gpt2 import GPT2Config
 from transformers.models.gpt2.modeling_gpt2 import GPT2LMHeadModel
 import colossalai
 from colossalai.nn.optimizer import HybridAdam
-from colossalai.nn.parallel import zero_model_wrapper, zero_optim_wrapper
+from colossalai.zero import zero_model_wrapper, zero_optim_wrapper
 from colossalai.utils.model.colo_init_context import ColoInitContext
 ```
 

examples/images/dreambooth/debug.py

@@ -5,7 +5,7 @@ torchrun --standalone --nproc_per_node=1 debug.py
 from diffusers import AutoencoderKL
 
 import colossalai
-from colossalai.utils.model.colo_init_context import ColoInitContext, post_process_colo_init_ctx
+from colossalai.zero import ColoInitContext, post_process_colo_init_ctx
 
 path = "/data/scratch/diffuser/stable-diffusion-v1-4"
 

examples/images/dreambooth/train_dreambooth_colossalai.py

@@ -21,10 +21,9 @@ import colossalai
 from colossalai.context.parallel_mode import ParallelMode
 from colossalai.core import global_context as gpc
 from colossalai.logging import disable_existing_loggers, get_dist_logger
-from colossalai.nn.optimizer.gemini_optimizer import GeminiAdamOptimizer
-from colossalai.nn.parallel.utils import get_static_torch_model
 from colossalai.utils import get_current_device
-from colossalai.utils.model.colo_init_context import ColoInitContext
+from colossalai.zero import ColoInitContext, GeminiAdamOptimizer
+from colossalai.zero.gemini import get_static_torch_model
 
 disable_existing_loggers()
 logger = get_dist_logger()

examples/images/dreambooth/train_dreambooth_colossalai_lora.py

@@ -23,10 +23,9 @@ import colossalai
 from colossalai.context.parallel_mode import ParallelMode
 from colossalai.core import global_context as gpc
 from colossalai.logging import disable_existing_loggers, get_dist_logger
-from colossalai.nn.optimizer.gemini_optimizer import GeminiAdamOptimizer
-from colossalai.nn.parallel.utils import get_static_torch_model
 from colossalai.utils import get_current_device
-from colossalai.utils.model.colo_init_context import ColoInitContext
+from colossalai.zero import ColoInitContext, GeminiAdamOptimizer
+from colossalai.zero.gemini import get_static_torch_model
 
 disable_existing_loggers()
 logger = get_dist_logger()

examples/images/vit/test_vit.py

@@ -18,7 +18,7 @@ from colossalai.tensor import ComputePattern, ComputeSpec, DistSpecManager, Proc
 from colossalai.testing import rerun_if_address_is_in_use
 from colossalai.utils import free_port
 from colossalai.utils.cuda import get_current_device
-from colossalai.utils.model.colo_init_context import ColoInitContext
+from colossalai.zero import ColoInitContext
 
 
 def set_seed(seed):

examples/images/vit/train.py

@@ -19,7 +19,7 @@ from colossalai.nn.optimizer import HybridAdam
 from colossalai.nn.parallel.data_parallel import ColoDDP
 from colossalai.tensor import ComputePattern, ComputeSpec, DistSpecManager, ProcessGroup, ShardSpec
 from colossalai.utils import get_current_device
-from colossalai.utils.model.colo_init_context import ColoInitContext
+from colossalai.zero import ColoInitContext
 
 
 def init_1d_row_for_linear_weight_spec(model, world_size: int):

examples/language/bert/train_bert_demo.py

@@ -12,10 +12,9 @@ from transformers import AlbertConfig, AlbertForSequenceClassification, BertConf
 import colossalai
 from colossalai.logging import disable_existing_loggers, get_dist_logger
 from colossalai.nn.optimizer import HybridAdam
-from colossalai.nn.parallel import zero_model_wrapper, zero_optim_wrapper
 from colossalai.tensor import ColoParameter, ComputePattern, ComputeSpec, ProcessGroup, ReplicaSpec, ShardSpec
 from colossalai.utils import get_current_device
-from colossalai.utils.model.colo_init_context import ColoInitContext
+from colossalai.zero import ColoInitContext, zero_model_wrapper, zero_optim_wrapper
 
 CAI_VERSION = colossalai.__version__
 

examples/language/gpt/gemini/train_gpt_demo.py

@@ -13,10 +13,9 @@ from torch.nn.parallel import DistributedDataParallel as DDP
 import colossalai
 from colossalai.logging import disable_existing_loggers, get_dist_logger
 from colossalai.nn.optimizer import HybridAdam
-from colossalai.nn.parallel import zero_model_wrapper, zero_optim_wrapper
 from colossalai.tensor import ColoParameter, ComputePattern, ComputeSpec, ProcessGroup, ReplicaSpec, ShardSpec
 from colossalai.utils import get_current_device
-from colossalai.utils.model.colo_init_context import ColoInitContext
+from colossalai.zero import ColoInitContext, zero_model_wrapper, zero_optim_wrapper
 
 CAI_VERSION = colossalai.__version__
 

examples/language/opt/train_gemini_opt.py

@@ -34,12 +34,9 @@ from transformers.utils.versions import require_version
 
 import colossalai
 from colossalai.logging import disable_existing_loggers, get_dist_logger
-from colossalai.nn.optimizer.gemini_optimizer import GeminiAdamOptimizer
-from colossalai.nn.parallel import GeminiDDP
-from colossalai.utils import get_current_device
-from colossalai.utils.model.colo_init_context import ColoInitContext
-
 from colossalai.tensor import ProcessGroup, ShardSpec
+from colossalai.utils import get_current_device
+from colossalai.zero import ColoInitContext, GeminiAdamOptimizer, GeminiDDP
 
 
 def get_data(batch_size, seq_len, vocab_size):
@@ -179,13 +176,15 @@ def main():
     # build model
     if args.model_name_or_path is None:
         logger.info("Train a new model from scratch", ranks=[0])
-        with ColoInitContext(device=init_dev, dtype=torch.half,
+        with ColoInitContext(device=init_dev,
+                             dtype=torch.half,
                              default_dist_spec=default_dist_spec,
                              default_pg=shard_pg):
             model = OPTForCausalLM(config)
     else:
         logger.info("Finetune a pre-trained model", ranks=[0])
-        with ColoInitContext(device=init_dev, dtype=torch.half,
+        with ColoInitContext(device=init_dev,
+                             dtype=torch.half,
                              default_dist_spec=default_dist_spec,
                              default_pg=shard_pg):
             model = OPTForCausalLM.from_pretrained(args.model_name_or_path,
@@ -198,8 +197,11 @@ def main():
 
     numel = sum([p.numel() for p in model.parameters()])
     PLACEMENT_POLICY = 'cpu'
-    model = GeminiDDP(model, device=get_current_device(), placement_policy=PLACEMENT_POLICY,
-                      pin_memory=True, strict_ddp_mode=args.shardinit)
+    model = GeminiDDP(model,
+                      device=get_current_device(),
+                      placement_policy=PLACEMENT_POLICY,
+                      pin_memory=True,
+                      strict_ddp_mode=args.shardinit)
     optimizer = GeminiAdamOptimizer(model, lr=args.learning_rate, initial_scale=2**14, gpu_margin_mem_ratio=0.0)
 
     SEQ_LEN = 1024

examples/language/palm/train.py

@@ -15,11 +15,9 @@ from torch.utils.data import DataLoader, Dataset
 
 import colossalai
 from colossalai.logging import disable_existing_loggers, get_dist_logger
-from colossalai.nn.optimizer.gemini_optimizer import GeminiAdamOptimizer
-from colossalai.nn.parallel import ZeroDDP
 from colossalai.tensor import ColoParameter, ComputePattern, ComputeSpec, ProcessGroup, ReplicaSpec, ShardSpec
 from colossalai.utils import MultiTimer, get_current_device
-from colossalai.utils.model.colo_init_context import ColoInitContext
+from colossalai.zero import ColoInitContext, GeminiAdamOptimizer, ZeroDDP
 
 # constants
 
@@ -127,7 +125,7 @@ def gemini_zero_dpp(model: torch.nn.Module, pg: ProcessGroup, placememt_policy:
     return model
 
 
-## Parameter Sharding Strategies for Tensor Parallelism
+# Parameter Sharding Strategies for Tensor Parallelism
 def split_param_single_dim_tp1d(dim: int, param: ColoParameter, pg: ProcessGroup):
     spec = (ShardSpec([dim], [pg.tp_world_size()]), ComputeSpec(ComputePattern.TP1D))
     param.set_tensor_spec(*spec)
@@ -232,7 +230,7 @@ if args.distplan == "colossalai":
     tensor_parallelize(model, pg)
     model = gemini_zero_dpp(model, pg, args.placement)
 
-    #optimizer
+    # optimizer
 
     #optimizer = GeminiAdamOptimizer(model, lr=1e-7, initial_scale=2**5)
     optimizer = GeminiAdamOptimizer(model, lr=LEARNING_RATE, initial_scale=2**5)

examples/language/roberta/pretraining/run_pretraining.py

-import colossalai
 import math
+import os
+import time
+from functools import partial
+
 import torch
-from colossalai.context import ParallelMode
-from colossalai.core import global_context as gpc
-import colossalai.nn as col_nn
 from arguments import parse_args
-from pretrain_utils import get_model, get_optimizer, get_lr_scheduler, save_ckpt
-from utils.exp_util import get_tflops, get_mem_info, throughput_calculator, log_args
-from utils.global_vars import set_global_variables, get_timers, get_tensorboard_writer
-from utils.logger import Logger
 from evaluation import evaluate
 from loss import LossForPretraining
-
-from colossalai.zero.init_ctx import ZeroInitContext
-from colossalai.zero.shard_utils import TensorShardStrategy
-from colossalai.zero.sharded_model import ShardedModelV2
-from colossalai.zero.sharded_optim import ShardedOptimizerV2
 from nvidia_bert_dataset_provider import NvidiaBertDatasetProvider
+from pretrain_utils import get_lr_scheduler, get_model, get_optimizer, save_ckpt
 from tqdm import tqdm
-import os
-import time
-from functools import partial
-
 from transformers import AutoTokenizer
+from utils.exp_util import get_mem_info, get_tflops, log_args, throughput_calculator
+from utils.global_vars import get_tensorboard_writer, get_timers, set_global_variables
+from utils.logger import Logger
 
-from colossalai.gemini import ChunkManager, GeminiManager
-from colossalai.utils.model.colo_init_context import ColoInitContext
-from colossalai.utils import get_current_device
+import colossalai
+import colossalai.nn as col_nn
+from colossalai.context import ParallelMode
+from colossalai.core import global_context as gpc
+from colossalai.nn.optimizer import HybridAdam
 from colossalai.nn.parallel import ZeroDDP
-from colossalai.zero import ZeroOptimizer
 from colossalai.tensor import ProcessGroup
-from colossalai.nn.optimizer import HybridAdam
+from colossalai.utils import get_current_device
+from colossalai.zero import ZeroOptimizer
+from colossalai.zero.gemini import ChunkManager, ColoInitContext, GeminiManager
+from colossalai.zero.legacy import ShardedModelV2, ShardedOptimizerV2, ZeroInitContext
+from colossalai.zero.legacy.shard_utils import TensorShardStrategy
 
 
 def main():
@@ -54,10 +50,12 @@ def main():
         args.local_rank = -1
         args.log_interval = 1
     else:
-        colossalai.launch_from_torch(args.colossal_config) #args.colossal_config
+        colossalai.launch_from_torch(args.colossal_config)    # args.colossal_config
         args.local_rank = int(os.environ["LOCAL_RANK"])
-        logger.info(f'launch_from_torch, world size: {torch.distributed.get_world_size()} | ' + 
-                    f'ParallelMode.MODEL: {ParallelMode.MODEL} | ParallelMode.DATA: {ParallelMode.DATA} | ParallelMode.TENSOR: {ParallelMode.TENSOR}')
+        logger.info(
+            f'launch_from_torch, world size: {torch.distributed.get_world_size()} | ' +
+            f'ParallelMode.MODEL: {ParallelMode.MODEL} | ParallelMode.DATA: {ParallelMode.DATA} | ParallelMode.TENSOR: {ParallelMode.TENSOR}'
+        )
 
     log_args(logger, args)
     args.tokenizer = tokenizer
@@ -84,7 +82,8 @@ def main():
     logger.info(f'Model numel: {numel}')
 
     get_tflops_func = partial(get_tflops, numel, args.train_micro_batch_size_per_gpu, args.max_seq_length)
-    steps_per_epoch = 144003367 // world_size // args.train_micro_batch_size_per_gpu // args.gradient_accumulation_steps // args.refresh_bucket_size #len(dataloader)
+    # len(dataloader)
+    steps_per_epoch = 144003367 // world_size // args.train_micro_batch_size_per_gpu // args.gradient_accumulation_steps // args.refresh_bucket_size
     total_steps = steps_per_epoch * args.epoch
 
     # build optimizer and lr_scheduler
@@ -98,7 +97,10 @@ def main():
         o_l_state_dict['lr_scheduler']['last_epoch'] = o_l_state_dict['lr_scheduler']['last_epoch'] - 1
         optimizer = get_optimizer(model, lr=args.lr)
         optimizer.load_state_dict(o_l_state_dict['optimizer'])
-        lr_scheduler = get_lr_scheduler(optimizer, total_steps=total_steps, last_epoch=o_l_state_dict['lr_scheduler']['last_epoch']) #o_l_state_dict['lr_scheduler']['last_epoch']
+        # o_l_state_dict['lr_scheduler']['last_epoch']
+        lr_scheduler = get_lr_scheduler(optimizer,
+                                        total_steps=total_steps,
+                                        last_epoch=o_l_state_dict['lr_scheduler']['last_epoch'])
         for state in optimizer.state.values():
             for k, v in state.items():
                 if isinstance(v, torch.Tensor):
@@ -109,7 +111,9 @@ def main():
         start_epoch = o_l_state_dict['epoch']
         start_shard = o_l_state_dict['shard'] + 1
         # global_step = o_l_state_dict['global_step'] + 1
-        logger.info(f'resume from epoch {start_epoch} shard {start_shard} step {lr_scheduler.last_epoch} lr {lr_scheduler.get_last_lr()[0]}')
+        logger.info(
+            f'resume from epoch {start_epoch} shard {start_shard} step {lr_scheduler.last_epoch} lr {lr_scheduler.get_last_lr()[0]}'
+        )
     else:
         optimizer = get_optimizer(model, lr=args.lr)
         lr_scheduler = get_lr_scheduler(optimizer, total_steps=total_steps, last_epoch=-1)
@@ -130,7 +134,6 @@ def main():
 
     logger.info(get_mem_info(prefix='After init model, '))
 
-
     best_loss = None
     eval_loss = 0
     train_loss = 0
@@ -146,7 +149,10 @@ def main():
             dataset_iterator, total_length = pretrain_dataset_provider.get_shard(shard)
             # pretrain_dataset_provider.prefetch_shard(shard + 1) # may cause cpu memory overload
             if torch.distributed.get_rank() == 0:
-                iterator_data = tqdm(enumerate(dataset_iterator), total=(total_length // args.train_micro_batch_size_per_gpu // world_size), colour='cyan', smoothing=1)
+                iterator_data = tqdm(enumerate(dataset_iterator),
+                                     total=(total_length // args.train_micro_batch_size_per_gpu // world_size),
+                                     colour='cyan',
+                                     smoothing=1)
             else:
                 iterator_data = enumerate(dataset_iterator)
 
@@ -179,7 +185,8 @@ def main():
                         and torch.distributed.get_rank() == 0:
                     elapsed_time = timers('interval_time').elapsed(reset=False)
                     elapsed_time_per_iteration = elapsed_time / global_step
-                    samples_per_sec, tflops, approx_parameters_in_billions = throughput_calculator(numel, args, config, elapsed_time, global_step, world_size)
+                    samples_per_sec, tflops, approx_parameters_in_billions = throughput_calculator(
+                        numel, args, config, elapsed_time, global_step, world_size)
 
                     cur_loss = train_loss / args.log_interval
                     current_lr = lr_scheduelr.get_last_lr()[0]
@@ -189,7 +196,8 @@ def main():
 
                     if args.wandb:
                         tensorboard_log = get_tensorboard_writer()
-                        tensorboard_log.log_train({
+                        tensorboard_log.log_train(
+                            {
                                 'lr': current_lr,
                                 'loss': cur_loss,
                                 'ppl': math.exp(cur_loss),
@@ -202,12 +210,14 @@ def main():
             logger.info('*' * 100)
 
             eval_loss += evaluate(engine, args, logger, global_step)
-            save_ckpt(engine.model, optimizer, lr_scheduelr, os.path.join(args.ckpt_path, launch_time, f'epoch-{epoch}_shard-{shard}_' + launch_time), epoch, shard, global_step)
-        
+            save_ckpt(engine.model, optimizer, lr_scheduelr,
+                      os.path.join(args.ckpt_path, launch_time, f'epoch-{epoch}_shard-{shard}_' + launch_time), epoch,
+                      shard, global_step)
 
         eval_loss /= len(os.listdir(args.data_path_prefix))
-        logger.info(f'epoch {epoch} | shard_length {len(os.listdir(args.data_path_prefix))} | elapsed_time: {timers("epoch_time").elapsed() / 60 :.3f} mins' + \
-                    f'eval_loss: {eval_loss} | ppl: {math.exp(eval_loss)}')
+        logger.info(
+            f'epoch {epoch} | shard_length {len(os.listdir(args.data_path_prefix))} | elapsed_time: {timers("epoch_time").elapsed() / 60 :.3f} mins'
+            + f'eval_loss: {eval_loss} | ppl: {math.exp(eval_loss)}')
         logger.info('-' * 100)
         if args.wandb and torch.distributed.get_rank() == 0:
             tensorboard_log = get_tensorboard_writer()

examples/tutorial/opt/opt/run_clm.py

@@ -30,24 +30,13 @@ from itertools import chain
 import datasets
 import torch
 import torch.distributed as dist
+import transformers
 from accelerate.utils import set_seed
 from context import barrier_context
 from datasets import load_dataset
 from packaging import version
 from torch.utils.data import DataLoader
 from tqdm.auto import tqdm
-
-import colossalai
-import transformers
-from colossalai.context import ParallelMode
-from colossalai.core import global_context as gpc
-from colossalai.logging import disable_existing_loggers, get_dist_logger
-from colossalai.nn.optimizer import HybridAdam
-from colossalai.nn.optimizer.zero_optimizer import ZeroOptimizer
-from colossalai.nn.parallel import ZeroDDP
-from colossalai.tensor import ProcessGroup
-from colossalai.utils import get_current_device, get_dataloader
-from colossalai.utils.model.colo_init_context import ColoInitContext
 from transformers import (
     CONFIG_MAPPING,
     MODEL_MAPPING,
@@ -61,6 +50,15 @@ from transformers import (
 )
 from transformers.utils.versions import require_version
 
+import colossalai
+from colossalai.context import ParallelMode
+from colossalai.core import global_context as gpc
+from colossalai.logging import disable_existing_loggers, get_dist_logger
+from colossalai.nn.optimizer import HybridAdam
+from colossalai.tensor import ProcessGroup
+from colossalai.utils import get_current_device, get_dataloader
+from colossalai.zero import ColoInitContext, ZeroDDP, ZeroOptimizer
+
 require_version("datasets>=1.8.0", "To fix: pip install -r examples/pytorch/language-modeling/requirements.txt")
 
 MODEL_CONFIG_CLASSES = list(MODEL_MAPPING.keys())

tests/test_auto_parallel/test_offload/test_perf.py

@@ -12,10 +12,9 @@ from colossalai.auto_parallel.offload.mem_optimize import memory_optimize
 from colossalai.auto_parallel.offload.solver import NOT_NVML
 from colossalai.fx.profiler import parameter_size
 from colossalai.nn.optimizer import HybridAdam
-from colossalai.nn.parallel import zero_model_wrapper, zero_optim_wrapper
 from colossalai.testing import parameterize
 from colossalai.utils import free_port, get_current_device
-from colossalai.utils.model.colo_init_context import ColoInitContext
+from colossalai.zero import ColoInitContext, zero_model_wrapper, zero_optim_wrapper
 from tests.test_auto_parallel.test_offload.model_utils import *
 from tests.test_tensor.common_utils import set_seed
 

tests/test_auto_parallel/test_tensor_shard/test_compatibility_with_gemini.py

@@ -11,12 +11,11 @@ from colossalai.device.device_mesh import DeviceMesh
 from colossalai.initialize import launch
 from colossalai.logging import disable_existing_loggers
 from colossalai.nn.optimizer import HybridAdam
-from colossalai.nn.parallel import zero_model_wrapper, zero_optim_wrapper
 from colossalai.tensor.process_group import ProcessGroup
 from colossalai.testing import assert_close, rerun_if_address_is_in_use
 from colossalai.testing.pytest_wrapper import run_on_environment_flag
 from colossalai.utils import free_port, get_current_device
-from colossalai.utils.model.colo_init_context import ColoInitContext, post_process_colo_init_ctx
+from colossalai.zero import ColoInitContext, post_process_colo_init_ctx, zero_model_wrapper, zero_optim_wrapper
 
 
 class MLP(torch.nn.Module):

tests/test_ddp/test_ddp_ignore_params.py

@@ -10,14 +10,14 @@ import torch.distributed as dist
 import torch.multiprocessing as mp
 
 import colossalai
-from colossalai.gemini.chunk import ChunkManager, search_chunk_configuration
-from colossalai.gemini.gemini_mgr import GeminiManager
-from colossalai.nn.parallel import ColoDDP, ZeroDDP
+from colossalai.nn.parallel import ColoDDP
 from colossalai.tensor import ProcessGroup
 from colossalai.testing import rerun_if_address_is_in_use
 from colossalai.utils import free_port
 from colossalai.utils.cuda import get_current_device
-from colossalai.utils.model.colo_init_context import ColoInitContext
+from colossalai.zero import ColoInitContext, ZeroDDP
+from colossalai.zero.gemini.chunk import ChunkManager, search_chunk_configuration
+from colossalai.zero.gemini.gemini_mgr import GeminiManager
 
 
 def set_seed(seed):

tests/test_ddp/test_ddp_state_dict.py

 import copy
+from collections import OrderedDict
+from functools import partial
 
 import pytest
-import colossalai
 import torch
 import torch.multiprocessing as mp
+
+import colossalai
+from colossalai.nn.parallel import ColoDDP
+from colossalai.tensor import ColoParameter, ProcessGroup
 from colossalai.testing import rerun_if_address_is_in_use
-from colossalai.utils.cuda import get_current_device
 from colossalai.utils import free_port
-from colossalai.utils.model.colo_init_context import ColoInitContext
-from functools import partial
+from colossalai.utils.cuda import get_current_device
+from colossalai.zero import ColoInitContext
 from tests.components_to_test.registry import non_distributed_component_funcs
-from colossalai.nn.parallel import ColoDDP
-from collections import OrderedDict
-from colossalai.tensor import ProcessGroup, ColoParameter
 
 
 def check_state_dict_equal(state_dict: OrderedDict, other_state_dict: OrderedDict):

tests/test_gemini/test_gemini_manager.py

 import pytest
 import torch
 
-from colossalai.gemini.stateful_tensor import TensorState, StatefulTensor
+from colossalai.zero.legacy.gemini.stateful_tensor import StatefulTensor, TensorState
 
 
 @pytest.mark.dist

tests/test_gemini/test_param_op.py

@@ -2,7 +2,7 @@ import copy
 
 import torch
 
-from colossalai.gemini.paramhooks import BaseParamHookMgr
+from colossalai.zero.legacy.gemini.paramhooks import BaseParamHookMgr
 from tests.components_to_test.registry import non_distributed_component_funcs