Fix a convergence issues in TP topology caused by incorrect grad_norm. #5411
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
|
@inkcherry, amazing contribution!!! @tohtana, @conglongli FYI |
tjruwase
requested review from
tohtana and
conglongli
and removed request for
mrwyattii
conglongli
suggested changes
Apr 15, 2024
There was a problem hiding this comment.
Overall looks good to me and left a few comments. Still need a second review from @tohtana
deepspeed/runtime/utils.py
Outdated
Comment on lines
420
to
429
| # Use grad_norm_mask to avoid redundant computation of flattened gradient norm | ||
| # # including, Pipeline parallelism may replicate parameters. | ||
| # # replicated tensors from tensor model parallelism | ||
|
|
||
| # A loop-free implementation to create a mask tensor based on a range list, | ||
| # which is logically equivalent to the following implementation. | ||
|
|
||
| # # mask_tensor_ = torch.zeros_like(p, device=p.device, dtype=bool) | ||
| # #for mask_idx in grad_norm_mask[idx]: | ||
| # # mask_tensor_[mask_idx[0]:mask_idx[1]] = True |
There was a problem hiding this comment.
Please clean up this block of comments
There was a problem hiding this comment.
cleaned up and kept some for readability
deepspeed/runtime/utils.py
Outdated
| mask_tensor = torch.zeros(p.shape[0] + 1, device=get_accelerator().current_device(), dtype=p.dtype) | ||
| mask_tensor = mask_tensor.scatter_(0, grad_norm_mask[idx].view(-1), | ||
| cum_sum_pairs.view(-1)).cumsum(0).bool()[:-1] | ||
| # assert torch.equal(mask_tensor_, mask_tensor) |
There was a problem hiding this comment.
Please delete this if no longer needed
tohtana
approved these changes
Apr 16, 2024
conglongli
approved these changes
Apr 16, 2024
rraminen
pushed a commit
to ROCm/DeepSpeed
that referenced
this pull request
May 9, 2024
microsoft#5411) Some users are concerned that changes in TP topology during MOE training may potentially cause interference with experiments when noticing similar issues microsoft/Megatron-DeepSpeed#151 https://github.com/microsoft/Megatron-DeepSpeed/pull/176/files We found a grad_norm calculation error after enabling TP. This error occurs because flattened grad of a params group is used, where the group contains both non-TP and TP parameters. Therefore, it is not possible to use a single attribute to determine whether flattened grad needs to compute the norm. In the current code logic, all params are assumed to be non-TP, resulting in only tp_rank0 grad participating in grad_norm computation. Other tp_rank grads have grad_norm_sum equal to 0. We tested and found that with TP=1 and TP=4, the difference in grad_norm is approximately twice (sqrt(4)). This aligns with the aforementioned issue. This problem should also affect dense models. Due to the absence of flattening params_group grad in bf16, this problem is avoided. We tested the loss curve on the 1.3B model. In cases where TP size increases the inconsistent gap should be larger. with this change 1.3B with EP=4 TP=4 &1 , fp16,mbs=1,gbs=16  without this change 1.3B with EP=4 TP=4&1 ,fp16,mbs=1,gbs=16  --------- Co-authored-by: Conglong Li <conglong.li@gmail.com>
umchand
pushed a commit
to umchand/DeepSpeed
that referenced
this pull request
May 20, 2024
microsoft#5411) Some users are concerned that changes in TP topology during MOE training may potentially cause interference with experiments when noticing similar issues microsoft/Megatron-DeepSpeed#151 https://github.com/microsoft/Megatron-DeepSpeed/pull/176/files We found a grad_norm calculation error after enabling TP. This error occurs because flattened grad of a params group is used, where the group contains both non-TP and TP parameters. Therefore, it is not possible to use a single attribute to determine whether flattened grad needs to compute the norm. In the current code logic, all params are assumed to be non-TP, resulting in only tp_rank0 grad participating in grad_norm computation. Other tp_rank grads have grad_norm_sum equal to 0. We tested and found that with TP=1 and TP=4, the difference in grad_norm is approximately twice (sqrt(4)). This aligns with the aforementioned issue. This problem should also affect dense models. Due to the absence of flattening params_group grad in bf16, this problem is avoided. We tested the loss curve on the 1.3B model. In cases where TP size increases the inconsistent gap should be larger. with this change 1.3B with EP=4 TP=4 &1 , fp16,mbs=1,gbs=16  without this change 1.3B with EP=4 TP=4&1 ,fp16,mbs=1,gbs=16  --------- Co-authored-by: Conglong Li <conglong.li@gmail.com>
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Some users are concerned that changes in TP topology during MOE training may potentially cause interference with experiments when noticing similar issues
microsoft/Megatron-DeepSpeed#151
https://github.com/microsoft/Megatron-DeepSpeed/pull/176/files
We found a grad_norm calculation error after enabling TP. This error occurs because flattened grad of a params group is used, where the group contains both non-TP and TP parameters. Therefore, it is not possible to use a single attribute to determine whether flattened grad needs to compute the norm. In the current code logic, all params are assumed to be non-TP, resulting in only tp_rank0 grad participating in grad_norm computation. Other tp_rank grads have grad_norm_sum equal to 0. We tested and found that with TP=1 and TP=4, the difference in grad_norm is approximately twice (sqrt(4)). This aligns with the aforementioned issue. This problem should also affect dense models.
Due to the absence of flattening params_group grad in bf16, this problem is avoided.
We tested the loss curve on the 1.3B model. In cases where TP size increases the inconsistent gap should be larger.
with this change 1.3B with EP=4 TP=4 &1 , fp16,mbs=1,gbs=16
without this change 1.3B with EP=4 TP=4&1 ,fp16,mbs=1,gbs=16
