[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_tensor/model/test_gpt2.py

@@ -11,7 +11,7 @@ from colossalai.tensor import ColoTensor, ColoTensorSpec, ComputePattern, Comput
 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
 from tests.components_to_test.registry import non_distributed_component_funcs
 from tests.test_tensor.common_utils import (
     debug_print,

tests/test_tensor/model/test_model.py

@@ -11,7 +11,7 @@ from colossalai.tensor.colo_parameter import ColoParameter
 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
 from tests.components_to_test.registry import non_distributed_component_funcs
 from tests.test_tensor.common_utils import (
     check_equal,

tests/test_tensor/model/test_module_spec.py

@@ -20,7 +20,7 @@ from colossalai.tensor import (
 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
 from tests.components_to_test.registry import non_distributed_component_funcs
 from tests.test_tensor.common_utils import set_seed, tensor_equal, tensor_shard_equal
 

tests/test_tensor/test_context.py

@@ -17,7 +17,7 @@ from colossalai.tensor import (
 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
 from tests.components_to_test.registry import non_distributed_component_funcs
 from tests.test_tensor.common_utils import set_seed
 

tests/test_tensor/test_tp_with_zero.py

@@ -7,14 +7,12 @@ from torch.nn.parallel import DistributedDataParallel as DDP
 
 import colossalai
 from colossalai.amp import convert_to_apex_amp
-from colossalai.gemini.chunk import search_chunk_configuration
-from colossalai.nn.optimizer.gemini_optimizer import GeminiAdamOptimizer
-from colossalai.nn.parallel import GeminiDDP, ZeroDDP
 from colossalai.tensor import ColoTensor, ColoTensorSpec, ComputePattern, ComputeSpec, ProcessGroup, ShardSpec
 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, GeminiAdamOptimizer, GeminiDDP, ZeroDDP
+from colossalai.zero.gemini import search_chunk_configuration
 from tests.components_to_test.registry import non_distributed_component_funcs
 from tests.test_tensor.common_utils import set_seed, tensor_shard_equal
 from tests.test_tensor.model.test_gpt2 import init_megatron_spec

tests/test_utils/test_colo_checkpoint.py

-import os, shutil
-import torch
-import pytest
+import os
+import shutil
 from copy import deepcopy
 from functools import partial
 
-import torch.multiprocessing as mp
+import pytest
+import torch
 import torch.distributed as dist
-
-from torch.optim.lr_scheduler import CosineAnnealingLR
-from torch.optim.lr_scheduler import MultiplicativeLR
-from colossalai.nn.lr_scheduler import CosineAnnealingWarmupLR
+import torch.multiprocessing as mp
+from torch.optim.lr_scheduler import CosineAnnealingLR, MultiplicativeLR
 
 import colossalai
+from colossalai.nn.lr_scheduler import CosineAnnealingWarmupLR
+from colossalai.nn.optimizer import ColossalaiOptimizer
+from colossalai.tensor import ColoTensor, ComputePattern, ComputeSpec, ProcessGroup, ShardSpec
 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 colossalai.tensor import ComputePattern, ComputeSpec, ColoTensor, ShardSpec, ProcessGroup
-from colossalai.utils.checkpoint import save_checkpoint, load_checkpoint
-from colossalai.nn.optimizer import ColossalaiOptimizer
-
+from colossalai.utils.checkpoint import load_checkpoint, save_checkpoint
+from colossalai.utils.cuda import get_current_device
+from colossalai.zero import ColoInitContext
 from tests.components_to_test.registry import non_distributed_component_funcs
 
 

tests/test_utils/test_commons.py

-from colossalai.utils import free_port
-from colossalai.testing import rerun_if_address_is_in_use
-from colossalai.zero.sharded_param import ShardedTensor
-from colossalai.gemini.tensor_utils import colo_model_data_tensor_move, colo_model_data_tensor_move_inline
-import colossalai
-
 import torch
-
 import torch.multiprocessing as mp
 
+import colossalai
+from colossalai.testing import rerun_if_address_is_in_use
+from colossalai.utils import free_port
+from colossalai.zero.legacy.gemini.tensor_utils import colo_model_data_tensor_move, colo_model_data_tensor_move_inline
+from colossalai.zero.legacy.sharded_param import ShardedTensor
+
 
 def run_tensor_move(rank):
     colossalai.launch(config={}, rank=0, world_size=1, host='localhost', port=free_port(), backend='nccl')

tests/test_utils/test_zero_gradient_clippling.py

@@ -2,21 +2,22 @@
 # -*- encoding: utf-8 -*-
 
 import copy
+from functools import partial
 
-import colossalai
-from colossalai.zero.sharded_model.sharded_model_v2 import ShardedModelV2
 import pytest
 import torch
 import torch.distributed as dist
 import torch.multiprocessing as mp
 import torch.nn as nn
-from colossalai.logging import disable_existing_loggers
-from colossalai.utils import checkpoint, clip_grad_norm_fp32, free_port
 from torch.nn.parallel import DistributedDataParallel as DDP
 from torch.nn.utils import clip_grad_norm_
-from colossalai.zero.shard_utils.tensor_shard_strategy import TensorShardStrategy
-from functools import partial
+
+import colossalai
+from colossalai.logging import disable_existing_loggers
 from colossalai.testing import parameterize, rerun_if_address_is_in_use
+from colossalai.utils import checkpoint, clip_grad_norm_fp32, free_port
+from colossalai.zero.legacy.shard_utils.tensor_shard_strategy import TensorShardStrategy
+from colossalai.zero.legacy.sharded_model.sharded_model_v2 import ShardedModelV2
 
 
 def checkpoint_wrapper(module, enable=True):

tests/test_zero/common.py

@@ -2,10 +2,11 @@ from functools import partial
 
 import torch
 import torch.distributed as dist
+
 from colossalai.logging import get_dist_logger
 from colossalai.utils import checkpoint
-from colossalai.zero.shard_utils import TensorShardStrategy
-from colossalai.zero.sharded_model import ShardedModelV2
+from colossalai.zero.legacy.shard_utils import TensorShardStrategy
+from colossalai.zero.legacy.sharded_model import ShardedModelV2
 
 LOGGER = get_dist_logger('zero_test')
 

tests/test_zero/low_level_zero/test_zero_init.py

@@ -9,8 +9,7 @@ import torch.nn as nn
 import colossalai
 from colossalai.tensor import ProcessGroup
 from colossalai.utils import free_port, get_current_device
-from colossalai.utils.model.colo_init_context import ColoInitContext
-from colossalai.zero import LowLevelZeroOptimizer
+from colossalai.zero import ColoInitContext, LowLevelZeroOptimizer
 
 
 class MlpModel(nn.Module):

tests/test_zero/low_level_zero/test_zero_tp.py

@@ -11,8 +11,7 @@ import colossalai
 from colossalai.tensor import ProcessGroup
 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 LowLevelZeroOptimizer
+from colossalai.zero import ColoInitContext, LowLevelZeroOptimizer
 from tests.test_tensor.common_utils import set_seed, split_param_col_tp1d, split_param_row_tp1d, tensor_shard_equal
 
 

tests/test_zero/test_found_inf.py

-from functools import partial
-
-import colossalai
-from colossalai.utils.cuda import get_current_device
-import pytest
-import torch
-import torch.multiprocessing as mp
-from colossalai.nn.optimizer import HybridAdam
-from colossalai.testing import 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
-from colossalai.zero.sharded_model import ShardedModelV2
-from colossalai.zero.sharded_optim import ShardedOptimizerV2
-from colossalai.zero.sharded_optim._utils import has_inf_or_nan
-from tests.components_to_test.registry import non_distributed_component_funcs
-from tests.test_zero.test_sharded_optim_v2 import _run_step
-
-from common import CONFIG
-
-
-@parameterize("cpu_offload", [True, False])
-@parameterize("shard_strategy_class", [BucketTensorShardStrategy])
-@parameterize("gpu_margin_mem_ratio", [0.0, 0.7])
-def _run_test_found_inf(cpu_offload, shard_strategy_class, gpu_margin_mem_ratio):
-    test_models = ['repeated_computed_layers']
-    shard_strategy = shard_strategy_class()
-
-    for model_name in test_models:
-        get_components_func = non_distributed_component_funcs.get_callable(model_name)
-        model_builder, train_dataloader, _, optimizer_class, criterion = get_components_func()
-
-        with ZeroInitContext(target_device=torch.device(f'cpu:0') if cpu_offload else get_current_device(),
-                             shard_strategy=shard_strategy,
-                             shard_param=True):
-            zero_model = model_builder(checkpoint=True)
-        zero_model = ShardedModelV2(
-            zero_model,
-            shard_strategy,
-            tensor_placement_policy='cpu' if cpu_offload else 'cuda',
-            reuse_fp16_shard=True,
-        )
-
-        sharded_optim = HybridAdam(zero_model.parameters(), lr=1e-3)
-        sharded_optim = ShardedOptimizerV2(zero_model, sharded_optim, gpu_margin_mem_ratio=gpu_margin_mem_ratio)
-
-        for i, (data, label) in enumerate(train_dataloader):
-            if i > 1:
-                break
-            assert zero_model.overflow_counter == 0
-            data, label = data.cuda(), label.cuda()
-            _run_step(zero_model, sharded_optim, data, label, criterion, False)
-            for param in zero_model.parameters():
-                assert not has_inf_or_nan(param.colo_attr.data_payload)
-
-
-def _run_dist(rank, world_size, port):
-    colossalai.launch(config=CONFIG, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl')
-    _run_test_found_inf()
-
-
-# use_cpuadam = True can be used with cpu_offload = False
-@pytest.mark.dist
-@pytest.mark.parametrize("world_size", [1, 2])
-@rerun_if_address_is_in_use()
-def test_found_inf(world_size):
-    run_func = partial(_run_dist, world_size=world_size, port=free_port())
-    mp.spawn(run_func, nprocs=world_size)
-
-
-if __name__ == '__main__':
-    test_found_inf(world_size=2)
+from functools import partial
+
+import pytest
+import torch
+import torch.multiprocessing as mp
+from common import CONFIG
+
+import colossalai
+from colossalai.nn.optimizer import HybridAdam
+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.zero.legacy.init_ctx import ZeroInitContext
+from colossalai.zero.legacy.shard_utils import BucketTensorShardStrategy
+from colossalai.zero.legacy.sharded_model import ShardedModelV2
+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_zero.test_sharded_optim_v2 import _run_step
+
+
+@parameterize("cpu_offload", [True, False])
+@parameterize("shard_strategy_class", [BucketTensorShardStrategy])
+@parameterize("gpu_margin_mem_ratio", [0.0, 0.7])
+def _run_test_found_inf(cpu_offload, shard_strategy_class, gpu_margin_mem_ratio):
+    test_models = ['repeated_computed_layers']
+    shard_strategy = shard_strategy_class()
+
+    for model_name in test_models:
+        get_components_func = non_distributed_component_funcs.get_callable(model_name)
+        model_builder, train_dataloader, _, optimizer_class, criterion = get_components_func()
+
+        with ZeroInitContext(target_device=torch.device(f'cpu:0') if cpu_offload else get_current_device(),
+                             shard_strategy=shard_strategy,
+                             shard_param=True):
+            zero_model = model_builder(checkpoint=True)
+        zero_model = ShardedModelV2(
+            zero_model,
+            shard_strategy,
+            tensor_placement_policy='cpu' if cpu_offload else 'cuda',
+            reuse_fp16_shard=True,
+        )
+
+        sharded_optim = HybridAdam(zero_model.parameters(), lr=1e-3)
+        sharded_optim = ShardedOptimizerV2(zero_model, sharded_optim, gpu_margin_mem_ratio=gpu_margin_mem_ratio)
+
+        for i, (data, label) in enumerate(train_dataloader):
+            if i > 1:
+                break
+            assert zero_model.overflow_counter == 0
+            data, label = data.cuda(), label.cuda()
+            _run_step(zero_model, sharded_optim, data, label, criterion, False)
+            for param in zero_model.parameters():
+                assert not has_inf_or_nan(param.colo_attr.data_payload)
+
+
+def _run_dist(rank, world_size, port):
+    colossalai.launch(config=CONFIG, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl')
+    _run_test_found_inf()
+
+
+# use_cpuadam = True can be used with cpu_offload = False
+@pytest.mark.dist
+@pytest.mark.parametrize("world_size", [1, 2])
+@rerun_if_address_is_in_use()
+def test_found_inf(world_size):
+    run_func = partial(_run_dist, world_size=world_size, port=free_port())
+    mp.spawn(run_func, nprocs=world_size)
+
+
+if __name__ == '__main__':
+    test_found_inf(world_size=2)

tests/test_zero/test_init_context.py

@@ -9,14 +9,14 @@ import torch.multiprocessing as mp
 from common import CONFIG
 
 import colossalai
-from colossalai.gemini.memory_tracer.utils import colo_model_mem_usage
 from colossalai.logging import get_dist_logger
 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.memory import colo_device_memory_used
-from colossalai.zero.init_ctx import ZeroInitContext
-from colossalai.zero.shard_utils import BucketTensorShardStrategy, TensorShardStrategy
+from colossalai.zero.gemini.memory_tracer.utils import colo_model_mem_usage
+from colossalai.zero.legacy.init_ctx import ZeroInitContext
+from colossalai.zero.legacy.shard_utils import BucketTensorShardStrategy, TensorShardStrategy
 from tests.components_to_test.registry import non_distributed_component_funcs
 
 

tests/test_zero/test_shard_model_v2.py

@@ -12,11 +12,11 @@ from torch.nn.parallel import DistributedDataParallel as DDP
 import colossalai
 from colossalai.testing import 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
-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
+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
 
 

tests/test_zero/test_shard_param.py

 from copy import deepcopy
 from functools import partial
 
-import colossalai
 import pytest
 import torch
 import torch.multiprocessing as mp
+
+import colossalai
 from colossalai.testing import parameterize, rerun_if_address_is_in_use
 from colossalai.utils import free_port
-from colossalai.zero.shard_utils import (BucketTensorShardStrategy, TensorShardStrategy)
-from colossalai.zero.sharded_param import ShardedTensor
-from colossalai.zero.sharded_param.sharded_param import ShardedParamV2
+from colossalai.zero.legacy.gemini.stateful_tensor import StatefulTensor
+from colossalai.zero.legacy.shard_utils import BucketTensorShardStrategy, TensorShardStrategy
+from colossalai.zero.legacy.sharded_param import ShardedTensor
+from colossalai.zero.legacy.sharded_param.sharded_param import ShardedParamV2
 from tests.test_zero.common import CONFIG, allclose
-from colossalai.gemini.stateful_tensor import StatefulTensor
 
 
 @parameterize("shard_strategy_class", [TensorShardStrategy, BucketTensorShardStrategy])

tests/test_zero/test_sharded_optim_state_dict.py

+from functools import partial
+
 import pytest
-import colossalai
 import torch
 import torch.multiprocessing as mp
-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 functools import partial
-from tests.test_tensor.common_utils import set_seed
-from tests.components_to_test.registry import non_distributed_component_funcs
-from colossalai.testing import parameterize
+
+import colossalai
 from colossalai.nn.optimizer import HybridAdam
-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 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.zero.legacy.init_ctx import ZeroInitContext
+from colossalai.zero.legacy.shard_utils import TensorShardStrategy
+from colossalai.zero.legacy.sharded_model import ShardedModelV2
+from colossalai.zero.legacy.sharded_optim import ShardedOptimizerV2
+from tests.components_to_test.registry import non_distributed_component_funcs
+from tests.test_tensor.common_utils import set_seed
 
 
 def init_zero(model_builder, placement_policy):

tests/test_zero/test_sharded_optim_v2.py

@@ -13,12 +13,12 @@ from colossalai.nn.optimizer import CPUAdam
 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.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
 
 

tests/test_zero/test_sharded_optim_with_sync_bn.py

@@ -3,18 +3,19 @@
 
 from functools import partial
 
-import colossalai
 import pytest
 import torch
 import torch.distributed as dist
 import torch.multiprocessing as mp
+from torchvision.models import resnet50
+
+import colossalai
 from colossalai.context.parallel_mode import ParallelMode
 from colossalai.core import global_context as gpc
 from colossalai.testing import 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 TensorShardStrategy
-from torchvision.models import resnet50
+from colossalai.zero.legacy.init_ctx import ZeroInitContext
+from colossalai.zero.legacy.shard_utils import TensorShardStrategy
 
 
 def run_dist(rank, world_size, port):

tests/test_zero/test_state_dict.py

@@ -4,20 +4,20 @@
 from copy import deepcopy
 from functools import partial
 
-import colossalai
 import pytest
 import torch
 import torch.multiprocessing as mp
+from common import CONFIG
+
+import colossalai
 from colossalai.testing import 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 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 col_model_deepcopy
 from tests.components_to_test.registry import non_distributed_component_funcs
 
-from common import CONFIG
-
 
 @parameterize("shard_strategy_class", [TensorShardStrategy, BucketTensorShardStrategy])
 def run_zero_state_dict(shard_strategy_class):

tests/test_zero/test_tensor_utils.py

+from functools import partial
+
 import pytest
+import torch
+import torch.multiprocessing as mp
 
 import colossalai
-from colossalai.utils.cuda import get_current_device
-from colossalai.gemini.tensor_utils import (colo_tensor_mem_usage, colo_model_data_tensor_move,
-                                            colo_model_data_tensor_move_inline, colo_model_data_move_to_cpu,
-                                            colo_model_tensor_clone)
-from colossalai.gemini.stateful_tensor import StatefulTensor
-from colossalai.utils import free_port
 from colossalai.testing import rerun_if_address_is_in_use
-
-import torch
-
-from functools import partial
-import torch.multiprocessing as mp
+from colossalai.utils import free_port
+from colossalai.utils.cuda import get_current_device
+from colossalai.zero.legacy.gemini.stateful_tensor import StatefulTensor
+from colossalai.zero.legacy.gemini.tensor_utils import (
+    colo_model_data_move_to_cpu,
+    colo_model_data_tensor_move,
+    colo_model_data_tensor_move_inline,
+    colo_model_tensor_clone,
+    colo_tensor_mem_usage,
+)
 
 
 def _run_colo_tensor_mem_usage():