import math from typing import Dict import torch import torch.distributed as dist from colossalai.interface.optimizer import DistributedOptim from colossalai.shardformer.layer._operation import _gather, _split from colossalai.tensor.d_tensor import get_sharding_spec, is_distributed_tensor # DistributedAdaFactor (with Tensor parallel and Zero stage 2) __all__ = ["DistributedAdaFactor"] class DistributedAdaFactor(DistributedOptim): def __init__( self, params, lr=None, eps=(1e-30, 1e-3), clip_threshold=1.0, decay_rate=-0.8, beta1=None, weight_decay=0.0, scale_parameter=True, relative_step=True, warmup_init=False, ): lr = None if lr is not None and relative_step: raise ValueError("Cannot combine manual `lr` and `relative_step=True` options") if warmup_init and not relative_step: raise ValueError("`warmup_init=True` requires `relative_step=True`") defaults = { "lr": lr, "eps": eps, "clip_threshold": clip_threshold, "decay_rate": decay_rate, "beta1": beta1, "weight_decay": weight_decay, "scale_parameter": scale_parameter, "relative_step": relative_step, "warmup_init": warmup_init, } self.tp_size = 1 self.tp_group = None self.dp_size = 1 self.dp_group = None self.shard_to_working_param = None # Dict{id:shape}, sample {id(param): torch.tensor} self.use_zero = True self.param_is_dtensor_dict = {} # {id(p): True/False} self.grad_shape_dict = {} # {id(p): master param shape} self.factored_dict = {} # {id(p): True/False} self.use_first_moment_dict = {} # {id(p): True/False} self.shard_spec_dict = {} # {id(p): ShardSpec} super().__init__(params, defaults) def setup_distributed( self, tp_group: dist.ProcessGroup = None, dp_group: dist.ProcessGroup = None, shard_to_working_param: Dict = {}, padding_map=None, use_zero: bool = True, ) -> None: """Setup process groups for TP and ZeRO 2. Inject features to the Optimizer Args: tp_group: The devices group for tensor parallel; dp_group: The devices group for data parallel; shard_to_working_param (Dict): ZeRO 2 feeds the optimizer a sharded param view as grads are sharded. This maps from id(view) to working params used in forward & backward. padding_map: An empty interface placeholder; use_zero: Whether or not to use zero; """ self.tp_group = tp_group # "Expected row process group" self.dp_group = dp_group if self.tp_group is not None: self.tp_size = dist.get_world_size(self.tp_group) if self.dp_group is not None: self.dp_size = dist.get_world_size(self.dp_group) self.use_zero = use_zero self.shard_to_working_param = shard_to_working_param if shard_to_working_param is not None else {} # grad is None, cause we dont setup now for group in self.param_groups: for p in group["params"]: self.shard_to_working_param[id(p)] = self.shard_to_working_param.get( id(p), p ) # If not ZeRO, working param is master param self.param_is_dtensor_dict[id(p)] = is_distributed_tensor(self.shard_to_working_param[id(p)]) self.grad_shape_dict[id(p)] = self.shard_to_working_param.get(id(p)).shape self.factored_dict[id(p)], self.use_first_moment_dict[id(p)] = self._get_options( group, self.grad_shape_dict[id(p)] ) if self.param_is_dtensor_dict[id(p)]: self.shard_spec_dict[id(p)] = get_sharding_spec(self.shard_to_working_param[id(p)]) else: self.shard_spec_dict[id(p)] = None @staticmethod def _get_lr(param_group, param_state): rel_step_sz = param_group["lr"] if param_group["relative_step"]: min_step = 1e-6 * param_state["step"] if param_group["warmup_init"] else 1e-2 rel_step_sz = min(min_step, 1.0 / math.sqrt(param_state["step"])) param_scale = 1.0 if param_group["scale_parameter"]: param_scale = max(param_group["eps"][1], param_state["RMS"]) return param_scale * rel_step_sz @staticmethod def _get_options(param_group, param_shape): """ Determines whether the current param is factored Args: param_group : param group param_shape : Original Shape of param """ factored = len(param_shape) >= 2 use_first_moment = param_group["beta1"] is not None return factored, use_first_moment @staticmethod def _rms(tensor, param_is_dtensor, use_zero, tp_size, dp_size, tp_group, dp_group): tensor_sum = tensor.pow(2).sum() num_of_element = tensor.numel() if param_is_dtensor: # reduce tensor_sum from tp_group dist.all_reduce(tensor_sum, group=tp_group) num_of_element = num_of_element * tp_size if use_zero: dist.all_reduce(tensor_sum, group=dp_group) num_of_element = num_of_element * dp_size rms = (tensor_sum / num_of_element).sqrt() return rms @staticmethod def _approx_sq_grad(exp_avg_sq_row, exp_avg_sq_col): r_factor = (exp_avg_sq_row / exp_avg_sq_row.mean(dim=-1, keepdim=True)).rsqrt_().unsqueeze(-1) c_factor = exp_avg_sq_col.unsqueeze(-2).rsqrt() return torch.mul(r_factor, c_factor) # approx_sq_grad for row parallel weight @staticmethod def _approx_sq_grad_row_parallel(exp_avg_sq_row, exp_avg_sq_col, sq_row_meam): # row_meam = sq_row_meam r_factor = (exp_avg_sq_row / sq_row_meam).rsqrt_().unsqueeze(-1) c_factor = exp_avg_sq_col.unsqueeze(-2).rsqrt() return torch.mul(r_factor, c_factor) def _col_parallel_factor(self, update, grad, state, grad_shape, beta2t): if grad_shape[0] % self.dp_size != 0: # gather update[flatten] along dp group then reshape to [H, W/tp] update = _gather(input_=update, dim=-1, process_group=self.dp_group) update_reshape = update.view(-1, grad_shape[1]) # gather grad[flatten] along dp group then reshape to [H, W/tp] grad = _gather(input_=grad, dim=-1, process_group=self.dp_group) grad_reshape = grad.view(-1, grad_shape[1]) exp_avg_sq_row = state["exp_avg_sq_row"] # [H] exp_avg_sq_col = state["exp_avg_sq_col"] # [W/tp] exp_avg_sq_row.mul_(beta2t).add_(update_reshape.mean(dim=-1), alpha=(1.0 - beta2t)) exp_avg_sq_col.mul_(beta2t).add_(update_reshape.mean(dim=-2), alpha=(1.0 - beta2t)) update_reshape = self._approx_sq_grad(exp_avg_sq_row, exp_avg_sq_col) update_reshape.mul_(grad_reshape) else: update_reshape = update.view(-1, grad_shape[1]) grad_reshape = grad.view(-1, grad_shape[1]) exp_avg_sq_row = state["exp_avg_sq_row"] # [H/dp] exp_avg_sq_col = state["exp_avg_sq_col"] # [W/tp] exp_avg_sq_row.mul_(beta2t).add_(update_reshape.mean(dim=-1), alpha=(1.0 - beta2t)) exp_avg_sq_col.mul_(beta2t).add_(update_reshape.mean(dim=-2), alpha=(1.0 - beta2t)) dist.all_reduce(exp_avg_sq_row, group=self.tp_group) exp_avg_sq_row.div_(self.tp_size) update_reshape = self._approx_sq_grad(exp_avg_sq_row, exp_avg_sq_col) update_reshape.mul_(grad_reshape) if self.use_zero: update = update_reshape.view(-1) else: update = update_reshape return update def _row_parallel_factor(self, update, grad, state, grad_shape, beta2t): if grad_shape[0] % self.dp_size != 0: # gather update[flatten] along dp group then reshape to [H/tp, W] update = _gather(input_=update, dim=-1, process_group=self.dp_group) # view update to origin[tp] shape update_reshape = update.view(-1, grad_shape[1]) # gather grad[flatten] along dp group then reshape to [H/tp, W] grad = _gather(input_=grad, dim=-1, process_group=self.dp_group) grad_reshape = grad.view(-1, grad_shape[1]) exp_avg_sq_row = state["exp_avg_sq_row"] # [H/tp] exp_avg_sq_col = state["exp_avg_sq_col"] # [W] exp_avg_sq_row.mul_(beta2t).add_(update_reshape.mean(dim=-1), alpha=(1.0 - beta2t)) exp_avg_sq_col.mul_(beta2t).add_(update_reshape.mean(dim=-2), alpha=(1.0 - beta2t)) # reduce col dist.all_reduce(exp_avg_sq_col, group=self.tp_group) exp_avg_sq_col.div_(self.tp_size) update_reshape = self._approx_sq_grad(exp_avg_sq_row, exp_avg_sq_col) update_reshape.mul_(grad_reshape) if self.use_zero: update = _split(input_=update_reshape.view(-1), dim=-1, process_group=self.dp_group) else: update = update_reshape else: update_reshape = update.view(-1, grad_shape[1]) grad_reshape = grad.view(-1, grad_shape[1]) exp_avg_sq_row = state["exp_avg_sq_row"] # [H/dp/tp] exp_avg_sq_col = state["exp_avg_sq_col"] # [W] exp_avg_sq_row.mul_(beta2t).add_(update_reshape.mean(dim=-1), alpha=(1.0 - beta2t)) exp_avg_sq_col.mul_(beta2t).add_(update_reshape.mean(dim=-2), alpha=(1.0 - beta2t)) # reduce col dist.all_reduce(exp_avg_sq_col, group=self.tp_group) exp_avg_sq_col.div_(self.tp_size) # gather row exp_avg_sq_row_gather = _gather(input_=exp_avg_sq_row, dim=-1, process_group=self.tp_group) sq_row_meam = exp_avg_sq_row_gather.mean(dim=-1, keepdim=True) update_reshape = self._approx_sq_grad_row_parallel(exp_avg_sq_row, exp_avg_sq_col, sq_row_meam) update_reshape.mul_(grad_reshape) if self.use_zero: update = update_reshape.view(-1) else: update = update_reshape return update def _base_factor(self, update, grad, state, grad_shape, beta2t): if self.use_zero: # only zero if grad_shape[0] % self.dp_size != 0: # view update to origin shape update.view(grad_shape[0]//self.data_parallel_size , grad_shape[1]) # row mean no change # col mean need reduce and div # gather update[flatten] along dp group then reshape to [H, W] update = _gather(input_=update, dim=-1, process_group=self.dp_group) # view update to origin[tp] shape update_reshape = update.view(-1, grad_shape[1]) # gather grad[flatten] along dp group then reshape to [H, W] grad = _gather(input_=grad, dim=-1, process_group=self.dp_group) grad_reshape = grad.view(-1, grad_shape[1]) exp_avg_sq_row = state["exp_avg_sq_row"] # [H/dp] exp_avg_sq_col = state["exp_avg_sq_col"] # [W] exp_avg_sq_row.mul_(beta2t).add_(update_reshape.mean(dim=-1), alpha=(1.0 - beta2t)) exp_avg_sq_col.mul_(beta2t).add_(update_reshape.mean(dim=-2), alpha=(1.0 - beta2t)) # reduce col dist.all_reduce(exp_avg_sq_col, group=self.tp_group) exp_avg_sq_col.div_(self.tp_size) update_reshape = self._approx_sq_grad(exp_avg_sq_row, exp_avg_sq_col) update_reshape.mul_(grad_reshape) update = _split(input_=update_reshape.view(-1), dim=-1, process_group=self.dp_group) else: # no residual row # view update to origin[tp] shape update_reshape = update.view(-1, grad_shape[1]) # [H/dp, W] grad_reshape = grad.view(-1, grad_shape[1]) # [H/dp, W] exp_avg_sq_row = state["exp_avg_sq_row"] # [H/tp] exp_avg_sq_col = state["exp_avg_sq_col"] # [W] exp_avg_sq_row.mul_(beta2t).add_(update_reshape.mean(dim=-1), alpha=(1.0 - beta2t)) exp_avg_sq_col.mul_(beta2t).add_(update_reshape.mean(dim=-2), alpha=(1.0 - beta2t)) # reduce col dist.all_reduce(exp_avg_sq_col, group=self.tp_group) exp_avg_sq_col.div_(self.tp_size) update_reshape = self._approx_sq_grad(exp_avg_sq_row, exp_avg_sq_col) update_reshape.mul_(grad_reshape) update = update_reshape.view(-1) else: # base factor; no tp, no dp exp_avg_sq_row = state["exp_avg_sq_row"] exp_avg_sq_col = state["exp_avg_sq_col"] # Exponential average of row indexes exp_avg_sq_row.mul_(beta2t).add_(update.mean(dim=-1), alpha=(1.0 - beta2t)) # Exponential average of columns indexes exp_avg_sq_col.mul_(beta2t).add_(update.mean(dim=-2), alpha=(1.0 - beta2t)) # Approximation of exponential moving average of square of gradient update = self._approx_sq_grad(exp_avg_sq_row, exp_avg_sq_col) update.mul_(grad) return update @torch.no_grad() def step(self, closure=None): """ Performs a single optimization steps Arguments: closure (callable, optional): A closure that reevaluates the model and returns the loss. """ loss = None if closure is not None: loss = closure() """ param_groups: Dict { "params":[weight, bias] "lr" "eps" "clip_threshold" "decay_rate" "beta1" "weight_decay" "scale_parameter" "relative_step" "warmup_init" } """ for group in self.param_groups: # update weight & bias for p in group["params"]: if p.grad is None: continue grad = p.grad if grad.is_sparse: raise RuntimeError("Adafactor does not support sparse gradients.") state = self.state[p] grad_shape = self.grad_shape_dict[id(p)] param_is_dtensor = self.param_is_dtensor_dict[id(p)] if param_is_dtensor: grad_shape = self.shard_to_working_param.get(id(p)).shape # tp shape (2 dim) factored, use_first_moment = self.factored_dict[id(p)], self.use_first_moment_dict[id(p)] shard_spec = self.shard_spec_dict[id(p)] if len(state) == 0: state["step"] = 0 if use_first_moment: # Exponential moving average of gradient values state["exp_avg"] = torch.zeros_like(p) if factored: if param_is_dtensor: if shard_spec.sharding_sequence[0] == "R": # Col Parallel if grad_shape[0] % self.dp_size != 0: state["exp_avg_sq_row"] = torch.zeros( grad_shape[0], device=p.device, dtype=p.dtype ) # [H] else: state["exp_avg_sq_row"] = torch.zeros( grad_shape[0] // self.dp_size, device=p.device, dtype=p.dtype ) # [H/dp] state["exp_avg_sq_col"] = torch.zeros( grad_shape[1], device=p.device, dtype=p.dtype ) # [W/TP] if shard_spec.sharding_sequence[-1] == "R": # Row Parallel # Row indivisible shape situation if grad_shape[0] % self.dp_size != 0: state["exp_avg_sq_row"] = torch.zeros( grad_shape[0], device=p.device, dtype=p.dtype ) # [H/tp] else: state["exp_avg_sq_row"] = torch.zeros( grad_shape[0] // self.dp_size, device=p.device, dtype=p.dtype ) # [H/dp/tp] state["exp_avg_sq_col"] = torch.zeros( grad_shape[1], device=p.device, dtype=p.dtype ) # [W] else: if self.use_zero: if grad_shape[0] % self.dp_size != 0: # save all exp_avg_sq_row [H] state["exp_avg_sq_row"] = torch.zeros( grad_shape[0], device=grad.device, dtype=p.dtype ) else: # exp_avg_sq_row [H // dp] state["exp_avg_sq_row"] = torch.zeros( grad_shape[0] // self.dp_size, device=grad.device, dtype=p.dtype ) else: # exp_avg_sq_row [H] state["exp_avg_sq_row"] = torch.zeros(grad_shape[0], device=grad.device, dtype=p.dtype) # exp_avg_sq_col alaways [W] state["exp_avg_sq_col"] = torch.zeros(grad_shape[1], device=grad.device, dtype=p.dtype) else: state["exp_avg_sq"] = torch.zeros_like(p) state["RMS"] = 0 else: if use_first_moment: state["exp_avg"] = state["exp_avg"] if factored: state["exp_avg_sq_row"] = state["exp_avg_sq_row"] state["exp_avg_sq_col"] = state["exp_avg_sq_col"] else: state["exp_avg_sq"] = state["exp_avg_sq"] state["step"] += 1 lr = self._get_lr(group, state) beta2t = 1.0 - math.pow(state["step"], group["decay_rate"]) update = (grad**2) + group["eps"][0] if factored: if param_is_dtensor: # ============================== # First Dim is R, Last Dim is S{} means split dim -1 ---> # Coloum Parallel ---> sq_row need Do (col) Reduce # ============================== if shard_spec.sharding_sequence[0] == "R": update = self._col_parallel_factor(update, grad, state, grad_shape, beta2t) # ============================== # Last Dim is R, First Dim is S{} means split dim 0 ---> # Row Parallel ---> sq_col need Do (row) Reduce # ============================== elif shard_spec.sharding_sequence[-1] == "R": update = self._row_parallel_factor(update, grad, state, grad_shape, beta2t) else: update = self._base_factor(update, grad, state, grad_shape, beta2t) else: exp_avg_sq = state["exp_avg_sq"] exp_avg_sq.mul_(beta2t).add_(update, alpha=(1.0 - beta2t)) update = exp_avg_sq.rsqrt().mul_(grad) # # (Line No.8) RMS rms = self._rms( update, param_is_dtensor, self.use_zero, self.tp_size, self.dp_size, self.tp_group, self.dp_group, ) update.div_((rms / group["clip_threshold"]).clamp_(min=1.0)) update.mul_(lr) if use_first_moment: exp_avg = state["exp_avg"] exp_avg.mul_(group["beta1"]).add_(update, alpha=(1 - group["beta1"])) update = exp_avg if group["weight_decay"] != 0: p.add_(p, alpha=(-group["weight_decay"] * lr)) p.add_(-update) return loss