MoE dispatch: fix silent Tutel + TP corruption, speed up native and Tutel paths - #8195
MoE dispatch: fix silent Tutel + TP corruption, speed up native and Tutel paths#8195jinyouzhi wants to merge 3 commits into
Conversation
Building a dense [s, e] mask per route made gating cost grow with k, so Tutel lost to the native path once k > 2. Gather the per-route index/location/gate values instead and hand Tutel contiguous int32 routes. Routing output is bit-identical. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: iLeGend <824040212@qq.com>
Tutel's encode() returns a flat [e * c, m] buffer, so the drop_tokens() that follows sliced the model dim instead of the capacity dim and experts silently ran on mismatched data. Reshape to [e, c, m] to match the dense path. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: iLeGend <824040212@qq.com>
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 28353d198d
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| if sparse_routes and TUTEL_INSTALLED: | ||
| locations1 = tutel_moe.fast_cumsum_sub_one(mask1) |
There was a problem hiding this comment.
For default k=1 MoE, MOELayer now asks the gate for sparse routes before it chooses the backend, so sparse_routes=True is passed even when the layer was constructed with use_tutel=False. Because this condition keys off TUTEL_INSTALLED instead of the selected backend, any environment that merely has Tutel installed still invokes tutel_moe.fast_cumsum_sub_one on the nominal native path, removing the documented way to opt out of Tutel kernels and polluting native-vs-Tutel behavior. Please keep this on torch.cumsum unless the Tutel backend is actually enabled, or pass the backend choice separately.
Useful? React with 👍 / 👎.
The dense [s, e, c] one-hot einsum cost O(s * e * c * m) to move only s * k * m elements. Route tokens through their capacity slots directly so dispatch and combine each become a single gather. The dense path stays for the inference gate. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: iLeGend <824040212@qq.com>
1 participant
Follow-up to #8174. Three independent changes to
deepspeed/moe/sharded_moe.py.1. Speed up Tutel route extraction for k > 2
topkgatingrebuilt a dense[s, e]one-hot mask for every route, so gating costgrew linearly with
k. On launch-bound configurations Tutel therefore lost to thenative 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 toTutel 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 pathproduces
[e, c, m]. Thedrop_tokens(dispatched_input, dim=1)that followstherefore sliced the model dim instead of the capacity dim. The result still
reshaped cleanly into the expert input (
e * c * m / tpdivided bye * mhappens toyield
c / tp), so there was no error — experts silently ran on a hidden-state slicereinterpreted as capacity, producing wrong results.
Reshaping the buffer to
[e, c, m]before sharding fixes it. This is a pre-existingbug on
master, not introduced by #8174, but #8174 widens its reach by lifting thek == 1restriction onuse_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 raneinsum("sec,sm->ecm"), costingO(s * e * c * m)to move onlys * k * melements.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
MOELayerselects a backend. The dense return path is kept becausedeepspeed/ops/transformer/inference/moe_inference.pystill 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):
Native path (change 3),
s=8192 m=2048 e=16 k=3: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 thedense einsum for both dispatch and combine.
TestMOETensorParallelTutelchecks Tutel against the dense path under tensorparallelism; it fails without change 2 and passes with it.
Cross-checked end-to-end against
master(fp32 max diff 1.2e-07) and verified thatbf16/fp16 accuracy versus an fp32 golden reference does not regress.