[QUESTION] Purpose of all2all communication for probs in MoEAlltoAllTokenDispatcher

[snowpeakz]

In the ADLR/megatron-lm!2668 - perf(MoE): Memory efficient token permutation, I noticed an all2all communication is applied to the probs tensor in MoEAlltoAllTokenDispatcher . Could someone tell me the motivation for doing this? I didn't observe a drop in memory usage compared to previous versions.

Appreciate any insights or references to discussions/design documents!

[snowpeakz]

@yanring @Victarry I wonder if you could help me understand, thanks

[yanring]

Hi @snowpeakz,
What kind of model are you using for testing? Specifically, what are the values for the model’s ⁠topK, ⁠moe_ffn_size, and ⁠hidden_size?
Generally speaking, this feature can help fine-grained MoE by using a large ⁠TopK value and a small ⁠moe_ffn_size.
The reason we need to communicate the ⁠probs is that the multiplication of tokens and probabilities occurs at the ⁠moe_act stage. At this point, the tokens have already undergone all-to-all (a2a) communication, so the corresponding ⁠probs also need to undergo the same communication process.

[snowpeakz]

@yanring Thanks a lot for your answer!

I'm testing deepseek v2 with following args

--moe-router-topk 6
--moe-ffn-hidden-size 1536
--moe-shared-expert-intermediate-size 3072

It seems that the released tensor accounts for a small proportion of the memory usage and is not obvious?