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

change the order of dp and sp all-reduce

pull/443/head
yingtongxiong 2023-10-25 13:27:47 +08:00
parent 1655a90f34
commit 1bc3c33b75
1 changed files with 2 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -318,15 +318,14 @@ class HybridZeroOptimizer(BaseOptimizer):
if self.skip_grad_reduce is False: if self.skip_grad_reduce is False:
reduction_sp_func() reduction_sp_func()
accum_grad_obj.register_hook(reduce_grad_hook)
# if sequence_parallel is True, the grad of norm should be all-reduce across the tp process group # if sequence_parallel is True, the grad of norm should be all-reduce across the tp process group
if gpc.config.parallel.sequence_parallel is True: if gpc.config.parallel.sequence_parallel is True:
if hasattr(param, IS_SEQUENCE_PARALLEL) and getattr(param, IS_SEQUENCE_PARALLEL) is True: if hasattr(param, IS_SEQUENCE_PARALLEL) and getattr(param, IS_SEQUENCE_PARALLEL) is True:
accum_grad_obj_sp = get_grad_accumulate_object(param) accum_grad_obj_sp = get_grad_accumulate_object(param)
accum_grad_obj_sp.register_hook(reduce_grad_hook_sp) accum_grad_obj_sp.register_hook(reduce_grad_hook_sp)
accum_grad_obj.register_hook(reduce_grad_hook)
_define_and_attach(param, reduce_rank) _define_and_attach(param, reduce_rank)
def belongs_to_current_rank(self, param) -> bool: def belongs_to_current_rank(self, param) -> bool: