[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

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

@@ -14,17 +14,16 @@ from transformers.tokenization_utils_base import PreTrainedTokenizerBase
 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 @@ from typing import Optional, Set
 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

-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 @@ class Region:
         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 @@ class Region:
     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)
\ No newline at end of file

colossalai/booster/plugin/gemini_plugin.py

@@ -14,12 +14,12 @@ from torch.utils.data.distributed import DistributedSampler
 
 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 torch.nn.modules.loss import _Loss
 
 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 @@ class PipelineSchedule(BaseSchedule):
         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/gemini/__init__.py

-from .chunk import ChunkManager, TensorInfo, TensorState, search_chunk_configuration
-from .gemini_mgr import GeminiManager
-from .stateful_tensor_mgr import StatefulTensorMgr
-from .tensor_placement_policy import TensorPlacementPolicyFactory
-
-__all__ = [
-    'StatefulTensorMgr', 'TensorPlacementPolicyFactory', 'GeminiManager', 'TensorInfo', 'TensorState', 'ChunkManager',
-    'search_chunk_configuration'
-]

colossalai/initialize.py

@@ -29,13 +29,12 @@ from colossalai.engine.schedule import (
     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 @@ import torch.nn as nn
 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.experts import Experts, MoeExperts
 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/optimizer/gemini_optimizer.py

-from typing import Any
-
-import torch
-
-from colossalai.nn.optimizer import HybridAdam
-from colossalai.nn.optimizer.zero_optimizer import ZeroOptimizer
-
-__all__ = ['GeminiAdamOptimizer']
-
-
-class GeminiAdamOptimizer(ZeroOptimizer):
-
-    def __init__(self, model: torch.nn.Module, **defaults: Any) -> None:
-        optimizer = HybridAdam(model.parameters(), **defaults)
-        super().__init__(optimizer, model, **defaults)

colossalai/nn/parallel/__init__.py

-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',
+]

colossalai/nn/parallel/data_parallel.py

-import itertools
 from collections import OrderedDict
 from functools import partial
-from typing import Dict, Iterable, List, Optional, Set
+from typing import Iterable, Optional, Set
 
 import torch
 import torch.distributed as dist
-import torch.nn as nn
 
-from colossalai.gemini.chunk import Chunk, ChunkManager, TensorState
-from colossalai.gemini.gemini_mgr import GeminiManager
-from colossalai.gemini.memory_tracer import OrderedParamGenerator
-from colossalai.logging import get_dist_logger
-from colossalai.nn.parallel.utils import get_temp_total_chunk_on_cuda
 from colossalai.tensor import ProcessGroup as ColoProcessGroup
-from colossalai.tensor import ReplicaSpec
-from colossalai.tensor.colo_parameter import ColoParameter, ColoTensor, ColoTensorSpec
-from colossalai.tensor.param_op_hook import ColoParamOpHookManager
-from colossalai.utils import get_current_device, is_ddp_ignored
-from colossalai.zero.utils.gemini_hook import GeminiZeROHook
+from colossalai.utils import is_ddp_ignored
 
 from .reducer import Reducer
-from .utils import get_static_torch_model
-
-try:
-    from torch.nn.modules.module import _EXTRA_STATE_KEY_SUFFIX, _IncompatibleKeys
-except ImportError:
-    _EXTRA_STATE_KEY_SUFFIX = '_extra_state'
 
 
 def free_storage(data: torch.Tensor) -> None:
@@ -189,507 +172,3 @@ class ColoDDP(torch.nn.Module):
 
     def load_state_dict(self, state_dict: 'OrderedDict[str, torch.Tensor]', strict: bool = True):
         return self.module.load_state_dict(state_dict, strict)
