github
/
InternLM
mirror of https://github.com/InternLM/InternLM
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

fix bugs with split group

pull/182/head
Wenwen Qu 2023-09-22 16:05:21 +08:00
parent 1fdc7107b4
commit f6cadcafa2
1 changed files with 1 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
internlm/train/utils.py
View File

@ -8,7 +8,6 @@ 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]: def split_params_into_different_groups_for_optimizer(param_groups: Tuple[Dict]) -> Tuple[Dict]:
"""Split parameters into different MoE groups for optimizer """Split parameters into different MoE groups for optimizer
Compatiable with muiltiple param groups, each should have a name
Args: Args:
param_groups (Tuple[Dict]): The list of parameter groups to split param_groups (Tuple[Dict]): The list of parameter groups to split
@ -41,7 +40,7 @@ def split_params_into_different_groups_for_optimizer(param_groups: Tuple[Dict])
new_groups = {} new_groups = {}
new_groups["fp32"] = {"name": "fp32", "params": []} new_groups["fp32"] = {"name": "fp32", "params": []}
for key in ["gate", "norm"]: for key in ["gate", "norm"]:
new_groups[key] = {"name": key, "sync_tp": True, "params": []} new_groups[key] = {"name": key, key: True, "params": []}
for key in gpc.expert_parallel_group_names: for key in gpc.expert_parallel_group_names:
new_groups[key] = {"name": key, "moe": True, "params": []} new_groups[key] = {"name": key, "moe": True, "params": []}