mirror of https://github.com/InternLM/InternLM
59 lines
1.9 KiB
Python
59 lines
1.9 KiB
Python
from typing import Dict, Tuple
|
|
|
|
import torch
|
|
|
|
|
|
def split_params_into_different_groups_for_optimizer(param_groups: Tuple[Dict]) -> Tuple[Dict]:
|
|
"""Split parameters into different groups for optimizer
|
|
Compatiable with muiltiple param groups, each should have a name
|
|
|
|
Args:
|
|
param_groups (Tuple[Dict]):
|
|
The list of parameter groups to split
|
|
|
|
Returns:
|
|
Tuple[Dict]:
|
|
list of fp16/fp32 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)}")
|
|
|
|
# Create fp32 and moe groups and copy origin attribute
|
|
for group_param in param_groups:
|
|
fp32_group = {}
|
|
|
|
# copy attribute for fp32 group
|
|
for ori_key in group_param.keys():
|
|
if ori_key == "name":
|
|
fp32_group["name"] = ori_key + "_fp32"
|
|
else:
|
|
if ori_key == "params":
|
|
fp32_group[ori_key] = []
|
|
else:
|
|
fp32_group[ori_key] = group_param[ori_key]
|
|
|
|
# Assign param
|
|
new_params = []
|
|
for param in group_param["params"]:
|
|
if param.dtype == torch.float32:
|
|
fp32_group["params"].append(param)
|
|
else:
|
|
new_params.append(param)
|
|
|
|
# origin group without fp32
|
|
group_param["params"] = new_params
|
|
# append to origin group
|
|
param_groups.append(fp32_group)
|
|
|
|
return tuple(param_groups)
|
|
|
|
|
|
def create_param_groups(model, weight_decay):
|
|
parameters = {"params": list(model.parameters()), "name": "default", "weight_decay": weight_decay}
|
|
|
|
return split_params_into_different_groups_for_optimizer(parameters)
|