Enable support for Tutel when k != 1 for shared moe - #8174
Merged
Conversation
Signed-off-by: iLeGend <824040212@qq.com>
tohtana
approved these changes
Jul 26, 2026
tohtana
left a comment
Collaborator
There was a problem hiding this comment.
Looks good to me, thank you @jinyouzhi!
Performance note from the validation on my side: on one 2×H100 BF16, the Tutel/native step-time ratio was 0.948 for k=2, but 1.119 for k=3. Correctness looks good so it is fine to merge, but Tutel might not have performance advantage for k=3 in this configuration.
Contributor
Author
Thank you for your review and perf metrics. With your insights, I further tried to optimize the k=3 case to reduce overhead. |
banxingmjj
pushed a commit
to openanolis/DeepSpeed
that referenced
this pull request
Aug 2, 2026
…utel paths (deepspeedai#8195) Follow-up to deepspeedai#8174. Three independent changes to `deepspeed/moe/sharded_moe.py`. ### 1. Speed up Tutel route extraction for k > 2 `topkgating` rebuilt a dense `[s, e]` one-hot mask for every route, so gating cost grew linearly with `k`. On launch-bound configurations Tutel therefore lost to the native path once `k > 2`. Since the top-k columns already name the selected experts, the per-route index/location/gate values are now a `gather`. Routes are handed to Tutel as contiguous int32 so its own casts become no-ops, and the capacity tensor is resolved once instead of forcing a device-to-host sync per use. Routing output is bit-identical. ### 2. Shard the Tutel dispatch buffer along capacity under tensor parallelism Tutel's `encode()` returns a flat `[e * c, m]` buffer, whereas the dense path produces `[e, c, m]`. The `drop_tokens(dispatched_input, dim=1)` that follows therefore sliced the **model** dim instead of the **capacity** dim. The result still reshaped cleanly into the expert input (`e * c * m / tp` divided by `e * m` happens to yield `c / tp`), so there was no error — experts silently ran on a hidden-state slice reinterpreted as capacity, producing wrong results. Reshaping the buffer to `[e, c, m]` before sharding fixes it. This is a pre-existing bug on `master`, not introduced by deepspeedai#8174, but deepspeedai#8174 widens its reach by lifting the `k == 1` restriction on `use_tutel`. ### 3. Dispatch MoE tokens by index instead of a dense one-hot einsum The native path materialised a dense `[s, e, c]` one-hot and ran `einsum("sec,sm->ecm")`, costing `O(s * e * c * m)` to move only `s * k * m` elements. Tokens are now routed through their capacity slots directly, so dispatch and combine each become a single gather. The sparse routing metadata already existed for Tutel, so the gate always returns it and `MOELayer` selects a backend. The dense return path is kept because `deepspeed/ops/transformer/inference/moe_inference.py` still consumes it. Combine accumulates in fp32, matching the einsum it replaces — the einsum accumulated in fp32 inside the tensor-core matmul, so a naive low-precision gather would have been *less* accurate than the code it replaced. ## Results Measured on 2 GPUs, bf16, forward+backward. Tutel/native step-time ratio (change 1): | k | before | after | |---|--------|-------| | 2 | 0.947 | 0.937 | | 3 | 1.078 | 1.01 | | 4 | 1.158 | 1.04 | Native path (change 3), `s=8192 m=2048 e=16 k=3`: | | step time | peak memory | |---|-----------|-------------| | before | 58.33 ms | 4353 MB | | after | 37.25 ms | 2918 MB | The native path now matches Tutel (37.25 vs 37.49 ms). ## Testing `tests/unit/moe/`: 50 passed with Tutel installed, 48 passed / 2 skipped without. - `test_sparse_dispatch_matches_dense` (k=1/2/3) checks the index path against the dense einsum for both dispatch and combine. - `TestMOETensorParallelTutel` checks Tutel against the dense path under tensor parallelism; it fails without change 2 and passes with it. Cross-checked end-to-end against `master` (fp32 max diff 1.2e-07) and verified that bf16/fp16 accuracy versus an fp32 golden reference does not regress. > `TestMOETensorParallelTutel` requires Tutel and is skipped in CI, which does not > install it. Change 2 was validated locally on 2 GPUs: max diff `7.27e-01` before the > fix, `0.00e+00` after. --------- Signed-off-by: iLeGend <824040212@qq.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Co-authored-by: Masahiro Tanaka <81312776+tohtana@users.noreply.github.com>
This file contains hidden or 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
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
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.
Remove the limited on top2 and topk >2 for Tutel.