wehub-resource-sync
183 lines
6.5 KiB
Python
183 lines
6.5 KiB
Python
# Copyright (c) Microsoft Corporation.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# DeepSpeed Team
|
|
|
|
from collections import defaultdict
|
|
from typing import Any, Dict, List, Set, Tuple, Union, cast
|
|
|
|
import torch
|
|
from torch import nn
|
|
|
|
from .layer import MoE
|
|
|
|
|
|
def has_moe_layers(m: nn.Module) -> Tuple[bool, int]:
|
|
has_moe = False
|
|
num_experts = 0
|
|
|
|
for module in m.modules():
|
|
if isinstance(module, MoE):
|
|
has_moe = True
|
|
num_experts = module.num_experts
|
|
break
|
|
return has_moe, num_experts
|
|
|
|
|
|
def is_moe_param(param: torch.Tensor) -> bool:
|
|
if hasattr(param, "allreduce") and not param.allreduce:
|
|
return True
|
|
return False
|
|
|
|
|
|
def split_params_into_shared_and_expert_params(
|
|
params: List[torch.nn.Parameter]) -> Tuple[List[torch.nn.Parameter], List[torch.nn.Parameter]]:
|
|
shared_params: List[nn.Parameter] = []
|
|
expert_params: List[nn.Parameter] = []
|
|
|
|
for p in params:
|
|
if is_moe_param(p):
|
|
expert_params.append(p)
|
|
else:
|
|
shared_params.append(p)
|
|
return shared_params, expert_params
|
|
|
|
|
|
def split_params_grads_into_shared_and_expert_params(
|
|
group: List[torch.nn.Parameter]) -> Tuple[List[torch.Tensor], List[torch.Tensor]]:
|
|
"""Split grad of parameters into grads of non-expert params
|
|
and grads of expert params. This is useful while computing
|
|
grad-norms for clipping and overflow detection
|
|
|
|
group (List[torch.nn.Parameter]):
|
|
Args:
|
|
The group of parameters to split
|
|
|
|
Returns:
|
|
Tuple[List[torch.Tensor], List[torch.Tensor]]:
|
|
list of gradients for non MoE params, list of gradients of MoE params
|
|
"""
|
|
expert_grads: List[torch.Tensor] = []
|
|
shared_grads: List[torch.Tensor] = []
|
|
|
|
for p in group:
|
|
if p.grad is not None:
|
|
if is_moe_param(p):
|
|
expert_grads.append(p.grad.to(p.dtype))
|
|
else:
|
|
shared_grads.append(p.grad.to(p.dtype))
|
|
return shared_grads, expert_grads
|
|
|
|
|
|
def split_params_into_different_moe_groups_for_optimizer(
|
|
param_groups: Union[Dict[str, Any], Tuple[Dict[str, Any], ...], List[Dict[str, Any]]],
|
|
max_group_size: Union[int, float] = 178956971) -> List[Dict[str, Any]]:
|
|
"""Split parameters into different MoE groups for optimizer
|
|
|
|
Args:
|
|
param_groups (Union[Dict[str, Any], Tuple[Dict[str, Any], ...], List[Dict[str, Any]]])
|
|
The list of parameter groups to split
|
|
|
|
Returns:
|
|
List[Dict[str, Any]]:
|
|
list of MoE/non-MoE groups for optimizer
|
|
"""
|
|
if isinstance(param_groups, tuple):
|
|
param_groups = list(param_groups) # Tuple cannot be modified
|
|
elif isinstance(param_groups, dict):
|
|
param_groups = [param_groups]
|
|
elif not isinstance(param_groups, list):
|
|
raise ValueError(f"Unknown param group type of {type(param_groups)}")
|
|
|
|
# gather all data parallel group names
|
|
data_parallel_group_names: Set[str] = set()
|
|
for param_group in param_groups:
|
|
for param in cast(List[nn.Parameter], param_group["params"]):
|
|
if is_moe_param(param):
|
|
data_parallel_group_names.add(param.group_name)
|
|
|
|
# Create the param MoE groups, leave param assign to next step
|
|
group_moe: Dict[str, Dict[str, Dict[str, Any]]] = defaultdict(lambda: defaultdict(dict))
|
|
for param_group in param_groups:
|
|
for key in data_parallel_group_names:
|
|
group_moe[param_group['name']][key] = {
|
|
**param_group,
|
|
'name': key,
|
|
'moe': True,
|
|
'params': [],
|
|
}
|
|
|
|
# Assign param
|
|
for param_group in param_groups:
|
|
new_params: List[nn.Parameter] = []
|
|
|
|
for param in cast(List[nn.Parameter], param_group['params']):
|
|
if is_moe_param(param):
|
|
group_moe[param_group['name']][param.group_name]['params'].append(param)
|
|
else:
|
|
new_params.append(param)
|
|
param_group['params'] = new_params
|
|
|
|
# Flatten the moe groups
|
|
if max_group_size is not None:
|
|
for moe_group in group_moe.values():
|
|
for param_group in moe_group.values():
|
|
cur_group: List[nn.Parameter] = []
|
|
all_groups: List[List[nn.Parameter]] = []
|
|
size_of_cur_group = 0
|
|
|
|
for param in cast(List[nn.Parameter], param_group['params']):
|
|
if size_of_cur_group + param.numel() <= max_group_size:
|
|
cur_group.append(param)
|
|
size_of_cur_group += param.numel()
|
|
else:
|
|
all_groups.append(cur_group)
|
|
cur_group = [param]
|
|
size_of_cur_group = param.numel()
|
|
|
|
if cur_group:
|
|
all_groups.append(cur_group)
|
|
|
|
for group in all_groups:
|
|
param_groups.append({**param_group, 'params': group})
|
|
else:
|
|
for moe_group in group_moe.values():
|
|
for param_group in moe_group.values():
|
|
param_groups.append(param_group)
|
|
|
|
return param_groups
|
|
|
|
|
|
def is_moe_param_group(param_group):
|
|
return param_group.get('moe', False)
|
|
|
|
|
|
def configure_moe_param_groups(model_parameters: List):
|
|
assert isinstance(model_parameters, list), "model_parameters must be a list"
|
|
|
|
for p in model_parameters:
|
|
# match torch.optim.Optimizer expectations,
|
|
# see: https://github.com/pytorch/pytorch/blob/2ffab6e663b9c6951048b8c8ba82d2cc5ca5c2fc/torch/optim/optimizer.py#L270-L272
|
|
if not isinstance(p, (torch.Tensor, dict)):
|
|
raise TypeError("param argument that would be given to the optimizer should be "
|
|
f"an iterable of Tensors or dicts, but got {type(p)}")
|
|
|
|
# peak at the first element to determine how to proceed
|
|
first = model_parameters[0]
|
|
|
|
# Case 1: model_parameters is a list of torch.nn.Parameter
|
|
# -> need to create moe compatible param groups
|
|
if isinstance(first, torch.nn.Parameter):
|
|
param_group = {'params': model_parameters, 'name': 'dense-params'}
|
|
return split_params_into_different_moe_groups_for_optimizer(param_group)
|
|
|
|
# Case 2: model_parameters is a list of param groups List[dict]
|
|
# -> moe compatible param groups might already exist, if not create them
|
|
elif isinstance(first, dict):
|
|
#there are no moe groups created
|
|
if not any(['moe' in param_group for param_group in model_parameters]):
|
|
return split_params_into_different_moe_groups_for_optimizer(model_parameters)
|
|
else:
|
|
# moe groups exist, nothing to do
|
|
return model_parameters
|