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.
111 lines
4.9 KiB
111 lines
4.9 KiB
import torch
|
|
import torch.nn as nn
|
|
from torch.optim import Optimizer
|
|
from colossalai.amp.naive_amp import NaiveAMPModel
|
|
from colossalai.utils import is_no_pp_or_last_stage
|
|
from colossalai.core import global_context as gpc
|
|
from colossalai.context.parallel_mode import ParallelMode
|
|
|
|
from .zero_redundancy_optimizer_level_2 import ZeroRedundancyOptimizer_Level_2
|
|
from .zero_redundancy_optimizer_level_3 import ZeroRedundancyOptimizer_Level_3
|
|
|
|
|
|
def convert_to_zero(model: nn.Module,
|
|
optimizer: Optimizer,
|
|
level: int,
|
|
zero_config):
|
|
"""
|
|
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
|
|
"""
|
|
import deepspeed
|
|
assert level == 2 or level == 3, 'Only ZERO Optimizer Level 2 and 3 are provided'
|
|
model = NaiveAMPModel(model, output_to_fp32=False)
|
|
|
|
if level == 2:
|
|
optimizer = ZeroRedundancyOptimizer_Level_2(init_optimizer=optimizer, **zero_config)
|
|
else:
|
|
optimizer = ZeroRedundancyOptimizer_Level_3(init_optimizer=optimizer, module=model, **zero_config)
|
|
return model, optimizer
|
|
|
|
|
|
def zero3_model_context(dtype=torch.half):
|
|
"""A context to enable massive model construction for training with
|
|
ZeRO-3. Models are automatically partitioned (or, sharded) across the
|
|
system and converted to half precision. Note that the config of ZeRO-3 will be loaded automatically from `gpc.config`.
|
|
|
|
Args:
|
|
dtype (``dtype``, optional): Can be used to change the data type of the parameters.
|
|
Supported options are ``torch.half`` and ``torch.float``. Defaults to ``torch.half``
|
|
|
|
This context accelerates model initialization and enables models that
|
|
are too large to allocate in their entirety in CPU memory. It has the
|
|
following effects:
|
|
|
|
#. allocates tensors to either GPU or CPU memory or NVMe
|
|
#. converts floating point tensors to half precision
|
|
#. immediately partitions tensors among the group of data-parallel devices
|
|
#. (*optional*) replaces ``torch.nn.functional.linear`` with a more
|
|
memory-efficient implementation
|
|
|
|
These modifications allow for models that exceed the size of local CPU/GPU
|
|
memory/NVMe, but fit within the total NVMe capacity (*i.e.*, aggregate CPU
|
|
or GPU memory or NVMe) across all nodes. Consider initializing a model with one
|
|
trillion parameters, whose weights occupy two terabytes (TB) in half
|
|
precision. The initial CPU allocation in full precision requires 4TB of
|
|
memory *per process*, and so a system with 8 GPUs per node would need 32TB of
|
|
CPU memory due to data-parallel redundancies. Instead, by immediately
|
|
partitioning tensors we remove the redundancies. The result is that
|
|
regardless of the number of GPUs, we still only require the original 4TB. This
|
|
allows for a linear increase in model size with the aggregate system memory.
|
|
For example, if a node has 1TB of memory and 8 GPUs, we could fit a trillion
|
|
parameter model with 4 nodes and 32 GPUs.
|
|
|
|
Important: If the fp16 weights of the model can't fit onto a single GPU memory
|
|
this feature must be used.
|
|
|
|
Examples
|
|
--------
|
|
|
|
#. Allocate a model and partition it among all processes:
|
|
|
|
.. code-block:: python
|
|
|
|
with zero3_model_context():
|
|
model = MyLargeModel()
|
|
|
|
"""
|
|
assert dtype == torch.half or dtype == torch.float, f'Invalid dtype, except torch.half or torch.float, got {dtype}'
|
|
import deepspeed
|
|
ds_config = {
|
|
"train_micro_batch_size_per_gpu": 1,
|
|
"gradient_accumulation_steps": 1,
|
|
"zero_optimization": {
|
|
"offload_param": getattr(gpc.config.zero, 'offload_param_config', None),
|
|
"offload_optimizer": getattr(gpc.config.zero, 'offload_optimizer_config'),
|
|
},
|
|
"aio": getattr(gpc.config.zero, 'aio_config', None)
|
|
}
|
|
remote_device = getattr(ds_config['zero_optimization']['offload_param'], 'device', None)
|
|
pin_memory = getattr(ds_config['zero_optimization']['offload_param'], 'pin_memory', False)
|
|
return deepspeed.zero.Init(data_parallel_group=gpc.get_group(ParallelMode.DATA),
|
|
remote_device=remote_device,
|
|
config_dict_or_path=ds_config,
|
|
pin_memory=pin_memory,
|
|
dtype=dtype)
|
|
|
|
|
|
__all__ = ['convert_to_zero', 'ZeroRedundancyOptimizer_Level_2',
|
|
'ZeroRedundancyOptimizer_Level_3', 'zero3_model_context']
|