Add DeepEP v2 flex dispatcher backend - #5153
Conversation
|
This PR has been automatically converted to draft because all PRs must start as drafts. When you are ready for review, click Ready for Review to begin the review process. This will:
See the contribution guide for more details. |
Autumn1998
changed the title
|
/ok to test 646cceb |
|
/ok to test c0a73a0 |
Signed-off-by: tongliu <tongliu@nvidia.com>
Signed-off-by: tongliu <tongliu@nvidia.com>
Signed-off-by: tongliu <tongliu@nvidia.com>
Signed-off-by: tongliu <tongliu@nvidia.com>
Signed-off-by: tongliu <tongliu@nvidia.com>
Signed-off-by: tongliu <tongliu@nvidia.com>
Signed-off-by: tongliu <tongliu@nvidia.com>
7b85b7d to
30e602f
Compare
|
/ok to test 30e602f |
YangFei1990
left a comment
YangFei1990
left a comment
There was a problem hiding this comment.
Do you think we should merge DeepEPv2 and v2 instead of having complete separate APIs? I understand separation brought benefit for easy management, but currently from what I see almost 80-90% of code are duplicated.
Also, can you update Dockerfile.ci.dev's DEEPEP_COMMIT to point to newer commit in hybrid-ep so DeepEP v2 is also supported?
| return _buffer | ||
|
|
||
|
|
||
| def get_elastic_buffer( |
There was a problem hiding this comment.
Wondering why do we have this buffer reuse mechanism for DeepEPv2, seems it is still full eager execution (shape is reading during the runtime). Also is it safe to have different layer/mb to reuse the same buffer
There was a problem hiding this comment.
I think it is the recommended way to use DeepEP v2:https://nvidia.slack.com/archives/C03V462SAMS/p1779421202250149
| def _get_buffer(self, hidden_states: torch.Tensor): | ||
| self.buffer = get_elastic_buffer( | ||
| self.group, | ||
| num_max_tokens_per_rank=hidden_states.shape[0], |
There was a problem hiding this comment.
Could the num_max_tokens_per_rank be different across ranks?
There was a problem hiding this comment.
I think different ranks should still pass in different values, so in the THD case we may need to pad.
I actually think it’s fine. The mirror PR in dev has already been merged. This kind of refactoring will cause conflicts when dev pulls from main later. If a fix is really needed, we can submit another PR afterward. Currently, hep hasn’t merged deepep’s main yet, so v2 isn’t supported. At the moment, we can’t support both deepep and hep in a single commit. |
Is there a plan to update hep? |
maybe not |
…-dispatcher-main Signed-off-by: tongliu <tongliu@nvidia.com> # Conflicts: # megatron/core/models/gpt/fine_grained_callables.py # megatron/core/transformer/moe/token_dispatcher.py # megatron/core/transformer/transformer_config.py # tests/unit_tests/transformer/moe/test_token_dispatcher.py
|
Unit test fix: SIGSEGV in DeepEP v2 tests due to uninitialized NCCL communicator error: Root cause: PyTorch lazily initializes NCCL communicators — ProcessGroupNCCL._comm_ptr() returns 0 (null) until the first collective operation is performed on that group. DeepEP v2's get_nccl_comm_handle() calls backend._comm_ptr() and passes the result directly to the C++ function calculate_elastic_buffer_size(). When the communicator hasn't been initialized yet, this dereferences a null pointer → SIGSEGV. In training this is never a problem because process groups are already warmed up by model initialization and gradient syncs before the first MoE forward pass. But in unit tests, initialize_model_parallel() creates fresh sub-groups and the first operation on tp_ep_group is ElasticBuffer.get_buffer_size_hint(), which crashes. Fix: Add a single dummy all_reduce on tp_ep_group in the test fixture (MoEModelTestContainer) when moe_flex_dispatcher_backend == "deepepv2". This forces NCCL communicator creation before DeepEP v2 reads the pointer. ● Update(tests/unit_tests/transformer/moe/test_token_dispatcher.py)
⎿ Added 6 lines
98 # init moe layer
99 self.moe_layer = self.new_moe_layer()
100
101 + # DeepEP v2 reads _comm_ptr() which is null until the first collective.
102 + # Warm up the tp_ep group so the NCCL communicator is initialized.
103 + if kwargs.get("moe_flex_dispatcher_backend") == "deepepv2":
104 + tp_ep_group = self.moe_layer.token_dispatcher.tp_ep_group
105 + torch.distributed.all_reduce(torch.zeros(1, device="cuda"), group=tp_ep_group)
106 +
107 def new_moe_layer(self, **kargs):
108 submodules = get_submodules(
109 get_gpt_layer_local_submodules(Reference: DeepEP's own test (tests/elastic/test_ep.py) avoids this by passing device_id to init_process_group(), which triggers eager NCCL initialization. We use the more targeted warmup approach to avoid symmetric memory pressure from repeated group creation/destruction in the test suite. |
| previous_event = buffer.capture() if async_finish and allocate_on_comm_stream else None | ||
| # Process the dispatch and keep the handle for the subsequent combine call. | ||
| recv_x, recv_token_indices, recv_token_probs, handle, event = buffer.dispatch( | ||
| x, |
There was a problem hiding this comment.
could support do_expand=True ? @Autumn1998
7 participants
What does this PR do ?
The DeepEP V2 support as a backend of flex dispatcher.
PR on dev: #4793
Issue tracking
For PRs from open-source community contributors:
Linked issue:
Contribution process
Pre-checks
Code review
Feel free to message or comment @NVIDIA/mcore-oncall to help accelerate your merge into main. The less complex your PR is, the faster it will be approved and merged!
All PRs start as draft. If you open a non-draft PR, it will be automatically converted to draft.
Step 1: Mark PR as "Ready for Review"
.github/CODEOWNERS.Final Review might get declined if these requirements are not fulfilled.
Step 2: Final Review
For PRs that change
megatron/core, once all expert reviewers have approved, theFinal Reviewlabel is applied automatically and final reviewers are assigned.For PRs outside
megatron/core, this step is skipped.Step 3: Approved
Once all required reviewers have approved, the
Approvedlabel is applied automatically.Merge
Any member of mcore-engineers will be able to merge your PR.