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

fix moe bugs in zero optimizer

pull/375/head
Wenwen Qu 2023-08-17 15:31:27 +08:00
parent 3bfaad895a
commit 9ee57e6c8a
1 changed files with 31 additions and 22 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -166,10 +166,6 @@ class HybridZeroOptimizer(BaseOptimizer):
# partition these param groups for data parallel training
# and add buffers to parameter store for future access
for group_id, param_group in enumerate(self.optim.param_groups):
if "moe" in param_group.keys() and param_group["moe"]:
print("true", flush=True)
continue
group_params = param_group["params"]
# add the fp16 params to fp16_param_groups for bookkeeping
@ -180,7 +176,10 @@ class HybridZeroOptimizer(BaseOptimizer):
self.param_group_no_params_ranks.append(no_params_ranks)
self.param_group_has_params.append(self._zero_local_rank not in no_params_ranks)
# store the mapping between param to rank each param should belong to only one rank
# store the mapping between param to rank each param should belong to only one rank.
# we can skip the moe param and do not keep them in _param_store to save memory
# (means we need to deal with moe param in a different way), but it will increase
# complexity and reduce code readablity.
for rank, params in enumerate(params_per_rank):
# check whether any rank is not assigned params.
if len(params) != 0:
@ -267,6 +266,14 @@ class HybridZeroOptimizer(BaseOptimizer):
numel_per_rank = [0 for _ in range(self._zero_world_size)]
self.params_per_rank_id_dict.append([[] for _ in range(self._zero_world_size)])
if "moe" in param_list.keys() and param_list["moe"]:
# just add current params to params_per_rank[_zero_local_rank]
params_per_rank[self._zero_local_rank] = list(param_list["params"])
self.params_per_rank_id_dict[-1][self._zero_local_rank].append(None)
no_params_ranks = list(range(self._zero_world_size))
no_params_ranks.pop(self._zero_world_size)
else:
sorted_params = sorted(param_list, key=lambda x: x.numel(), reverse=True)
for i, param in enumerate(sorted_params):
global_id = str(i)
@ -296,6 +303,7 @@ class HybridZeroOptimizer(BaseOptimizer):
for group_id in range(self.num_param_groups):
param_group = self._fp16_param_groups[group_id]
for param in param_group:
# we should not reduce the param in moe
if param.requires_grad and not is_moe_param(param):
reduce_rank = None
@ -496,6 +504,7 @@ class HybridZeroOptimizer(BaseOptimizer):
if not self._overlap_communication:
for group_id in range(len(self._fp16_param_groups)):
for param in self._fp16_param_groups[group_id]:
# we should not reduce the param in moe
if param.grad is not None and not is_moe_param(param):
self._store_and_try_reduce_grads_by_bucket(param)