mirror of https://github.com/hpcaitech/ColossalAI
You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
93 lines
3.5 KiB
93 lines
3.5 KiB
from asyncio.log import logger
|
|
from distutils.command.config import config
|
|
from colossalai.zero.sharded_model.sharded_model_v2 import ShardedModelV2
|
|
from colossalai.zero.sharded_optim.sharded_optim_v2 import ShardedOptimizerV2
|
|
from colossalai.zero.shard_utils import TensorShardStrategy
|
|
import torch
|
|
import torch.nn as nn
|
|
from colossalai.amp.naive_amp import NaiveAMPModel
|
|
from colossalai.context.parallel_mode import ParallelMode
|
|
from colossalai.core import global_context as gpc
|
|
from torch.optim import Optimizer
|
|
from .sharded_model import ShardedModel
|
|
from .sharded_optim import ShardedOptimizer
|
|
from colossalai.zero.init_ctx import ZeroInitContext
|
|
from typing import Callable, Type
|
|
from colossalai.core import global_context as gpc
|
|
from colossalai.logging import get_dist_logger
|
|
|
|
|
|
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
|
|
|
|
: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')
|
|
|
|
# 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=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, **model_config)
|
|
zero_optimizer = ShardedOptimizerV2(zero_model, **optimizer_config)
|
|
return zero_model, zero_optimizer
|
|
|
|
|
|
def convert_to_zero(model: nn.Module, optimizer: Optimizer, level: int, zero_config: dict):
|
|
"""
|
|
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: Your optimizer object
|
|
:type optimizer: :class:`torch.optim.Optimizer`
|
|
:param level: Optimizer level, can be 2 or 3
|
|
:type level: int
|
|
:param zero_config: Configuration for zero
|
|
:type zero_config: dict
|
|
|
|
:return: (model, optimizer)
|
|
:rtype: Tuple
|
|
"""
|
|
assert 1 <= level <= 3, 'Only ZERO Optimizer Level 1-3 are provided'
|
|
if level in [1, 2]:
|
|
if level == 2:
|
|
if 'partition_grad' in zero_config:
|
|
assert zero_config['partition_grad'], \
|
|
'Sharded Optimizer requires partition_grad to be True'
|
|
else:
|
|
zero_config['partiton_grad'] = True
|
|
model = NaiveAMPModel(model, output_to_fp32=True)
|
|
optimizer = ShardedOptimizer(optimizer, **zero_config)
|
|
else:
|
|
model = ShardedModel(module=model, **zero_config)
|
|
return model, optimizer
|
|
|
|
|
|
__all__ = ['convert_to_zero', 'ShardedModel', 'ShardedOptimizer']
|