mirror of
https://github.com/hpcaitech/ColossalAI.git
synced 2026-04-12 23:22:42 +00:00
3280869358
* add zero1 (#209) * add zero1 * add test zero1 * update zero stage 1 develop (#212) * Implement naive zero3 (#240) * naive zero3 works well * add zero3 param manager * add TODOs in comments * add gather full param ctx * fix sub module streams * add offload * fix bugs of hook and add unit tests * fix bugs of hook and add unit tests (#252) * add gather full param ctx * fix sub module streams * add offload * fix bugs of hook and add unit tests * polish code and add state dict hook * fix bug * update unit test * refactor reconstructed zero code * clip_grad support zero3 and add unit test * add unit test for Zero3ParameterManager * [WIP] initialize the shard param class * [WIP] Yet another sharded model implementation (#274) * [WIP] initialize the shard param class * [WIP] Yes another implementation of shardModel. Using a better hook method. * torch.concat -> torch.cat * fix test_zero_level_1.py::test_zero_level_1 unitest * remove deepspeed implementation and refactor for the reconstructed zero module * polish zero dp unittests Co-authored-by: ver217 <lhx0217@gmail.com> Co-authored-by: Frank Lee <somerlee.9@gmail.com>
42 lines
1.5 KiB
Python
42 lines
1.5 KiB
Python
from distutils.command.config import config
|
|
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
|
|
|
|
|
|
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:
|
|
assert config['partition_grad'], 'ZeRO Optimizer requires partition_grad to be True'
|
|
model = NaiveAMPModel(model, output_to_fp32=True)
|
|
optimizer = ShardedOptimizer(model.parameters(), *zero_config)
|
|
else:
|
|
model = ShardedModel(module=model, **zero_config)
|
|
return model, optimizer
|
|
|
|
__all__ = ['convert_to_zero', 'ShardedModel', 'ShardedOptimizer']
|