-
-
-class ZeroDDP(ColoDDP):
-    """ZeRO DDP for ColoTensor.
-    Warning: Nested ZeroDDP is not supported now.
-    It is designed to be used with ChunkManager and GeminiManager.
-    For more details, see the API reference of ``ChunkManager`` and ``GeminiManager``.
-
-    Args:
-        module (torch.nn.Module): Module to apply ZeRO-DP.
-        gemini_manager (GeminiManager): Manages the chunk manager and heterogeneous momery space.
-            For more details, see the API reference of ``GeminiManager``.
-        pin_memory (bool): Chunks on CPU Memory use pin-memory.
-        force_outputs_fp32 (bool): If set to True, outputs will be fp32. Otherwise, outputs will be fp16.
-            Defaults to False.
-        strict_ddp_mode (bool): If set to True, there is no tensor sharding, each tensor is replicated.
-            Defaults to False. Users can set it to True, when they clearly know that they only need DDP.
-    """
-
-    def __init__(self,
-                 module: torch.nn.Module,
-                 gemini_manager: GeminiManager,
-                 pin_memory: bool = False,
-                 force_outputs_fp32: bool = False,
-                 strict_ddp_mode: bool = False) -> None:
-        super().__init__(module, process_group=ColoProcessGroup())
-        self.gemini_manager = gemini_manager
-        self.chunk_manager: ChunkManager = gemini_manager.chunk_manager
-        self.force_outputs_fp32 = force_outputs_fp32
-        self.param_op_hook = GeminiZeROHook(gemini_manager)
-        self.fp32_params: List[ColoTensor] = list()
-        self.fp16_params: List[ColoParameter] = list()
-        self.overflow_counter = 0
-        self.grads_device: Dict[torch.Tensor, torch.device] = dict()
-        self.param2name: Dict[nn.Parameter, str] = dict()
-        self.name2param: Dict[str, nn.Parameter] = dict()
-
-        self._cast_buffers()
-        self._logger = get_dist_logger()
-
-        if self.gemini_manager._premade_memstats_:
-            # build chunk in param runtime visited order.
-            param_order = self.gemini_manager.memstats()._param_runtime_order
-        else:
-            # build chunk in param initialized order.
-            # Note: in this way, it can not get filter unused params during runtime.
-            param_order = OrderedParamGenerator()
-            for p in module.parameters():
-                param_order.append(p)
-
-        self._init_chunks(param_order=param_order,
-                          strict_ddp_mode=strict_ddp_mode,
-                          cpu_offload=self.gemini_manager.policy_name != 'cuda',
-                          pin_memory=pin_memory)
-
-        for name, param in module.named_parameters():
-            self.param2name[param] = name
-        for m_name, m_var in module.named_modules():
-            for p_name, p_var in m_var.named_parameters(recurse=False):
-                param_name = m_name + '.' + p_name if m_name else p_name
-                self.name2param[param_name] = p_var
-
-    def _post_forward(self):
-        """This function is only triggered for inference.
-        """
-        access_list = list(self.chunk_manager.accessed_chunks)
-        # we need to scatter all accessed chunks and move them to their original places
-        for chunk in access_list:
-            if chunk.keep_gathered:
-                self.chunk_manager.fake_release_chunk(chunk)
-            else:
-                assert chunk.can_release
-                self.chunk_manager.release_chunk(chunk)
-            first_param = next(iter(chunk.tensors_info))
-            self.chunk_manager.move_chunk(chunk, self.grads_device[first_param])
-        assert self.chunk_manager.accessed_mem == 0
-        # reset all recorded attributes
-        self.gemini_manager.reset_attributes()
-
-    def forward(self, *args, **kwargs):
-        # check whether we are in a inference mode
-        grad_flag = torch.is_grad_enabled()
-        if not grad_flag:
-            assert not self.gemini_manager.need_warmup or not self.gemini_manager.is_warmup(
-            ), "You should run a completed iteration as your warmup iter"
-
-        args, kwargs = _cast_float(args, torch.half), _cast_float(kwargs, torch.half)
-        self.module.zero_grad(set_to_none=True)
-        self.gemini_manager.pre_iter(*args)
-        with ColoParamOpHookManager.use_hooks(self.param_op_hook):
-            outputs = self.module(*args, **kwargs)
-        # scatter chunks in the inference mode
-        if not grad_flag:
-            self._post_forward()
-
-        if self.force_outputs_fp32:
-            return _cast_float(outputs, torch.float)
-        return outputs
-
-    def _setup_grads_ptr(self):
-        for p in self.module.parameters():
-            if is_ddp_ignored(p):
-                continue
-            p.grad = None
-
-    def _pre_backward(self):
-        # set a visit label for all parameters
-        # the label is used to check whether the parameter is correctly reduced
-        for param in self.param2name:
-            if not is_ddp_ignored(param):
-                setattr(param, "_gemini_reduced", False)
-
-    def _post_backward(self):
-        if self.chunk_manager.accessed_mem != 0:
-            error_params = ["Reduction failed at followed parameters:"]
-            for param in self.param2name:
-                if not is_ddp_ignored(param) and not getattr(param, "_gemini_reduced"):
-                    error_params.append(self.param2name[param])
-            error_str = "\n\t".join(error_params)
-            raise RuntimeError("ZERO DDP error: the synchronization of gradients doesn't exit properly.",
-                               "The most possible reason is that the model is not compatible with ZeroDDP.\n",
-                               f"{error_str}")
-        self._setup_grads_ptr()
-        self._logger.debug(
-            f'comp cuda demand time: {self.gemini_manager._comp_cuda_demand_time}, layout time: {self.gemini_manager._layout_time}, evict time: {self.gemini_manager._evict_time}, CPU->CUDA vol: {self.gemini_manager._h2d_volume}B, CUDA->CPU vol: {self.gemini_manager._d2h_volume}'
-        )
-        self.gemini_manager.post_iter()
-
-    def backward(self, loss: torch.Tensor):
-        self._pre_backward()
-        with self.param_op_hook.switch_to_backward(), ColoParamOpHookManager.use_hooks(self.param_op_hook):
-            loss.backward()
-        self._post_backward()
-
-    def backward_by_grad(self, tensor, grad):
-        with self.param_op_hook.switch_to_backward(), ColoParamOpHookManager.use_hooks(self.param_op_hook):
-            torch.autograd.backward(tensor, grad)
-        self._post_backward()
-
-    def grad_handle(self, p, grad):
-        empty_grad = torch.empty_like(grad)
-        free_storage(empty_grad)
-        with torch._C.DisableTorchFunction():
-            chunk = self.chunk_manager.get_chunk(p)
-            if chunk.tensors_info[p].state != TensorState.HOLD_AFTER_BWD:
-                raise RuntimeError(f"Parameter `{self.param2name[p]}` failed at the gradient reduction. "
-                                   "Some unsupported torch function is operated upon this parameter.")
-            self.chunk_manager.trans_tensor_state(p, TensorState.READY_FOR_REDUCE)
-            chunk.copy_tensor_to_chunk_slice(p, grad)
-            reduced = self.chunk_manager.reduce_chunk(chunk)
-            if reduced:
-                if chunk.is_gathered:
-                    chunk.cuda_global_chunk.div_(chunk.pg_size)
-                else:
-                    chunk.cuda_shard.div_(chunk.pg_size)
-                # check overflow elements
-                self.overflow_counter += chunk.has_inf_or_nan
-                # record l2 norm for gradient clipping
-                if chunk.l2_norm_flag:
-                    chunk.set_l2_norm()
-                self.chunk_manager.move_chunk(chunk, self.grads_device[p], force_copy=True)
-        return empty_grad
-
-    def zero_grad(self, set_to_none: bool = False) -> None:
-        self.module.zero_grad(set_to_none=True)
-
-    def set_chunk_grad_device(self, chunk: Chunk, device: torch.device) -> None:
-        for tensor in chunk.get_tensors():
-            self.grads_device[tensor] = device
-
-    def state_dict(self, destination=None, prefix='', keep_vars=False, only_rank_0: bool = True):
-        """Returns a dictionary containing a whole state of the module.
-
-        Both parameters and persistent buffers (e.g. running averages) are included.
-        Keys are corresponding parameter and buffer names.
-        Parameters and buffers set to ``None`` are not included.
-
-        Warning: The non strict state dict would ignore the parameters if the tensors of the parameters
-            are shared with other parameters which have been included in the dictionary.
-            When you need to load the state dict, you should set the argument `strict` to False.
-
-        Returns:
-            dict:
-                a dictionary containing a whole state of the module
-        """
-        if destination is None:
-            destination = OrderedDict()
-            destination._metadata = OrderedDict()
-        destination._metadata[prefix[:-1]] = local_metadata = dict(version=self._version)
-        self._save_to_state_dict(destination, prefix, keep_vars, only_rank_0)
-
-        for hook in self._state_dict_hooks.values():
-            hook_result = hook(self, destination, prefix, local_metadata)
-            if hook_result is not None:
-                destination = hook_result
-        return destination
-
-    def _get_param_to_save_data(self, param_list: List[torch.nn.Parameter], only_rank_0: bool) -> Dict:
-        """
-        get param content from chunks.
-
-        Args:
-            param_list (_type_): a list of torch.nn.Parameters
-            only_rank_0 (_type_): _description_
-
-        Returns:
-            Dict: a dict whose key is param name and value is param with correct payload
-        """
-        # save parameters
-        param_to_save_data = dict()
-        chunk_list = self.chunk_manager.get_chunks(param_list)
-        for chunk in chunk_list:
-            temp_chunk = get_temp_total_chunk_on_cuda(chunk)
-
-            for tensor, tensor_info in chunk.tensors_info.items():
-                record_tensor = torch.empty([0])
-                record_flag = (not only_rank_0) | (dist.get_rank(chunk.torch_pg) == 0)
-                if record_flag:
-                    record_tensor = temp_chunk[tensor_info.offset:tensor_info.end].view(tensor.shape).cpu()
-
-                assert tensor not in param_to_save_data
-                param_to_save_data[tensor] = record_tensor
-
-            del temp_chunk
-        return param_to_save_data
-
-    def _save_to_state_dict(self, destination, prefix, keep_vars, only_rank_0=True):
-        r"""Saves module state to `destination` dictionary, containing a state
-        of the module, but not its descendants. This is called on every
-        submodule in :meth:`~torch.nn.Module.state_dict`.
-
-        In rare cases, subclasses can achieve class-specific behavior by
-        overriding this method with custom logic.
-
-        Args:
-            destination (dict): a dict where state will be stored
-            prefix (str): the prefix for parameters and buffers used in this
-                module
-        """
-        assert keep_vars is False, "`state_dict` with parameter, `keep_vars=True`, is not supported now."
-
-        # get copies of fp32 parameters in CPU
-        param_to_save_data = self._get_param_to_save_data(self.fp32_params, only_rank_0)
-        # get the mapping between copies and fp16 parameters
-        p_mapping = dict()
-        for p, fp32_p in zip(self.fp16_params, self.fp32_params):
-            name = self.param2name[p]
-            assert fp32_p in param_to_save_data, "Parameter '{}' is neglected in the chunk list".format(name)
-            record_parameter = param_to_save_data[fp32_p]
-            p_mapping[p] = record_parameter
-        for name, param in self.name2param.items():
-            if param is not None:
-                if is_ddp_ignored(param):
-                    # deal with ddp ignored parameters
-                    destination[prefix + name] = param if keep_vars else param.detach()
-                else:
-                    destination[prefix + name] = p_mapping[param]
-        del p_mapping
-        del param_to_save_data
-
-        # save all buffers
-        for name, buf in self.named_buffers():
-            if buf is not None and name not in self._non_persistent_buffers_set:
-                destination[prefix + name] = buf if keep_vars else buf.detach()
-        # save extra states
-        extra_state_key = prefix + _EXTRA_STATE_KEY_SUFFIX
-        if getattr(self.__class__, "get_extra_state",
-                   torch.nn.Module.get_extra_state) is not torch.nn.Module.get_extra_state:
-            destination[extra_state_key] = self.get_extra_state()
-
-    def load_state_dict(self, state_dict: 'OrderedDict[str, torch.Tensor]', strict: bool = True):
-        r"""Copies parameters and buffers from :attr:`state_dict` into
-        this module and its descendants. If :attr:`strict` is ``True``, then
-        the keys of :attr:`state_dict` must exactly match the keys returned
-        by this module's :meth:`~torch.nn.Module.state_dict` function.
-
-        Args:
-            state_dict (dict): a dict containing parameters and
-                persistent buffers.
-            strict (bool, optional): whether to strictly enforce that the keys
-                in :attr:`state_dict` match the keys returned by this module's
-                :meth:`~torch.nn.Module.state_dict` function. Default: ``True``
-
-        Returns:
-            ``NamedTuple`` with ``missing_keys`` and ``unexpected_keys`` fields:
-                * **missing_keys** is a list of str containing the missing keys
-                * **unexpected_keys** is a list of str containing the unexpected keys
-
-        Note:
-            If a parameter or buffer is registered as ``None`` and its corresponding key
-            exists in :attr:`state_dict`, :meth:`load_state_dict` will raise a
-            ``RuntimeError``.
-        """
-        missing_keys: List[str] = []
-        unexpected_keys: List[str] = []
-        error_msgs: List[str] = []
-
-        # copy state_dict so _load_from_state_dict can modify it
-        metadata = getattr(state_dict, '_metadata', None)
-        state_dict = state_dict.copy()
-        if metadata is not None:
-            # mypy isn't aware that "_metadata" exists in state_dict
-            state_dict._metadata = metadata    # type: ignore[attr-defined]
-
-        prefix = ''
-        local_metadata = {} if metadata is None else metadata.get(prefix[:-1], {})
-        self._load_from_state_dict(state_dict, prefix, local_metadata, True, missing_keys, unexpected_keys, error_msgs)
-
-        if strict:
-            if len(unexpected_keys) > 0:
-                error_msgs.insert(
-                    0, 'Unexpected key(s) in state_dict: {}. '.format(', '.join(
-                        '"{}"'.format(k) for k in unexpected_keys)))
-            if len(missing_keys) > 0:
-                error_msgs.insert(
-                    0, 'Missing key(s) in state_dict: {}. '.format(', '.join('"{}"'.format(k) for k in missing_keys)))
-
-        if len(error_msgs) > 0:
-            raise RuntimeError('Error(s) in loading state_dict for {}:\n\t{}'.format(
-                self.__class__.__name__, "\n\t".join(error_msgs)))
-        return _IncompatibleKeys(missing_keys, unexpected_keys)
-
-    def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys,
-                              error_msgs):
-        r"""Copies parameters and buffers from :attr:`state_dict` into only
-        this module, but not its descendants. This is called on every submodule
-        in :meth:`~torch.nn.Module.load_state_dict`. Metadata saved for this
-        module in input :attr:`state_dict` is provided as :attr:`local_metadata`.
-        For state dicts without metadata, :attr:`local_metadata` is empty.
-        Subclasses can achieve class-specific backward compatible loading using
-        the version number at `local_metadata.get("version", None)`.
-
-        .. note::
-            :attr:`state_dict` is not the same object as the input
-            :attr:`state_dict` to :meth:`~torch.nn.Module.load_state_dict`. So
-            it can be modified.
-
-        Args:
-            state_dict (dict): a dict containing parameters and
-                persistent buffers.
-            prefix (str): the prefix for parameters and buffers used in this
-                module
-            local_metadata (dict): a dict containing the metadata for this module.
-                See
-            strict (bool): whether to strictly enforce that the keys in
-                :attr:`state_dict` with :attr:`prefix` match the names of
-                parameters and buffers in this module
-            missing_keys (list of str): if ``strict=True``, add missing keys to
-                this list
-            unexpected_keys (list of str): if ``strict=True``, add unexpected
-                keys to this list
-            error_msgs (list of str): error messages should be added to this
-                list, and will be reported together in
-                :meth:`~torch.nn.Module.load_state_dict`
-        """
-        for hook in self._load_state_dict_pre_hooks.values():
-            hook(state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs)
-
-        persistent_buffers = {k: v for k, v in self.named_buffers() if k not in self._non_persistent_buffers_set}
-        local_name_params = itertools.chain(self.named_parameters(), persistent_buffers.items())
-        local_state = {k: v for k, v in local_name_params if v is not None}
-
-        def load(param_name, dest_tensor, copy_func):
-            state_key = prefix + param_name
-            if state_key in state_dict:
-                input_param = state_dict[state_key]
-                # Backward compatibility: loading 1-dim tensor from 0.3.* to version 0.4+
-                if len(dest_tensor.shape) == 0 and len(input_param.shape) == 1:
-                    input_param = input_param[0]
-                if input_param.shape != dest_tensor.shape:
-                    # local shape should match the one in checkpoint
-                    error_msgs.append('size mismatch for {}: copying a param with shape {} from checkpoint, '
-                                      'the shape in current model is {}.'.format(state_key, input_param.shape,
-                                                                                 dest_tensor.shape))
-                    return
-                try:
-                    with torch.no_grad():
-                        copy_func(input_param)
-                except Exception as ex:
-                    error_msgs.append('While copying the parameter named "{}", '
-                                      'whose dimensions in the model are {} and '
-                                      'whose dimensions in the checkpoint are {}, '
-                                      'an exception occurred : {}.'.format(state_key, dest_tensor.size(),
-                                                                           input_param.size(), ex.args))
-            elif strict:
-                missing_keys.append(state_key)
-
-        def load_fp32_parameter(chunk_slice, data):
-            chunk_slice.copy_(data.flatten())
-
-        for name, param in self.named_parameters():
-            if is_ddp_ignored(param):
-                # deal with ddp ignored parameters
-                load(name, param, param.copy_)
-
-        fp32_to_name = dict()
-        for p, fp32_p in zip(self.fp16_params, self.fp32_params):
-            if p is not None:
-                name = self.param2name[p]
-                fp32_to_name[fp32_p] = name
-
-        chunk_list = self.chunk_manager.get_chunks(self.fp32_params)
-        for chunk in chunk_list:
-            temp_chunk = get_temp_total_chunk_on_cuda(chunk)
-
-            for tensor, tensor_info in chunk.tensors_info.items():
-                parameter_name = fp32_to_name[tensor]
-                parameter_slice = temp_chunk[tensor_info.offset:tensor_info.end]
-                load(parameter_name, tensor, partial(load_fp32_parameter, parameter_slice))
-
-            if chunk.is_gathered:
-                chunk.cuda_global_chunk.copy_(temp_chunk)
-            elif chunk.cuda_shard is not None:
-                chunk.cuda_shard.copy_(temp_chunk[chunk.shard_begin:chunk.shard_end])
-            else:
-                chunk.cpu_shard.copy_(temp_chunk[chunk.shard_begin:chunk.shard_end])
-
-            del temp_chunk
-
-        for chunk_32 in chunk_list:
-            chunk_16 = chunk_32.paired_chunk
-            assert chunk_16 is not None
-            chunk_16.optim_update()
-
-        for name, buf in persistent_buffers.items():
-            if buf is not None:
-                load(name, buf, buf.copy_)
-
-        extra_state_key = prefix + _EXTRA_STATE_KEY_SUFFIX
-        if getattr(self.__class__, "set_extra_state",
-                   torch.nn.Module.set_extra_state) is not torch.nn.Module.set_extra_state:
-            if extra_state_key in state_dict:
-                self.set_extra_state(state_dict[extra_state_key])
-            elif strict:
-                missing_keys.append(extra_state_key)
-        elif strict and (extra_state_key in state_dict):
-            unexpected_keys.append(extra_state_key)
-
-        if strict:
-            for key in state_dict.keys():
-                if key.startswith(prefix) and key != extra_state_key:
-                    input_name = key[len(prefix):]
-                    if input_name not in local_state:
-                        unexpected_keys.append(key)
-
-    def _init_chunks(self, param_order, strict_ddp_mode: bool, cpu_offload: bool, pin_memory: bool):
-        ddp_pg = ColoProcessGroup()
-        for p in param_order.generate():
-            assert isinstance(p, ColoParameter)
-
-            # gather sharded parameters in the strict ddp mode
-            if strict_ddp_mode:
-                if not p.is_replicate():
-                    p.set_dist_spec(ReplicaSpec())
-                p.set_process_group(pg=ddp_pg)
-
-            # ignore the parameters with no gradient
-            if not p.requires_grad:
-                self.set_params_to_ignore([p])
-
-            # move ignored parameters to CUDA
-            if is_ddp_ignored(p):
-                p.data = p.data.to(device=get_current_device(), dtype=torch.float16)
-                continue
-
-            # create a fp32 parameter
-            fp32_data = p.data.float()
-            fp32_p = ColoTensor(fp32_data, spec=ColoTensorSpec(p.process_group))
-            # create a fp16 parameter
-            p.data = p.data.half()
-
-            # register the fp16 parameter and fp32 parameter in the chunk manager
-            dp_world_size = p.process_group.dp_world_size()
-            self.chunk_manager.register_tensor(tensor=p,
-                                               group_type='fp16_param',
-                                               config_key=dp_world_size,
-                                               cpu_offload=cpu_offload,
-                                               pin_memory=pin_memory)
-            self.chunk_manager.register_tensor(tensor=fp32_p,
-                                               group_type='fp32_param',
-                                               config_key=dp_world_size,
-                                               cpu_offload=cpu_offload,
-                                               pin_memory=pin_memory)
-
-            self.fp16_params.append(p)
-            self.fp32_params.append(fp32_p)
-            self.grads_device[p] = self.gemini_manager.default_device
-
-        self.chunk_manager.close_all_groups()
-
-        for p, fp32_p in zip(self.fp16_params, self.fp32_params):
-            chunk_16 = self.chunk_manager.get_chunk(p)
-            chunk_32 = self.chunk_manager.get_chunk(fp32_p)
-            chunk_32.init_pair(chunk_16)
-
-            # keep gathered chunks are in CUDA
-            if chunk_16.keep_gathered:
-                self.grads_device[p] = get_current_device()
-
-    def _cast_buffers(self):
-        for buffer in self.module.buffers():
-            buffer.data = buffer.cuda()
-            if torch.is_floating_point(buffer):
-                buffer.data = buffer.half()

