[QUESTION] Purpose of all2all communication for probs in MoEAlltoAllTokenDispatcher
Open
community-request
module: moe
question
- Dominant language
- Python
- Stars
- 17.9k
- Forks
- 4.5k
- Avg merge
- 4d 6h
- Merged PRs (30d)
- 271
Description
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!
Contributor guide
Assessment
This issue has not been assessed yet.