ColossalAI/colossalai/zero/sharded_optim/sharded_optim_v2.py

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)