colossalai/nn/parallel/gemini_parallel.py

-from typing import Optional
-
-import torch
-
-from colossalai.gemini.chunk import init_chunk_manager
-from colossalai.gemini.gemini_mgr import GeminiManager
-from colossalai.gemini.memory_tracer import MemStats
-
-from .data_parallel import ZeroDDP
-
-
-class GeminiDDP(ZeroDDP):
-
-    def __init__(self,
-                 module: torch.nn.Module,
-                 device: torch.device,
-                 placement_policy: str = "cpu",
-                 pin_memory: bool = False,
-                 force_outputs_fp32: bool = False,
-                 strict_ddp_mode: bool = False,
-                 search_range_mb: int = 32,
-                 hidden_dim: Optional[int] = None,
-                 min_chunk_size_mb: float = 32,
-                 memstats: Optional[MemStats] = None) -> None:
-        """
-        A torch.Module warpper using ZeRO-DP and Genimi.
-        ZeRO is for parallel. Gemini is for memory management.
-        WARNING: The class will modify the module inline!
-
-        Example:
-            model is initialized under the context of ColoInitContext
-            >>> model = GeminiDDP(model, torch.cuda.current_device(), "cuda")
-            >>> logits = model(x)
-            >>> loss = criterion(logits, labels)
-            >>> model.backward(loss)
-
-        Args:
-            module (torch.nn.Module): the model to be wrapped.
-            device (torch.device): device to place the model.
-            placement_policy (str, optional): "cpu", "cuda", "auto". Defaults to "cpu".
-            pin_memory (bool, optional): use pin memory on CPU. Defaults to False.
-            force_outputs_fp32 (bool, optional): force outputs are fp32. Defaults to False.
-            search_range_mb (int, optional): chunk size searching range in MegaByte. Defaults to 32.
-            hidden_dim (int, optional): the hidden dimension of DNN.
-                Users can provide this argument to speed up searching.
-                If users do not know this argument before training, it is ok. We will use a default value 1024.
-            min_chunk_size_mb (float, optional): the minimum chunk size in MegaByte.
-                If the aggregate size of parameters is still samller than the minimum chunk size,
-                all parameters will be compacted into one small chunk.
-            memstats (MemStats, optional) the memory statistics collector by a runtime memory tracer.
-        """
-        # some ugly hotfix for the compatibility with Lightning
-        if search_range_mb is None:
-            search_range_mb = 32
-
-        chunk_manager = init_chunk_manager(model=module,
-                                           init_device=device,
-                                           hidden_dim=hidden_dim,
-                                           search_range_mb=search_range_mb,
-                                           min_chunk_size_mb=min_chunk_size_mb,
-                                           strict_ddp_flag=strict_ddp_mode)
-        gemini_manager = GeminiManager(placement_policy, chunk_manager, memstats)
-        super().__init__(module, gemini_manager, pin_memory, force_outputs_fp32, strict_ddp_mode)

