DeepSpeed

619e3199 - 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 (#8195) Follow-up to #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 #8174, but #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>