Support MoE for pipeline models #5338
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
mosheisland
requested review from
tjruwase,
mrwyattii,
duli2012,
awan-10,
arashb and
loadams
as code owners
|
Please note that Megatron-DeepSpeed PR#373 https://github.com/microsoft/Megatron-DeepSpeed/pull/373 is dependent on this PR. |
Signed-off-by: Moshe Island <misland@habana.ai>
Currently MoE uses Megatron-DeepSpeed APIs to get tensor-parallel info (rank, world_size, group). In order to enable MoE for PipelineModule, modify to use backward-compatible methods that can access either Megatron, DeepSpeed Topology or Old Megatron APIs. Since MoE is not part of deepspeed runtime, move backward compatible methods to deepseed.utils and modify imports as required. Signed-off-by: Moshe Island <misland@habana.ai>
Signed-off-by: Moshe Island <misland@habana.ai>
Currently, only "total_loss" is displayed. If model has additional losses (e.g. MoE loss), display them as well. Similar to "total loss", additional losses are displayed for the full batch after mean reduced across DP ranks. Signed-off-by: Moshe Island <misland@habana.ai>
Signed-off-by: Moshe Island <misland@habana.ai>
Currently, when using no-drop tokens, we calculate locally the capacity and then all-reduce(op=max) on world group. This fails when using pipeline parallel (with micro batches), since different stage workers are handling different model layers (or at warmup, where first stage workers are processing while last stage workers are idle). Fix it by running the all-reduce on the expert group. Signed-off-by: Moshe Island <misland@habana.ai>
This commit enhances expert group creation for both modes: - DP + PP + EP - DP + TP + PP + EP Signed-off-by: Moshe Island <misland@habana.ai>
When using MoE with MoE-TP disabled, use pipeline parallel group to max or sum MoE gradients. This also fixes the behavior for following configuration: No pipeline, TP enabled, MoE TP disabled. Signed-off-by: Moshe Island <misland@habana.ai>
mosheisland
force-pushed
the
moe/pipe
branch
2 times, most recently
from
7a5e888
to
0f9d2b5
Compare
tohtana
approved these changes
Apr 4, 2024
|
I left a comment about a very small part but have already approved this PR. |
Signed-off-by: Moshe Island <misland@habana.ai>
Done |
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
rraminen
pushed a commit
to ROCm/DeepSpeed
that referenced
this pull request
May 9, 2024
This PR enhances DeepSpeed to support MoE for pipeline models (e.g. GPTModelPipe from Megatron-DeepSpeed). Main changes: - Enhance expert groups creation for pipeline (enhance both flavors: DP/PP/EP and DP/TP/PP/EP) - Fix MoE save/load checkpoint for PipelineModule based models. - Display MoE loss for PipelineModule based models. - Support gradients reduce for BF16_Optimizer for PipelineModule.<br>Note that same commit also fixes gradients reduction error when using Megatron-DeepSpeed GPTModelPipe with BF16_Optimizer also for a dense (no MOE) model. - When using no-drop tokens, all-reduce the capacity (op=max) using expert parallel group instead of world group --------- Signed-off-by: Moshe Island <misland@habana.ai> Co-authored-by: Moshe Island <misland@habana.ai>
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.
This PR enhances DeepSpeed to support MoE for pipeline models (e.g. GPTModelPipe from Megatron-DeepSpeed).
Main changes:
Note that same commit also fixes gradients reduction error when using Megatron-DeepSpeed GPTModelPipe with BF16_Optimizer also for a dense (no MOE) model.