colossalai/zero/__init__.py

-from typing import Tuple
-
-import torch
-import torch.nn as nn
-
-from colossalai.logging import get_dist_logger
-from colossalai.zero.sharded_model.sharded_model_v2 import ShardedModelV2
-from colossalai.zero.sharded_optim import LowLevelZeroOptimizer, ShardedOptimizerV2
-
-from ..nn.optimizer.zero_optimizer import ZeroOptimizer
-
-
-def convert_to_zero_v2(model: nn.Module, optimizer: torch.optim.Optimizer, model_config,
-                       optimizer_config) -> Tuple[ShardedModelV2, ShardedOptimizerV2]:
-    """
-    A helper function to integrate the model and optimizer with ZeRO optimizer and off-loading
-
-    :param model: Your model object
-    :type model: :class:`torch.nn.Module`
-    :param optimizer_config: Your optimizer object
-    :type optimizer_config: :class:`dict`
-
-    :return: (model, optimizer)
-    :rtype: Tuple
-    """
-
-    logger = get_dist_logger('convert_to_zero_v2')
-
-    logger.info(f'optimizer_config is {optimizer_config}', ranks=[0])
-    if optimizer_config is None:
-        optimizer_config = dict()
-    logger.info(f'model_config is {model_config}', ranks=[0])
-    if model_config is None:
-        model_config = dict()
-
-    zero_model = ShardedModelV2(model, **model_config)
-    zero_optimizer = ShardedOptimizerV2(zero_model, optimizer, **optimizer_config)
-    return zero_model, zero_optimizer
-
-
-__all__ = ['convert_to_zero_v2', 'LowLevelZeroOptimizer', 'ShardedModelV2', 'ShardedOptimizerV2', 'ZeroOptimizer']
+from .gemini import (
+    ColoInitContext,
+    GeminiAdamOptimizer,
+    GeminiDDP,
+    ZeroDDP,
+    ZeroOptimizer,
+    get_static_torch_model,
+    post_process_colo_init_ctx,
+)
+from .low_level import LowLevelZeroOptimizer
+from .wrapper import zero_model_wrapper, zero_optim_wrapper
+
+__all__ = [
+    'ZeroDDP', 'GeminiDDP', 'ZeroOptimizer', 'GeminiAdamOptimizer', 'zero_model_wrapper', 'zero_optim_wrapper',
+    'LowLevelZeroOptimizer', 'ColoInitContext', 'post_process_colo_init_ctx', 'get_static_torch_model'
+]

