It was interesting to compare MoE Parallel Folding in Megatron with JAX/TPU architecture.
MoE Parallel Folding lets attention layers and MoE layers use different parallelism settings in the same model.
a. Attention works over the whole sequence, so every GPU in a tensor parallel or context parallel group has to exchange activations with every other GPU in that group.
b. MoE layers work on one token at a time, so expert parallelism only sends each token to the GPU holding its expert, which costs far less.
c. High tensor parallelism helps attention but hurts MoE, because each expert is small and the tensor parallel overhead is large relative to the work.
d. High context parallelism helps long context attention but does nothing for MoE, since tokens there need to be processed independently.
e. Previous setups nest the expert parallel group inside the data parallel group, so expert parallelism cannot not grow beyond the data parallel size.
Folding
a. Megatron separates out attention to its own group of tensor, context, data and pipeline parallelism,
and MoE gets a separate group of expert tensor, expert, expert data and pipeline parallelism.
b. Pipeline parallelism is the one setting they share, since pipeline stages are physical GPU placements used by both layer types.
c. The expert parallel group can be nested inside the attention layer's tensor or context parallel group, so its traffic stays inside a node on NVLink instead of crossing the slower network between nodes.
d.As an exmple keeping the attention config unchanged and setting MoE to expert tensor parallel 1 and expert parallel 64 gives 8 times more expert parallelism.
Token dispatcher for working with any combination
Once attention and MoE can pick independent parallel strategies, there needs to be a dispatcher that correctly routes tokens to experts no matter which combination was chosen
Megatron dispatcher handles this with a five-step communication pattern per MoE layer: permute tokens locally → All-to-All-V across the EP group → AllGather-V within the ETP (expert tensor parallel) group → local expert computation → ReduceScatter-V within ETP → All-to-All-V back → un-permute.
JAX never needed a folding feature, because sharding works very differently.
a. Megatron builds fixed communication groups at startup, so a second set of groups had to be added for MoE layers.
c. JAX declares one device mesh with named axes and shards each tensor separately against it, so attention weights can use one axis and expert weights and they are independent.
d. The XLA compiler then inserts the all gathers, reduce scatters and all to all calls, instead of a hand written dispatcher.