[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/zero/sharded_model/sharded_model_v2.py → colossalai/zero/legacy/sharded_model/sharded_model_v2.py

@@ -13,19 +13,18 @@ from torch.nn.parameter import Parameter
 
 from colossalai.context.parallel_mode import ParallelMode
 from colossalai.core import global_context as gpc
-from colossalai.gemini.memory_tracer import MemStatsCollector, StaticMemStatsCollector
-from colossalai.gemini.ophooks import register_ophooks_recursively
-from colossalai.gemini.paramhooks import BaseParamHookMgr
-from colossalai.gemini.stateful_tensor import TensorState
-from colossalai.gemini.stateful_tensor_mgr import StatefulTensorMgr
-from colossalai.gemini.tensor_placement_policy import TensorPlacementPolicy, TensorPlacementPolicyFactory
-from colossalai.gemini.tensor_utils import colo_model_data_move_to_cpu
 from colossalai.logging import get_dist_logger
 from colossalai.utils import disposable, get_current_device
 from colossalai.utils.memory import colo_device_memory_capacity
-from colossalai.zero.shard_utils import BaseShardStrategy
-from colossalai.zero.sharded_model.reduce_scatter import ReduceScatterBucketer
-from colossalai.zero.utils import ZeroHook
+from colossalai.zero.gemini.memory_tracer import MemStatsCollector, StaticMemStatsCollector
+from colossalai.zero.legacy.gemini.ophooks import register_ophooks_recursively
+from colossalai.zero.legacy.gemini.paramhooks import BaseParamHookMgr
+from colossalai.zero.legacy.gemini.stateful_tensor import TensorState
+from colossalai.zero.legacy.gemini.stateful_tensor_mgr import StatefulTensorMgr
+from colossalai.zero.legacy.gemini.tensor_placement_policy import TensorPlacementPolicy, TensorPlacementPolicyFactory
+from colossalai.zero.legacy.gemini.tensor_utils import colo_model_data_move_to_cpu
+from colossalai.zero.legacy.shard_utils import BaseShardStrategy
+from colossalai.zero.legacy.sharded_model.reduce_scatter import ReduceScatterBucketer
 
 from ._utils import (
     cast_float_arguments,
@@ -35,6 +34,7 @@ from ._utils import (
     free_storage,
     get_gradient_predivide_factor,
 )
+from .zero_hook import ZeroHook
 
 try:
     from torch.nn.modules.module import _EXTRA_STATE_KEY_SUFFIX

colossalai/zero/sharded_model/utils.py → colossalai/zero/legacy/sharded_model/utils.py

+import copy
+
 import torch
-from colossalai.zero.sharded_model import ShardedModelV2
 
-import copy
+from colossalai.zero.legacy.sharded_model import ShardedModelV2
 
 
 def col_model_deepcopy(sharded_model: ShardedModelV2, other_model: torch.nn.Module):

colossalai/zero/utils/zero_hook.py → colossalai/zero/legacy/sharded_model/zero_hook.py

@@ -3,14 +3,14 @@ from typing import Optional
 import torch
 import torch.distributed as dist
 
-from colossalai.gemini.memory_tracer import MemStatsCollector
-from colossalai.gemini.ophooks import BaseOpHook
-from colossalai.gemini.stateful_tensor import TensorState
-from colossalai.gemini.stateful_tensor_mgr import StatefulTensorMgr
 from colossalai.logging import get_dist_logger
 from colossalai.registry import OPHOOKS
 from colossalai.utils import get_current_device
-from colossalai.zero.shard_utils import BaseShardStrategy
+from colossalai.zero.gemini.memory_tracer import MemStatsCollector
+from colossalai.zero.legacy.gemini.ophooks import BaseOpHook
+from colossalai.zero.legacy.gemini.stateful_tensor import TensorState
+from colossalai.zero.legacy.gemini.stateful_tensor_mgr import StatefulTensorMgr
+from colossalai.zero.legacy.shard_utils import BaseShardStrategy
 
 
 @OPHOOKS.register_module

colossalai/zero/legacy/sharded_optim/__init__.py

+from .sharded_optim_v2 import ShardedOptimizerV2
+
+__all__ = ['ShardedOptimizerV2']

colossalai/zero/sharded_optim/sharded_optim_v2.py → colossalai/zero/legacy/sharded_optim/sharded_optim_v2.py

@@ -14,13 +14,13 @@ from torch.optim import Optimizer
 from colossalai.amp.naive_amp.grad_scaler import DynamicGradScaler
 from colossalai.context.parallel_mode import ParallelMode
 from colossalai.core import global_context as gpc
-from colossalai.gemini.stateful_tensor import StatefulTensor, TensorState
-from colossalai.gemini.tensor_placement_policy import AutoTensorPlacementPolicy
-from colossalai.gemini.tensor_utils import colo_model_data_tensor_move_inline, colo_tensor_mem_usage
 from colossalai.logging import get_dist_logger
 from colossalai.nn.optimizer import ColossalaiOptimizer
-from colossalai.zero.sharded_model import ShardedModelV2
-from colossalai.zero.sharded_model._utils import cast_tensor_to_fp32
+from colossalai.zero.legacy.gemini.stateful_tensor import StatefulTensor, TensorState
+from colossalai.zero.legacy.gemini.tensor_placement_policy import AutoTensorPlacementPolicy
+from colossalai.zero.legacy.gemini.tensor_utils import colo_model_data_tensor_move_inline, colo_tensor_mem_usage
+from colossalai.zero.legacy.sharded_model import ShardedModelV2
+from colossalai.zero.legacy.sharded_model._utils import cast_tensor_to_fp32
 
 
 class OptimState(Enum):

colossalai/zero/legacy/sharded_param/__init__.py

+from .sharded_param import ShardedParamV2
+from .sharded_tensor import ShardedTensor
+
+__all__ = ['ShardedTensor', 'ShardedParamV2']

colossalai/zero/sharded_param/sharded_param.py → colossalai/zero/legacy/sharded_param/sharded_param.py

+from typing import List, Optional, Tuple
+
 import torch
-from typing import Optional, Tuple
-from colossalai.zero.sharded_param.sharded_tensor import ShardedTensor
-from colossalai.gemini.tensor_utils import colo_tensor_mem_usage
-from colossalai.gemini.stateful_tensor import StatefulTensor, TensorState
-from typing import List
+
+from colossalai.zero.legacy.gemini.stateful_tensor import StatefulTensor, TensorState
+from colossalai.zero.legacy.gemini.tensor_utils import colo_tensor_mem_usage
+
+from .sharded_tensor import ShardedTensor
 
 EMPTY_TENSOR_DICT = {}
 

colossalai/zero/sharded_param/sharded_tensor.py → colossalai/zero/legacy/sharded_param/sharded_tensor.py

 import torch
-from colossalai.gemini.stateful_tensor import StatefulTensor, TensorState
+
+from colossalai.zero.legacy.gemini.stateful_tensor import StatefulTensor, TensorState
 
 
 class ShardedTensor(StatefulTensor):

colossalai/zero/sharded_optim/__init__.py → colossalai/zero/low_level/__init__.py

 from .low_level_optim import LowLevelZeroOptimizer
-from .sharded_optim_v2 import ShardedOptimizerV2
 
-__all__ = ['ShardedOptimizerV2', 'LowLevelZeroOptimizer']
+__all__ = ['LowLevelZeroOptimizer']

colossalai/zero/sharded_param/__init__.py

-from colossalai.zero.sharded_param.sharded_tensor import ShardedTensor
-from colossalai.zero.sharded_param.sharded_param import ShardedParamV2
-
-__all__ = ['ShardedTensor', 'ShardedParamV2']

colossalai/zero/utils/__init__.py

-from .zero_hook import ZeroHook
-
-__all__ = ['ZeroHook']
\ No newline at end of file