colossalai/zero/gemini/__init__.py

+from .chunk import ChunkManager, TensorInfo, TensorState, search_chunk_configuration
+from .colo_init_context import ColoInitContext, post_process_colo_init_ctx
+from .gemini_ddp import GeminiDDP, ZeroDDP
+from .gemini_mgr import GeminiManager
+from .gemini_optimizer import GeminiAdamOptimizer, ZeroOptimizer
+from .utils import get_static_torch_model
+
+__all__ = [
+    'GeminiManager', 'TensorInfo', 'TensorState', 'ChunkManager', 'search_chunk_configuration', 'ZeroDDP', 'GeminiDDP',
+    'get_static_torch_model', 'GeminiAdamOptimizer', 'ZeroOptimizer', 'ColoInitContext', 'post_process_colo_init_ctx'
+]

colossalai/gemini/chunk/manager.py → colossalai/zero/gemini/chunk/manager.py

@@ -3,10 +3,11 @@ from typing import Deque, Dict, Iterable, List, Optional, Set, Tuple
 
 import torch
 
-from colossalai.gemini.chunk import Chunk, ChunkFullError, TensorState
 from colossalai.tensor import ColoTensor
 from colossalai.utils import get_current_device
 
+from .chunk import Chunk, ChunkFullError, TensorState
+
 
 class ChunkManager:
     """

colossalai/gemini/chunk/search_utils.py → colossalai/zero/gemini/chunk/search_utils.py

@@ -5,9 +5,9 @@ import numpy as np
 import torch.distributed as dist
 import torch.nn as nn
 
-from colossalai.gemini.memory_tracer import MemStats, OrderedParamGenerator
 from colossalai.tensor import ColoParameter
 from colossalai.utils import is_ddp_ignored
+from colossalai.zero.gemini.memory_tracer import MemStats, OrderedParamGenerator
 
 
 def _filter_exlarge_params(model: nn.Module, size_dict: Dict[int, List[int]]) -> None: