mirror of https://github.com/hpcaitech/ColossalAI
110 lines
4.8 KiB
Python
110 lines
4.8 KiB
Python
from copy import copy
|
|
from typing import Dict, Optional
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
from .gemini_parallel import GeminiDDP
|
|
|
|
|
|
def zero_model_wrapper(model: nn.Module, zero_stage: int = 1, gemini_config: Optional[Dict] = None):
|
|
"""This wrapper function is used to wrap your training model for ZeRO DDP.
|
|
|
|
Example:
|
|
|
|
>>> with ColoInitContext():
|
|
>>> my_model = Bert()
|
|
>>> my_optim = SGD(my_model.parameters(), lr = 1e-3)
|
|
>>> zero_model = zero_model_wrapper(my_model, zero_stage=1)
|
|
>>> zero_optim = zero_optim_wrapper(zero_model, my_optim)
|
|
|
|
Args:
|
|
model (nn.Module): The model used in ZeRO DDP.
|
|
zero_stage (int, optional): The stage of ZeRO DDP. You can find more information in ZeRO's paper.
|
|
https://arxiv.org/abs/1910.02054
|
|
gemini_config (dict, optional): The configuration dictionary of `GeminiDDP`. `GeminiDDP` is enabled
|
|
when the stage is set to 3. You can set the arguemnts of `GeminiDDP` in the gemini_config.
|
|
Here is an example where we set the device of the model, the placement policy of Gemini, and the
|
|
size of hidden dimension to help Gemini find out a unified chunk size.
|
|
|
|
Example:
|
|
|
|
>>> config_dict = dict(device=torch.cuda.current_device(), hidden_dim=1024, placement_policy='auto')
|
|
>>> model = zero_model_wrapper(model, zero_stage=3, gemini_config=config_dict)
|
|
"""
|
|
assert zero_stage in [1, 2, 3], "The stage of ZeRO should be 1, 2 or 3"
|
|
|
|
if gemini_config is None:
|
|
gemini_config = dict()
|
|
|
|
if zero_stage in [1, 2]:
|
|
wrapped_model = model
|
|
else:
|
|
wrapped_model = GeminiDDP(model, **gemini_config)
|
|
|
|
setattr(wrapped_model, "_colo_zero_stage", zero_stage)
|
|
|
|
return wrapped_model
|
|
|
|
|
|
def zero_optim_wrapper(model: nn.Module,
|
|
optimizer: torch.optim.Optimizer,
|
|
initial_scale: float = 2**16,
|
|
growth_factor: float = 2,
|
|
backoff_factor: float = 0.5,
|
|
growth_interval: int = 1000,
|
|
hysteresis: int = 2,
|
|
min_scale: float = 1,
|
|
max_scale: float = 2**32,
|
|
max_norm: float = 0.0,
|
|
norm_type: float = 2.0,
|
|
optim_config: Optional[Dict] = None):
|
|
"""This wrapper function is used to wrap your training optimizer for ZeRO DDP.
|
|
|
|
Args:
|
|
model (nn.Module): Your model wrapped by `zero_model_wrapper`
|
|
optimizer (torch.optim.Optimizer): Your initialized optimizer
|
|
initial_scale (float, optional): initial_scale used by DynamicGradScaler.
|
|
min_scale (float, optional): min_scale used by DynamicGradScaler.
|
|
growth_factor (float, optional): growth_factor used by DynamicGradScaler.
|
|
backoff_factor (float, optional): backoff_factor used by DynamicGradScaler.
|
|
growth_interval (float, optional): growth_interval used by DynamicGradScaler.
|
|
hysteresis (float, optional): hysteresis used by DynamicGradScaler.
|
|
max_scale (int, optional): max_scale used by DynamicGradScaler.
|
|
max_norm (float, optional): max_norm used for `clip_grad_norm`. You should notice that you shall not do
|
|
clip_grad_norm by yourself when using ZeRO DDP. The ZeRO optimizer will take care of clip_grad_norm.
|
|
norm_type (float, optional): norm_type used for `clip_grad_norm`.
|
|
optim_config (dict, optinoal): The configuration used for the ZeRO optimizer.
|
|
Example:
|
|
|
|
>>> zero2_config = dict(reduce_bucket_size=12 * 1024 * 1024, overlap_communication=True)
|
|
>>> optim = zero_optim_wrapper(model, optim, optim_config=zero2_config)
|
|
"""
|
|
assert hasattr(model, "_colo_zero_stage"), "You should use `zero_ddp_wrapper` first"
|
|
zero_stage = getattr(model, "_colo_zero_stage")
|
|
|
|
assert norm_type == 2.0, "Current ZeRO optimizers only support 'norm_type=2'"
|
|
|
|
if optim_config is None:
|
|
config_dict = dict()
|
|
else:
|
|
config_dict = copy(optim_config)
|
|
|
|
config_dict['initial_scale'] = initial_scale
|
|
config_dict['growth_factor'] = growth_factor
|
|
config_dict['backoff_factor'] = backoff_factor
|
|
config_dict['growth_interval'] = growth_interval
|
|
config_dict['hysteresis'] = hysteresis
|
|
config_dict['min_scale'] = min_scale
|
|
config_dict['max_scale'] = max_scale
|
|
|
|
if zero_stage in [1, 2]:
|
|
from colossalai.zero.sharded_optim.low_level_optim import LowLevelZeroOptimizer
|
|
config_dict['partition_grad'] = zero_stage == 2
|
|
config_dict['clip_grad_norm'] = max_norm
|
|
return LowLevelZeroOptimizer(optimizer, **config_dict)
|
|
else:
|
|
from colossalai.nn.optimizer.zero_optimizer import ZeroOptimizer
|
|
config_dict['clipping_norm'] = max_norm
|
|
return ZeroOptimizer(optimizer, model, **config_dict)
|