Fix overlap communication of ZeRO stage 1 and 2 #5606
Merged
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
penn513
force-pushed
the
master
branch
2 times, most recently
from
6a001f9 to
1319870
Compare
|
@microsoft-github-policy-service agree company="Huawei" |
|
Hi @tjruwase @GuanhuaWang , would you please help to review this PR? |
|
Hi @penn513 thx for the nice figure and pr. I think your fix on compute stream wait back reduce stream make sense to me, especially when compute is shorter. To make PR more concise, could you remove your modification on npu fuseadam in current pr, and make a new pr on fused adam? (mainly because it is irrelevant with this PR title as "fix overlap communication...") |
Co-authored-by: CurryRice233 <nmeia@qq.com>
Thanks for your reply. It's been updated. |
jomayeri
approved these changes
Jun 7, 2024
github-merge-queue
bot
removed this pull request from the merge queue due to failed status checks
github-merge-queue
bot
removed this pull request from the merge queue due to failed status checks
github-merge-queue
bot
removed this pull request from the merge queue due to failed status checks
github-merge-queue
bot
removed this pull request from the merge queue due to failed status checks
sfc-gh-reyazda
pushed a commit
to Snowflake-Labs/DeepSpeed
that referenced
this pull request
Jun 10, 2024
`deepspeed.runtime.zero.stage_1_and_2.DeepSpeedZeroOptimizer.average_tensor` only sets reduction stream waiting for default stream. This is ok in cases where the computation time is longer than the communication time, but when the communication time is longer, it may result in a rewrite of the ipg_buffer when the communication is not completed.  To fix this bug, the easiest way is just add default stream to wait for reduction stream at the **same point**. For example, in point 1, the `reduction stream` needs to wait for '2', so we add a wait_stream to `reduction stream` waiting for `default stream`. Also, the `default stream` needs to wait for 'A', so we need to add a wait_stream to `default stream` waiting for `reduction stream` before the 'B'.  Compared with the modification of deepspeedai#5523, wait_stream does not cause host synchronization. Compared with the modification of deepspeedai#5545, the modification is more simple and the logic is the same, just waiting for what needs to wait. --- With this modification, losses of Qwen-1.5 with and without overlap_comm are totally identical.  --- On the contrary, there is an obvious gap with a small sequence length, which means a short computation time.  Co-authored-by: gp513 <guopeng34@huawei.com> Co-authored-by: CurryRice233 <nmeia@qq.com> Co-authored-by: Joe Mayer <114769929+jomayeri@users.noreply.github.com> Co-authored-by: Logan Adams <114770087+loadams@users.noreply.github.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.
deepspeed.runtime.zero.stage_1_and_2.DeepSpeedZeroOptimizer.average_tensoronly sets reduction stream waiting for default stream. This is ok in cases where the computation time is longer than the communication time, but when the communication time is longer, it may result in a rewrite of the ipg_buffer when the communication is not completed.To fix this bug, the easiest way is just add default stream to wait for reduction stream at the same point. For example, in point 1, the
reduction streamneeds to wait for '2', so we add a wait_stream toreduction streamwaiting fordefault stream. Also, thedefault streamneeds to wait for 'A', so we need to add a wait_stream todefault streamwaiting forreduction streambefore the 'B'.Compared with the modification of #5523, wait_stream does not cause host synchronization.
Compared with the modification of #5545, the modification is more simple and the logic is the same, just waiting for what needs to wait.
With this modification, losses of Qwen-1.5 with and without overlap_comm are totally identical.
On the contrary, there is an obvious gap with a small sequence length, which means a short computation time.