refactor code for split group

pull/182/head
Wenwen Qu 2023-09-22 13:00:26 +08:00
parent 17bc5f562b
commit 3607548265
1 changed files with 29 additions and 24 deletions

View File

@ -1,7 +1,6 @@
from typing import Dict, Tuple
from internlm.core.context.parallel_context import global_context as gpc
from internlm.model.utils import is_gate_param, is_moe_param, is_norm_param
def split_params_into_different_groups_for_optimizer(param_groups: Tuple[Dict]) -> Tuple[Dict]:
@ -27,6 +26,19 @@ def split_params_into_different_groups_for_optimizer(param_groups: Tuple[Dict])
>>> ...,
>>> )
"""
def _get_group(param):
group_keys = ["is_expert", "is_gate", "is_norm"]
for i, key in enumerate(group_keys):
if hasattr(param, key) and getattr(param, key):
# experts param should return its group name
if i == 0:
return param.group_name
else:
return key[3:]
# TODO: deal with fp32 group
return None
if isinstance(param_groups, tuple):
param_groups = list(param_groups) # Tuple cannot be modified
elif isinstance(param_groups, dict):
@ -34,46 +46,39 @@ def split_params_into_different_groups_for_optimizer(param_groups: Tuple[Dict])
elif not isinstance(param_groups, list):
raise ValueError(f"Unknown param group type of {type(param_groups)}")
new_groups = {}
new_groups = []
for pgroup in param_groups:
new_groups[pgroup["name"]] = {}
current_groups = {}
# create new groups for gate and norm
for key in ["gate", "norm"]:
new_groups[pgroup["name"]][key] = {}
new_groups[pgroup["name"]][key]["name"] = key
new_groups[pgroup["name"]][key][key] = True
current_groups[key] = {"name": key, key: True, "params": []}
# create moe groups
for key in gpc.expert_parallel_group_names:
new_groups[pgroup["name"]][key] = {}
new_groups[pgroup["name"]][key]["name"] = key
new_groups[pgroup["name"]][key]["moe"] = True
current_groups[key] = {"name": key, "moe": True, "params": []}
# copy attribute from origin group
for ori_key in pgroup.keys():
for key in new_groups[pgroup["name"]].keys():
if ori_key != "name":
if ori_key == "params":
new_groups[pgroup["name"]][key][ori_key] = []
else:
new_groups[pgroup["name"]][key][ori_key] = pgroup[ori_key]
if ori_key not in ("name", "params"):
for _, group in current_groups.items():
group[ori_key] = pgroup[ori_key]
# Assign param
origin_params = []
for param in pgroup["params"]:
if is_moe_param(param):
new_groups[pgroup["name"]][param.group_name]["params"].append(param)
elif is_norm_param(param):
new_groups[pgroup["name"]]["norm"]["params"].append(param)
elif is_gate_param(param):
new_groups[pgroup["name"]]["gate"]["params"].append(param)
group = _get_group(param)
if group is not None:
current_groups[group]["params"].append(param)
else:
origin_params.append(param)
pgroup["params"] = origin_params
for _, v in new_groups.items():
for _, v1 in v.items():
param_groups.append(v1)
new_groups.append(current_groups)
for g in new_groups:
for _, v in g.items():
param_groups.append(v)
return tuple(param_groups)