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

fix group_norms computing in hybrid_zero_optim

pull/182/head
Wenwen Qu 2023-08-31 18:46:13 +08:00
parent 2ad5f512b5
commit 7ca5da27e8
1 changed files with 4 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
internlm/solver/optimizer/hybrid_zero_optim.py
View File

@ -574,6 +574,9 @@ class HybridZeroOptimizer(BaseOptimizer):
# compute norm for gradients in the before bucket # compute norm for gradients in the before bucket
groups_norms = [] groups_norms = []
for group_id in range(self.num_param_groups): for group_id in range(self.num_param_groups):
if self._is_moe_group(self.optim.param_groups[group_id]):
groups_norms.append([])
else:
groups_norms.append(self._compute_norm_with_stage(group_id=group_id)) groups_norms.append(self._compute_norm_with_stage(group_id=group_id))
# clear reduced grads # clear reduced grads