[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

tests/test_gemini/test_runtime_mem_tracer.py

@@ -3,8 +3,8 @@ from copy import deepcopy
 import numpy as np
 import torch
 
-from colossalai.gemini.memory_tracer.runtime_mem_tracer import RuntimeMemTracer
-from colossalai.utils.model.colo_init_context import ColoInitContext
+from colossalai.zero import ColoInitContext
+from colossalai.zero.gemini.memory_tracer.runtime_mem_tracer import RuntimeMemTracer
 from tests.components_to_test import run_fwd_bwd
 from tests.components_to_test.registry import non_distributed_component_funcs
 

tests/test_gemini/update/test_chunk_mgrv2.py

@@ -5,10 +5,10 @@ import torch
 import torch.multiprocessing as mp
 
 import colossalai
-from colossalai.gemini.chunk import ChunkManager
 from colossalai.tensor import ColoTensor, ColoTensorSpec, ProcessGroup
 from colossalai.testing import parameterize, rerun_if_address_is_in_use
 from colossalai.utils import free_port
+from colossalai.zero.gemini.chunk import ChunkManager
 from tests.test_tensor.common_utils import debug_print
 
 CUDA_MEM_0 = {False: 512, True: 1024}

tests/test_gemini/update/test_chunkv2.py

@@ -6,12 +6,12 @@ import torch.distributed as dist
 import torch.multiprocessing as mp
 
 import colossalai
-from colossalai.gemini import TensorState
-from colossalai.gemini.chunk import Chunk
 from colossalai.tensor import ColoParameter
 from colossalai.tensor import ProcessGroup as ColoProcessGroup
 from colossalai.testing import parameterize, rerun_if_address_is_in_use
 from colossalai.utils import free_port, get_current_device
+from colossalai.zero.gemini import TensorState
+from colossalai.zero.gemini.chunk import Chunk
 
 
 def dist_sum(x):

tests/test_gemini/update/test_fwd_bwd.py

@@ -8,16 +8,14 @@ from torch.testing import assert_close
 
 import colossalai
 from colossalai.amp import convert_to_apex_amp
-from colossalai.gemini.chunk import ChunkManager, search_chunk_configuration
-from colossalai.gemini.gemini_mgr import GeminiManager
 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.testing import parameterize, 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, ZeroOptimizer
+from colossalai.zero.gemini.chunk import ChunkManager, search_chunk_configuration
+from colossalai.zero.gemini.gemini_mgr import GeminiManager
 from tests.components_to_test import run_fwd_bwd
 from tests.components_to_test.registry import non_distributed_component_funcs
 from tests.test_tensor.common_utils import set_seed

tests/test_gemini/update/test_gemini_use_rmt.py

@@ -5,15 +5,13 @@ import torch
 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.gemini.memory_tracer.runtime_mem_tracer import RuntimeMemTracer
-from colossalai.nn.optimizer.gemini_optimizer import GeminiAdamOptimizer
-from colossalai.nn.parallel import GeminiDDP, ZeroDDP
 from colossalai.tensor import ProcessGroup
 from colossalai.testing import parameterize, rerun_if_address_is_in_use
 from colossalai.utils import free_port
-from colossalai.utils.model.colo_init_context import ColoInitContext
+from colossalai.zero import ColoInitContext, GeminiAdamOptimizer, GeminiDDP, ZeroDDP
+from colossalai.zero.gemini.chunk import ChunkManager, search_chunk_configuration
+from colossalai.zero.gemini.gemini_mgr import GeminiManager
+from colossalai.zero.gemini.memory_tracer.runtime_mem_tracer import RuntimeMemTracer
 from tests.components_to_test import run_fwd_bwd
 from tests.components_to_test.registry import non_distributed_component_funcs
 from tests.test_tensor.common_utils import set_seed

tests/test_gemini/update/test_get_torch_model.py

@@ -6,13 +6,12 @@ import torch
 import torch.multiprocessing as mp
 
 import colossalai
-from colossalai.nn.parallel import GeminiDDP
-from colossalai.nn.parallel.utils import get_static_torch_model
 from colossalai.tensor import ColoParameter
 from colossalai.testing import parameterize, 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, GeminiDDP
+from colossalai.zero.gemini.utils import get_static_torch_model
 from tests.components_to_test.registry import non_distributed_component_funcs
 
 

tests/test_gemini/update/test_grad_clip.py

@@ -10,15 +10,13 @@ from torch.testing import assert_close
 
 import colossalai
 from colossalai.amp import convert_to_apex_amp
-from colossalai.gemini.chunk import ChunkManager, search_chunk_configuration
-from colossalai.gemini.gemini_mgr import GeminiManager
 from colossalai.nn.optimizer import HybridAdam
-from colossalai.nn.optimizer.zero_optimizer import ZeroOptimizer
-from colossalai.nn.parallel import ZeroDDP
 from colossalai.testing import parameterize, 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, ZeroOptimizer
+from colossalai.zero.gemini.chunk import ChunkManager, search_chunk_configuration
+from colossalai.zero.gemini.gemini_mgr import GeminiManager
 from tests.components_to_test import run_fwd_bwd
 from tests.components_to_test.registry import non_distributed_component_funcs
 from tests.test_tensor.common_utils import debug_print, set_seed

tests/test_gemini/update/test_inference.py

@@ -10,15 +10,13 @@ from torch.testing import assert_close
 
 import colossalai
 from colossalai.amp import convert_to_apex_amp
-from colossalai.gemini.chunk import ChunkManager, init_chunk_manager, search_chunk_configuration
-from colossalai.gemini.gemini_mgr import GeminiManager
 from colossalai.nn.optimizer import HybridAdam
-from colossalai.nn.optimizer.zero_optimizer import ZeroOptimizer
-from colossalai.nn.parallel import ZeroDDP, zero_model_wrapper
 from colossalai.testing import parameterize, 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, post_process_colo_init_ctx
+from colossalai.zero import ColoInitContext, ZeroDDP, ZeroOptimizer, post_process_colo_init_ctx, zero_model_wrapper
+from colossalai.zero.gemini.chunk import ChunkManager, init_chunk_manager, search_chunk_configuration
+from colossalai.zero.gemini.gemini_mgr import GeminiManager
 from tests.components_to_test import run_fwd_bwd
 from tests.components_to_test.registry import non_distributed_component_funcs
 from tests.test_tensor.common_utils import debug_print, set_seed

tests/test_gemini/update/test_optim.py

@@ -9,16 +9,14 @@ from torch.testing import assert_close
 
 import colossalai
 from colossalai.amp import convert_to_apex_amp
-from colossalai.gemini.chunk import ChunkManager, init_chunk_manager, search_chunk_configuration
-from colossalai.gemini.gemini_mgr import GeminiManager
 from colossalai.nn.optimizer import HybridAdam
-from colossalai.nn.optimizer.zero_optimizer import ZeroOptimizer
-from colossalai.nn.parallel import ZeroDDP
 from colossalai.tensor import ColoParameter, ColoTensor
 from colossalai.testing import parameterize, 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, post_process_colo_init_ctx
+from colossalai.zero import ColoInitContext, ZeroDDP, ZeroOptimizer, post_process_colo_init_ctx
+from colossalai.zero.gemini.chunk import ChunkManager, init_chunk_manager, search_chunk_configuration
+from colossalai.zero.gemini.gemini_mgr import GeminiManager
 from tests.components_to_test import run_fwd_bwd
 from tests.components_to_test.registry import non_distributed_component_funcs
 from tests.test_tensor.common_utils import debug_print, set_seed

tests/test_gemini/update/test_search.py

@@ -6,11 +6,11 @@ import torch.distributed as dist
 import torch.multiprocessing as mp
 
 import colossalai
-from colossalai.gemini.chunk import init_chunk_manager, search_chunk_configuration
 from colossalai.tensor import ComputePattern, ComputeSpec, ProcessGroup, ShardSpec
 from colossalai.testing import rerun_if_address_is_in_use
 from colossalai.utils import free_port, get_current_device
-from colossalai.utils.model.colo_init_context import ColoInitContext
+from colossalai.zero import ColoInitContext
+from colossalai.zero.gemini.chunk import init_chunk_manager, search_chunk_configuration
 from tests.components_to_test.registry import non_distributed_component_funcs
 
 

tests/test_gemini/update/test_zeroddp_state_dict.py

@@ -7,13 +7,12 @@ import torch.multiprocessing as mp
 from torch.testing import assert_close
 
 import colossalai
-from colossalai.gemini.chunk import ChunkManager, search_chunk_configuration
-from colossalai.gemini.gemini_mgr import GeminiManager
-from colossalai.nn.parallel import ZeroDDP
 from colossalai.testing import parameterize, 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
 from tests.components_to_test.registry import non_distributed_component_funcs
 from tests.test_tensor.common_utils import debug_print, set_seed
 

tests/test_gemini/update/test_zerooptim_state_dict.py

@@ -6,15 +6,13 @@ 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.optimizer import HybridAdam
-from colossalai.nn.optimizer.zero_optimizer import ZeroOptimizer
-from colossalai.nn.parallel import ZeroDDP
 from colossalai.testing import parameterize, 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, ZeroOptimizer
+from colossalai.zero.gemini.chunk import ChunkManager, search_chunk_configuration
+from colossalai.zero.gemini.gemini_mgr import GeminiManager
 from tests.components_to_test.registry import non_distributed_component_funcs
 from tests.test_tensor.common_utils import debug_print, set_seed
 

tests/test_moe/test_moe_checkpoint.py

@@ -11,7 +11,7 @@ from colossalai.context import MOE_CONTEXT
 from colossalai.nn.layer.moe import load_moe_model, save_moe_model
 from colossalai.testing import parameterize, rerun_if_address_is_in_use
 from colossalai.utils import free_port, get_current_device
-from colossalai.utils.model.colo_init_context import ColoInitContext
+from colossalai.zero import ColoInitContext
 from tests.test_moe.test_moe_zero_init import MoeModel
 from tests.test_tensor.common_utils import debug_print
 from tests.test_zero.common import CONFIG

tests/test_moe/test_moe_colo_init.py

 from functools import partial
 
-import colossalai
 import pytest
 import torch
-import torch.multiprocessing as mp
 import torch.distributed as dist
-from colossalai.testing import parameterize
-from colossalai.utils import free_port
+import torch.multiprocessing as mp
+
+import colossalai
 from colossalai.context import MOE_CONTEXT
 from colossalai.tensor import ColoParameter
-from colossalai.utils.model.colo_init_context import ColoInitContext
-
-from colossalai.testing import rerun_if_address_is_in_use
-from colossalai.utils import get_current_device
-
-from tests.test_zero.common import CONFIG
+from colossalai.testing import parameterize, rerun_if_address_is_in_use
+from colossalai.utils import free_port, get_current_device
+from colossalai.zero import ColoInitContext
 from tests.test_moe.test_moe_zero_init import MoeModel
 from tests.test_tensor.common_utils import debug_print
+from tests.test_zero.common import CONFIG
 
 
 @parameterize("init_device_type", ['cpu', 'cuda'])

tests/test_moe/test_moe_zero_init.py

 from functools import partial
 
-import colossalai
 import pytest
 import torch
 import torch.multiprocessing as mp
 import torch.nn as nn
-from colossalai.nn import CheckpointModule
-from colossalai.logging import get_dist_logger
-from colossalai.testing import parameterize
-from colossalai.utils import free_port
+
+import colossalai
 from colossalai.context import MOE_CONTEXT
+from colossalai.logging import get_dist_logger
+from colossalai.nn import CheckpointModule
 from colossalai.nn.layer import MoeModule
-from colossalai.zero.init_ctx import ZeroInitContext
-from colossalai.zero.shard_utils import (BucketTensorShardStrategy, TensorShardStrategy)
-
-from colossalai.testing import rerun_if_address_is_in_use
-from colossalai.utils import get_current_device
+from colossalai.testing import parameterize, rerun_if_address_is_in_use
+from colossalai.utils import free_port, get_current_device
+from colossalai.zero.legacy.init_ctx import ZeroInitContext
+from colossalai.zero.legacy.shard_utils import BucketTensorShardStrategy, TensorShardStrategy
 from tests.test_zero.common import CONFIG
 
 

tests/test_moe/test_moe_zero_model.py

@@ -10,11 +10,11 @@ from colossalai.engine.gradient_handler import MoeGradientHandler
 from colossalai.nn import MoeLoss
 from colossalai.testing import assert_equal_in_group, parameterize, rerun_if_address_is_in_use
 from colossalai.utils import free_port
-from colossalai.zero.init_ctx import ZeroInitContext
-from colossalai.zero.shard_utils import BucketTensorShardStrategy, TensorShardStrategy
-from colossalai.zero.sharded_model import ShardedModelV2
-from colossalai.zero.sharded_model._utils import cast_tensor_to_fp16
-from colossalai.zero.sharded_model.utils import col_model_deepcopy
+from colossalai.zero.legacy.init_ctx import ZeroInitContext
+from colossalai.zero.legacy.shard_utils import BucketTensorShardStrategy, TensorShardStrategy
+from colossalai.zero.legacy.sharded_model import ShardedModelV2
+from colossalai.zero.legacy.sharded_model._utils import cast_tensor_to_fp16
+from colossalai.zero.legacy.sharded_model.utils import col_model_deepcopy
 from tests.components_to_test.registry import non_distributed_component_funcs
 from tests.test_moe.test_moe_zero_init import MoeModel
 from tests.test_zero.common import CONFIG, check_grads_padding, run_fwd_bwd

tests/test_moe/test_moe_zero_optim.py

@@ -12,12 +12,12 @@ from colossalai.nn import MoeLoss
 from colossalai.nn.optimizer import CPUAdam
 from colossalai.testing import assert_equal_in_group, parameterize, rerun_if_address_is_in_use
 from colossalai.utils import free_port, get_current_device
-from colossalai.zero.init_ctx import ZeroInitContext
-from colossalai.zero.shard_utils import BucketTensorShardStrategy, TensorShardStrategy
-from colossalai.zero.sharded_model import ShardedModelV2
-from colossalai.zero.sharded_model.utils import col_model_deepcopy
-from colossalai.zero.sharded_optim import ShardedOptimizerV2
-from colossalai.zero.sharded_optim._utils import has_inf_or_nan
+from colossalai.zero.legacy.init_ctx import ZeroInitContext
+from colossalai.zero.legacy.shard_utils import BucketTensorShardStrategy, TensorShardStrategy
+from colossalai.zero.legacy.sharded_model import ShardedModelV2
+from colossalai.zero.legacy.sharded_model.utils import col_model_deepcopy
+from colossalai.zero.legacy.sharded_optim import ShardedOptimizerV2
+from colossalai.zero.low_level._utils import has_inf_or_nan
 from tests.components_to_test.registry import non_distributed_component_funcs
 from tests.test_moe.test_moe_zero_init import MoeModel
 from tests.test_zero.common import CONFIG, check_sharded_model_params

tests/test_optimizer/test_cpu_adam.py

@@ -56,7 +56,7 @@ def test_cpu_adam(adamw, step, p_dtype, g_dtype):
     eps = 1e-8
     weight_decay = 0
 
-    for i in range(1024):
+    for i in range(3):
         p_data = torch.rand(64, dtype=p_dtype)
         p_data_copy = p_data.clone().float()
         p_grad = torch.rand(64, dtype=g_dtype)

tests/test_optimizer/test_fused_adam_kernel.py

@@ -54,7 +54,7 @@ def test_adam(adamw, step, p_dtype, g_dtype):
 
     count = 0
 
-    for i in range(1024):
+    for i in range(3):
         p = torch.rand(64, dtype=p_dtype).cuda()
         p_copy = p.clone().float()
         g = torch.rand(p.shape, dtype=g_dtype).cuda()

tests/test_optimizer/test_hybrid_adam.py

 import torch
 import torch.nn as nn
-from torch.optim.adam import Adam
 from torch.optim import AdamW
+from torch.optim.adam import Adam
 
 from colossalai.nn.optimizer.hybrid_adam import HybridAdam
 from colossalai.testing import parameterize
 
-RE = 1024
+RE = 3
 
 
 @parameterize('adamw', [False, True])