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.