[hotfix] fix initialize bug with zero (#442)

colossalai/amp/__init__.py

@@ -11,17 +11,13 @@ from .apex_amp import convert_to_apex_amp
 from .naive_amp import convert_to_naive_amp
 
 
-def convert_to_amp(model: nn.Module,
-                   optimizer: Optimizer,
-                   criterion: _Loss,
-                   mode: AMP_TYPE,
-                   amp_config: Config = None):
+def convert_to_amp(model: nn.Module, optimizer: Optimizer, criterion: _Loss, mode: AMP_TYPE, amp_config: Config = None):
     """A helper function to wrap training components with Torch AMP modules
 
     :param model: your model object
     :type model: :class:`torch.nn.Module`
     :param optimizer: your optimizer object
-    :type optimizer: :class:`torch.optim.Optimzer`
+    :type optimizer: :class:`torch.optim.Optimizer`
     :param criterion: your loss function object
     :type criterion: :class:`torch.nn.modules.loss._Loss`
     :param mode: amp mode

colossalai/amp/apex_amp/__init__.py

@@ -3,15 +3,13 @@ import torch.nn as nn
 from torch.optim import Optimizer
 
 
-def convert_to_apex_amp(model: nn.Module,
-                        optimizer: Optimizer,
-                        amp_config):
+def convert_to_apex_amp(model: nn.Module, optimizer: Optimizer, amp_config):
     """A helper function to wrap training components with Apex AMP modules
 
     :param model: your model object
     :type model: :class:`torch.nn.Module`
     :param optimizer: your optimizer object
-    :type optimizer: :class:`torch.optim.Optimzer`
+    :type optimizer: :class:`torch.optim.Optimizer`
     :param amp_config: configuration for nvidia apex
     :type amp_config: :class:`colossalai.context.Config` or dict
 

colossalai/amp/naive_amp/__init__.py

@@ -12,7 +12,7 @@ def convert_to_naive_amp(model: nn.Module, optimizer: Optimizer, amp_config):
     :param model: your model object
     :type model: :class:`torch.nn.Module`
     :param optimizer: your optimizer object
-    :type optimizer: :class:`torch.optim.Optimzer`
+    :type optimizer: :class:`torch.optim.Optimizer`
     :param amp_config: configuration for naive mode amp
     :type amp_config: :class:`colossalai.context.Config` or dict
 

colossalai/amp/torch_amp/__init__.py

@@ -15,7 +15,7 @@ def convert_to_torch_amp(model: nn.Module,
     :param model: your model object
     :type model: :class:`torch.nn.Module`
     :param optimizer: your optimizer object
-    :type optimizer: :class:`torch.optim.Optimzer`
+    :type optimizer: :class:`torch.optim.Optimizer`
     :param criterion: your loss function object
     :type criterion: :class:`torch.nn.modules.loss._Loss`, optional
     :param amp_config: configuration for different amp modes

colossalai/initialize.py

@@ -268,6 +268,7 @@ def initialize(model: Union[Callable, nn.Module],
     if verbose:
         logger.info(f"cuDNN benchmark = {cudnn_benchmark}, deterministic = {cudnn_deterministic}", ranks=[0])
 
+    # zero
     use_zero = hasattr(gpc.config, 'zero')
     if use_zero:
         zero_cfg = gpc.config.get('zero', None)
@@ -275,10 +276,13 @@ def initialize(model: Union[Callable, nn.Module],
             cfg_ = zero_cfg.copy()
         else:
             cfg_ = {}
-        optimizer_config = zero_cfg.get('optimzer', None)
-        model, optimizer = convert_to_zero_v2(model_builder=model, optimizer_config=optimizer_config)
+        optimizer_config = zero_cfg.get('optimizer_config', None)
+        model_config = zero_cfg.get('model_config', None)
+        model, optimizer = convert_to_zero_v2(model_builder=model,
+                                              model_config=model_config,
+                                              optimizer_config=optimizer_config)
 
-        logger.info("Initializing ZeRO model and optimzer finished!", ranks=[0])
+        logger.info("Initializing ZeRO model and optimizer finished!", ranks=[0])
         #FIXME() throw a warning if using zero with MP
         if gpc.get_world_size(ParallelMode.MODEL) > 1:
             logger.warning("ZeRO currently has not been tested with model parallelism.", ranks=[0])
@@ -289,6 +293,11 @@ def initialize(model: Union[Callable, nn.Module],
         elif isinstance(model, Callable):
             model = model().to(get_current_device())
 
+        # optimizer maybe a optimizer_cls
+        logger.warning("Initializing an non ZeRO model with optimizer class")
+        if isinstance(optimizer, Callable):
+            optimizer = optimizer(model.parameters())
+
     if not moe_env.is_initialized() and not use_zero:
         if is_using_sequence():
             sync_model_param(model, ParallelMode.SEQUENCE_DP)

colossalai/utils/commons/memory.py

@@ -1,4 +1,3 @@
-import imp
 import torch
 from colossalai.utils import get_current_device
 

colossalai/zero/__init__.py

@@ -17,7 +17,7 @@ from colossalai.core import global_context as gpc
 from colossalai.logging import get_dist_logger
 
 
-def convert_to_zero_v2(model_builder: Callable, optimizer_config) -> (ShardedModelV2, ShardedOptimizerV2):
+def convert_to_zero_v2(model_builder: Callable, model_config, optimizer_config) -> (ShardedModelV2, ShardedOptimizerV2):
     """
     A helper function to integrate the model and optimizer with ZeRO optimizer and off-loading
 
@@ -35,28 +35,26 @@ def convert_to_zero_v2(model_builder: Callable, optimizer_config) -> (ShardedMod
     # FIXME() pass shard strategy from config
     shard_strategy = TensorShardStrategy()
 
+    logger.info(f'optimizer_config is {optimizer_config}')
+    if optimizer_config is None:
+        optimizer_config = dict()
+    logger.info(f'model_config is {model_config}')
+    if model_config is None:
+        model_config = dict()
+
     if isinstance(model_builder, nn.Module):
         model = model_builder
     elif isinstance(model_builder, Callable):
         with ZeroInitContext(convert_fp16='fp16' in gpc.config,
                              target_device=torch.cuda.current_device(),
                              shard_strategy=shard_strategy,
-                             shard_param=True):
+                             shard_param=model_config.get('shard_param', True)):
             model = model_builder()
     else:
         raise TypeError(f"convert_to_zero_v2 dose not support model_builder of type {type(convert_to_zero_v2)}")
 
-    zero_model = ShardedModelV2(model, shard_strategy=shard_strategy)
-
-    optimizer_class = optimizer_config.get('optimizer_type', None)
-    if optimizer_class is None:
-        raise RuntimeError("Set optimizer_class in zero_config")
-    logger.info(f'optimizer class is {optimizer_class}')
-
-    cfg = optimizer_config.get('optimizer_config', None)
-    logger.info(f'optimizer_config is {cfg}')
-
-    zero_optimizer = ShardedOptimizerV2(zero_model, optimizer_class, **optimizer_config.get('optimizer_config', None))
+    zero_model = ShardedModelV2(model, shard_strategy=shard_strategy, **model_config)
+    zero_optimizer = ShardedOptimizerV2(zero_model, **optimizer_config)
     return zero_model, zero_optimizer
 
 

colossalai/zero/sharded_model/sharded_model_v2.py

@@ -1,5 +1,4 @@
 import functools
-from asyncio.log import logger
 from collections import OrderedDict
 from typing import Any, Optional