RFC: Primus-Megatron SyncFree MoE

[zhenhuang12]

Summary

The dynamic shape of MoE (such as, required hip/cuda kernel results to allocate device memory) can lead to D2H synchronization and cause significant CPU overhead. This may increase device idle time and impose considerable performance impact on training, which becomes particularly noticeable when context-parallelism is enabled.

We have eliminated all CPU synchronization throughout MoE pipeline—from Router to Dispatcher, Permutation, and GroupMLP —reducing the idle time.

The Sync-Free MoE implementation involves multiple parameters and options, each with trade-offs. Optimal performance requires careful tuning based on specific training scenarios, making usage complex. This document proposes provides primus users with configurable options supporting multiple levels about Sync-free MoE, allowing to select the appropriate level based on their actual needs.

We hope that this design will:

• Provide simple and clear options with comprehensive documentation.
• Provide reasonable multi-level of Sync-Free MoE to cover most of primus users's needs.

SyncFree-MoE Workflow

As we can see from the above diagram, almost all component of Megatron MoE requires CPU synchronization:

• Router
• Dispatcher
• Permutation
• GroupMLP

TODO

Proposed Options for Primus-Megatron

We provide primus user turbo_sync_free_moe_stage option and divided into 4 levels (0-3):

• 0: Disable the Sync-Free MoE (default)
• 1: Remove synchronization for Router and Permutation
• 2: Remove synchronization for Router, DeepEP, and GroupMLP
• 3: Remove all MoE synchronization

Sync-Free-related function parameters or options in Primus-Megatron MoE, as following:

• Router:

  • fused_group_topk_routing_with_aux_score: Primus-Megatron option to enable router fusion.

• DeepEP

  • use_cuda_num_token_per_expert: Turbo-DeepEP dispatch parameter to return num_tokens_per_experts by CUDA tensor
  • num_worst_token: Turbo-DeepEP dispatch parameter to eliminate notify-dispatch CPU busy-wait

• permuatation

  • moe_permutation_fusion: Primus-Megatron option to enable permutation fusion.

• groupmlp

  • use_turbo_groupmlp: Primus-Megatron option to use Turbo's GroupMLP which accepted CUDA num_token_per_experts as parameter.
  • use_turbo_groupmlp_act: Primus-Megatron option to use Turbo's activation which accepted CUDA num_token_per_experts as parameter.

syncfree_moe_stage	Involved options or parameters	CPU Sync	Single Rank Memory Usage
0	/	/	/
1	fused_group_topk_routing_with_aux_score moe_permutation_fusion	host sync at both deepep and groupmlp	/
2	fused_group_topk_routing_with_aux_score use_cuda_num_token_per_expert num_worst_token moe_permutation_fusion use_turbo_groupmlp	host sync once before permutation, CPU bottleneck mainly related to CPU overhead of permutation and group MLP.	T * P (temporary)
3	fused_group_topk_routing_with_aux_score use_cuda_num_token_per_expert num_worst_token moe_permutation_fusion use_turbo_groupmlp use_turbo_groupmlp_act	No	3 *K * T * P (activation)

• Assuming T = token bytes is batch_size * seqlen * hidden_size * element_size
• Assuming P = world size of EP
• Assuming K = topk

Note:
Assuming deepseekv3 with ep16 as example, each rank of stage2 only has 4096*7168*16*sizeof(bfloat16)=896MB of memory overhead. However, each rank of stage3 requires activation overhead is 4096*7168*16*sizeof (bfloat16) *8*61 *3 = 1281GB

[RuibinCheung]

Great work !

There are some suggestions and questions:

1. I think sync_free_moe is more clear than syncfree_moe. Can we rename syncfree_moe to sync_free_moe ?
2. I suggest you can highlight the memory usage is total usage rather than single rank usage.
3. Could you explain why memory usage need times 3 in the sync-free moe 3 level ?

Thanks a lot~

[zhenhuang12]

Great work !

There are some suggestions and questions:

1. I think sync_free_moe is more clear than syncfree_moe. Can we rename syncfree_moe to sync_free_moe ?
2. I suggest you can highlight the memory usage is total usage rather than single rank usage.
3. Could you explain why memory usage need times 3 in the sync-free moe 3 level ?

Thanks a lot~

Thanks for your suggestion!

For suggestion 1, I will rename turbo_syncfree_moe to turbo_sync_free_moe, that's completely fine.
For suggestion 2, the memory usage is a single rank memory usage, I will highlight it.
For the last Question, memory usage, which times 3 in sync-free moe level 3, is the activation of 2 groupmlp and 1 act