The MoE spends too much time on the gate network.

[zhaozheng09]

Model is GR(generative recommendation)SMoE
MoE select top 4 ep, total 12ep, bs 64.
token's num is 32768 in one steps.

deepspeed version is 0.17.4

        self.ff = FeedForward(input_dim, inner_dim, activation, "deepspeed_expert")
        self.ff_out = MoE(
            hidden_size=input_dim,  # 注意这里要用hidden_size
            expert=self.ff,
            num_experts=num_experts,
            k=top_k,
            ep_size=1,
        )

why gate speed too much time ?

[Haiderkhan64]

Hey, I dug into this a bit since I hit something similar. I reproduced your config (num_experts=12, k=4, ep_size=1, 32768 tokens/step, deepspeed 0.17.4) to see where the time was actually going.
Short answer: it's probably not the gate itself that's slow, it's the dispatch mask that gets built right after it.
Since ep_size=1, there's no cross-GPU alltoall happening, so it can't be communication. For k>2, DeepSpeed goes through topkgating() in sharded_moe.py, and capacity there is computed as:
capacity = ceil((tokens / num_experts) * capacity_factor * k)
Plug in your numbers and you get:
ceil((32768 / 12) * 1.0 * 4) = 10,923
That capacity then gets used to build a dense one-hot tensor of shape [32768, 12, 10923] for the dispatch mask, and this is built at full size regardless of how many tokens each expert actually ends up with. So the cost is basically driven by capacity_factor × k, and with k=4 that adds up fast.
I didn't just take this from the source though, I actually ran it: on a 16GB GPU, capacity_factor=1.0 OOMs trying to allocate exactly 32.00 GiB, and capacity_factor=0.5 OOMs at exactly 16.00 GiB. That's a clean linear match with the formula above, so I'm fairly confident this is really what's eating your time/memory.
A couple things that should help:

• Lower capacity_factor. Given the 1.0→32GB, 0.5→16GB pattern, dropping to 0.25 should land around ~8GB. Obviously this means dropping more overflow tokens (drop_tokens defaults to True), so keep an eye on exp_counts after you change it.
• Don't bother with use_tutel at k=4, I checked the source and Tutel only actually kicks in when k=1 (gate.k == 1). At k=4 it silently falls back to the same dense path, you'd just get a warning saying to use top-1 instead of top-2. If you can drop to k=1 or k=2, you'd both shrink the capacity formula and get Tutel's sparse dispatch, which skips this whole dense tensor.
• Worth turning on "wall_clock_breakdown": true in your DeepSpeed config too, it gives you separate timers for the moe layer vs the two alltoalls, so you can confirm in your actual training run that this is really where the time goes, not just in my isolated repro.
  Let me know what capacity_factor you're running and whether dropping tokens is acceptable for your setup, happy to dig further.

[jinyouzhi]

Very insightful! #8174 made use_tutel available for k > 1, and #8195 added some optimizations for MoE. @Haiderkhan64 @zhaozheng09

If you have any insights on optimization directions, please feel free to share your thoughts.

[Haiderkhan64]

Thanks for the pointer -- reproduced the original config (tokens=32768, experts=12, ep_size=1, k=1-4, capacity_factor=1.0/0.5/0.25) against current master (DeepSpeed 0.19.4, post-#8174/#8195) on a 16GB GPU.
Total step time still scales almost exactly linearly with capacity_factor at every k (halving capacity_factor roughly halves step time, consistently from k=1 through k=4). k=4 at the default capacity_factor=1.0 -- the original issue's exact settings -- still OOMs, reproduced directly on today's code.
To isolate why, I ran the expert FFN alone on a synthetic [experts, capacity, hidden] tensor with no gate/dispatch/combine at all, at the same capacity the full MoE layer uses for each (k, capacity_factor). Expert-only time accounts for 87-99% of total MoE step time across every configuration tested, and that share increases with k. At k=4/capacity_factor=1.0 specifically, the bare expert FFN alone fits under ~10.4GB, while the full MoE layer at that same capacity still OOMs -- so there is additional memory consumption beyond the FFN buffer itself, although I haven't isolated which specific allocation (indices, combine buffers, workspace, etc.) accounts for it.
So #8195's dispatch optimization looks correct and effective at what it targeted -- on current 0.19.4, gate/dispatch/combine is now only a small fraction of total cost. The remaining scaling is dominated by the expert FFN operating over an experts × capacity buffer, and capacity = (tokens / experts) × capacity_factor × k is linear in k by construction. In this implementation, k therefore influences the cost twice: it increases the number of expert assignments per token, while also increasing the capacity allocated to each expert. That scaling is unchanged by #8195, so the original default-settings configuration still doesn't fit in 16GB today.
If useful, options from here are lowering capacity_factor (more token drop), lowering k, or moving away from the fixed-capacity buffer toward padding-free/grouped-GEMM expert execution so compute scales with actual routed tokens rather than worst-case per-expert allocation. Happy to share the benchmark scripts if it'd help reproduce this.

[jinyouzhi]

Thank you very much. Your experiments helped us confirm that the dispatch overhead has been significantly reduced.

Your observation that memory consumption for the expert-capacity buffer grows linearly with k is consistent with DeepSpeed’s current implementation, as far as I understand.

By the way, if I remember correctly, Megatron-LM also supports an alternative approach based on tokens_per_expert, whereas DeepSpeed currently uses fixed-capacity buffers. If this optimization is important for your use case, you’re welcome to propose it as a roadmap item.

Thanks again for the detailed investigation.

[Haiderkhan64]

Thanks for the clarification. The tokens_per_expert direction looks like a promising way to address the remaining cost.
My current measurements suggest that after #8195, dispatch is no longer the dominant bottleneck; the main cost comes from expert execution over the fixed experts × capacity × hidden buffer. Since capacity scales with k, this also causes the expert-side cost and memory footprint to grow with k.
I’d like to investigate a small prototype of variable-length / padding-free expert execution, starting with ep_size=1, where the expert computation is based on the actual routed tokens_per_expert rather than the fixed capacity.
I’ll benchmark it against the current implementation in terms of runtime, peak memory, and correctness, and can share the results as a possible roadmap proposal.

[jinyouzhi]

Great, looking forward to your next update! I’d personally suggest opening a new issue when you have a concrete plan or prototype.