mirror of https://github.com/hpcaitech/ColossalAI
226 lines
9.2 KiB
Python
226 lines
9.2 KiB
Python
from enum import Enum
|
|
from typing import Callable, Dict, Optional, Union
|
|
|
|
import torch
|
|
import torch.distributed as dist
|
|
import torch.nn as nn
|
|
from colossalai.amp.naive_amp.grad_scaler import DynamicGradScaler
|
|
from colossalai.context.parallel_mode import ParallelMode
|
|
from colossalai.core import global_context as gpc
|
|
from colossalai.nn.optimizer import ColossalaiOptimizer
|
|
from colossalai.zero.shard_utils import BaseShardStrategy
|
|
from colossalai.zero.sharded_model import ShardedModelV2
|
|
from colossalai.zero.sharded_model._zero3_utils import cast_tensor_to_fp32
|
|
from torch import Tensor
|
|
from torch.distributed import ProcessGroup
|
|
from torch.nn.parameter import Parameter
|
|
from torch.optim import Optimizer
|
|
from typing import Type, Any
|
|
from ._utils import has_inf_or_nan
|
|
|
|
|
|
class OptimState(Enum):
|
|
SCALED = 1
|
|
UNSCALED = 2
|
|
|
|
|
|
class ShardedOptimizerV2(ColossalaiOptimizer):
|
|
|
|
def __init__(self,
|
|
sharded_model: ShardedModelV2,
|
|
optimizer_class: Type[Optimizer],
|
|
cpu_offload: bool = False,
|
|
initial_scale: float = 2**32,
|
|
min_scale: float = 1,
|
|
growth_factor: float = 2,
|
|
backoff_factor: float = 0.5,
|
|
growth_interval: float = 1000,
|
|
hysteresis: float = 2,
|
|
max_scale: int = 2**32,
|
|
dp_process_group: Optional[ProcessGroup] = None,
|
|
mp_process_group: Optional[ProcessGroup] = None,
|
|
**defaults: Any) -> None:
|
|
"""
|
|
:param sharded_model: A sharded model initialized by class ShardedModelV2. The optimizer will use the
|
|
shard strategy provided by sharded model to shard param fp32 tensors.
|
|
:type sharded_model: sharded_model
|
|
|
|
:param optimizer_class: A class type of Optimizer
|
|
:type optimizer_class: Type[Optimizer]
|
|
|
|
:param cpu_offload: is offloading the optimizer states to CPU.
|
|
:type cpu_offload: bool
|
|
|
|
:param initial_scale: initial scale used by DynamicGradScaler
|
|
:type initial_scale: float
|
|
|
|
:param min_scale: min scale used by DynamicGradScaler
|
|
:type min_scale: float
|
|
|
|
:param growth_factor: growth_factor used by DynamicGradScaler
|
|
:type growth_factor: float
|
|
|
|
:param backoff_factor: backoff_factor used by DynamicGradScaler
|
|
:type backoff_factor: float
|
|
|
|
:param growth_interval: growth_interval used by DynamicGradScaler
|
|
:type growth_interval: float
|
|
|
|
:param hysteresis: hysteresis used by DynamicGradScaler
|
|
:type hysteresis: float
|
|
|
|
:param max_scale: max_scale used by DynamicGradScaler
|
|
:type max_scale: float
|
|
|
|
:param dp_process_group: data paralle process group
|
|
:type dp_process_group: Optional[ProcessGroup]
|
|
|
|
:param mp_process_group: model paralle process group
|
|
:type mp_process_group: Optional[ProcessGroup]
|
|
|
|
:**defaults: any trailing arguments, which are forwarded to the local optimizer.
|
|
:type defaults: dict()
|
|
"""
|
|
assert isinstance(sharded_model, ShardedModelV2), 'model must be wrapped with ShardedModel'
|
|
|
|
self._optim_defaults = defaults
|
|
# initialize the M, V as zeros tensors and initialize param fp32 from sharded_model.parameters()
|
|
|
|
self.optimizer = optimizer_class(sharded_model.parameters(), **self._optim_defaults)
|
|
|
|
super().__init__(self.optimizer)
|
|
self.shard_strategy = sharded_model.shard_strategy
|
|
self.model: ShardedModelV2 = sharded_model
|
|
if cpu_offload and not sharded_model.cpu_offload:
|
|
raise RuntimeError(
|
|
f"ShardedOptimizerV2 using cpu_offload, but the sharded_model used to initialize it dose not use cpu_offload"
|
|
)
|
|
self.device = torch.cuda.current_device() if not cpu_offload else torch.device('cpu')
|
|
self.optim_state: OptimState = OptimState.UNSCALED
|
|
self.dp_process_group = dp_process_group or gpc.get_group(ParallelMode.DATA)
|
|
self.mp_process_group = mp_process_group or gpc.get_group(ParallelMode.MODEL)
|
|
# Grad scaler
|
|
self.grad_scaler = DynamicGradScaler(initial_scale=initial_scale,
|
|
min_scale=min_scale,
|
|
growth_factor=growth_factor,
|
|
backoff_factor=backoff_factor,
|
|
growth_interval=growth_interval,
|
|
hysteresis=hysteresis,
|
|
max_scale=max_scale)
|
|
self._found_overflow: Tensor = torch.FloatTensor([0]).to(torch.cuda.current_device())
|
|
|
|
# Store fp32 param shards
|
|
self.master_params: Dict[Parameter, Tensor] = {}
|
|
|
|
for group in self.optimizer.param_groups:
|
|
for p in group['params']:
|
|
assert hasattr(p, 'col_attr'), 'The parameter must be wrapped with ShardedParam'
|
|
is_param_sharded = p.col_attr.data.is_sharded
|
|
if not is_param_sharded:
|
|
# TODO (ver217): we may not use shard / gather here
|
|
# Param is no sharded, which means we use ZeRO-2 here
|
|
# As we only store param shard, we shard it here
|
|
self.shard_strategy.shard([p.col_attr.data])
|
|
self.master_params[p] = cast_tensor_to_fp32(p.col_attr.data.payload).to(self.device)
|
|
if not is_param_sharded:
|
|
# In this branch, there's no need to shard param
|
|
# So we gather here
|
|
self.shard_strategy.gather([p.col_attr.data])
|
|
|
|
def step(self, *args, **kwargs):
|
|
# unscale grads if scaled
|
|
if self.optim_state == OptimState.SCALED:
|
|
self._unscale_grads()
|
|
|
|
found_inf = self._check_overflow()
|
|
self.grad_scaler.update(found_inf)
|
|
|
|
if found_inf:
|
|
self.zero_grad()
|
|
return
|
|
|
|
# assign master param pointers to p.data.
|
|
# We will not trigger data copy here.
|
|
for group in self.optim.param_groups:
|
|
for p in group['params']:
|
|
p.data = self.master_params[p]
|
|
# Now p.data is sharded
|
|
# So optimizer states are sharded naturally
|
|
|
|
ret = self.optim.step(*args, **kwargs)
|
|
|
|
# Copy master param data (fp32) to payload of col_attr (fp16)
|
|
# TODO() improve efficiency by gathering tensors into a chunk and transfering
|
|
# a chunk.
|
|
for group in self.optim.param_groups:
|
|
for p in group['params']:
|
|
is_param_sharded = p.col_attr.data.is_sharded
|
|
if not is_param_sharded:
|
|
# We use ZeRO-2 here
|
|
# The `p.col_attr.data` saves full fp16 param
|
|
# But we only have updated fp32 param shard here
|
|
# So we first shard full fp16 param and copy fp32 param shard to it
|
|
# Then we will gather them
|
|
self.shard_strategy.shard([p.col_attr.data])
|
|
# We have to use `copy_payload` instead of `reset_payload`
|
|
# Since p.data is fp32 and p.col_attr.data is fp16
|
|
|
|
# TODO() optimize this line CPU (fp32) -> GPU (fp16)
|
|
p.col_attr.data.copy_payload(p.data)
|
|
|
|
if not is_param_sharded:
|
|
# We gather full fp16 param here
|
|
self.shard_strategy.gather([p.col_attr.data])
|
|
p.data = p.col_attr.data.payload
|
|
return ret
|
|
|
|
def backward(self, loss: Tensor) -> None:
|
|
loss = self.loss_scale * loss
|
|
self.optim_state = OptimState.SCALED
|
|
self.model.backward(loss)
|
|
|
|
def backward_by_grad(self, tensor: Tensor, grad: Tensor) -> None:
|
|
self.model.backward_by_grad(tensor, grad)
|
|
|
|
def clip_grad_norm(self, model: nn.Module, max_norm: float):
|
|
if self.optim_state == OptimState.SCALED:
|
|
self._unscale_grads()
|
|
return super().clip_grad_norm(model, max_norm)
|
|
|
|
@property
|
|
def loss_scale(self):
|
|
return self.grad_scaler.scale.item()
|
|
|
|
def _check_overflow(self):
|
|
# clear previous overflow record
|
|
self._found_overflow.fill_(0.0)
|
|
|
|
# check for overflow
|
|
for group in self.optim.param_groups:
|
|
for p in group['params']:
|
|
if has_inf_or_nan(p.grad):
|
|
self._found_overflow.fill_(1.0)
|
|
break
|
|
|
|
# all-reduce across dp group
|
|
dist.all_reduce(self._found_overflow, op=dist.ReduceOp.MAX, group=self.dp_process_group)
|
|
|
|
# all-reduce over model parallel group
|
|
dist.all_reduce(self._found_overflow, op=dist.ReduceOp.MAX, group=self.mp_process_group)
|
|
|
|
return self._found_overflow.item() > 0
|
|
|
|
def _unscale_grads(self):
|
|
assert self.optim_state == OptimState.SCALED
|
|
for group in self.optim.param_groups:
|
|
for p in group['params']:
|
|
if p.grad is not None:
|
|
p.grad.data.div_(self.loss_scale)
|
|
self.optim_state = OptimState.UNSCALED
|
|
|
|
def zero_grad(self, *args, **kwargs):
|
|
# We must set grad to None
|
|
# Because we will judge whether local grad accumulation
|
|
# is enabled by wheter grad is None
|
|
self.optim.zero_grad(set_to_none=True)
|