perf: Update moe_token_dispatcher_type default to alltoall (#2004)

parthmannan · terrykong · web-flow · commit 919e373b112c · 2026-03-06T09:17:28.000Z
Signed-off-by: Parth Mannan &lt;pmannan@nvidia.com&gt;
Co-authored-by: Terry Kong &lt;terrycurtiskong@gmail.com&gt;
diff --git a/examples/configs/distillation_math.yaml b/examples/configs/distillation_math.yaml
@@ -109,7 +109,7 @@ policy: &POLICY_BASE
         defer_fp32_logits: False
         moe_per_layer_logging: False
         moe_enable_deepep: false
-        moe_token_dispatcher_type: "allgather"
+        moe_token_dispatcher_type: "alltoall"
         moe_shared_expert_overlap: false
         
         optimizer:
diff --git a/examples/configs/distillation_math_megatron.yaml b/examples/configs/distillation_math_megatron.yaml
@@ -60,7 +60,7 @@ policy: &POLICY_BASE
         moe_per_layer_logging: False
         defer_fp32_logits: False
         moe_enable_deepep: false
-        moe_token_dispatcher_type: "allgather"
+        moe_token_dispatcher_type: "alltoall"
         moe_shared_expert_overlap: false
         peft:
             enabled: false
diff --git a/examples/configs/dpo.yaml b/examples/configs/dpo.yaml
@@ -133,7 +133,7 @@ policy:
     defer_fp32_logits: False
     moe_per_layer_logging: False
     moe_enable_deepep: false
-    moe_token_dispatcher_type: "allgather"
+    moe_token_dispatcher_type: "alltoall"
     moe_shared_expert_overlap: false
     
     optimizer:
diff --git a/examples/configs/grpo_math_1B.yaml b/examples/configs/grpo_math_1B.yaml
@@ -144,7 +144,7 @@ policy:
     defer_fp32_logits: False
     moe_per_layer_logging: False
     moe_enable_deepep: false
-    moe_token_dispatcher_type: "allgather"
+    moe_token_dispatcher_type: "alltoall"
     moe_shared_expert_overlap: false
 
     peft:
diff --git a/examples/configs/grpo_math_1B_megatron.yaml b/examples/configs/grpo_math_1B_megatron.yaml
@@ -96,7 +96,7 @@ policy:
     moe_router_bias_update_rate: 0.0 # by default, disable bias updates for grpo
     moe_permute_fusion: false
     moe_enable_deepep: false
-    moe_token_dispatcher_type: "allgather"
+    moe_token_dispatcher_type: "alltoall"
     moe_shared_expert_overlap: false
     #gives ~20% training perf speedup with sequence packing 
     apply_rope_fusion: True
diff --git a/examples/configs/sft.yaml b/examples/configs/sft.yaml
@@ -115,7 +115,7 @@ policy:
     defer_fp32_logits: False
     moe_per_layer_logging: False
     moe_enable_deepep: false
-    moe_token_dispatcher_type: "allgather"
+    moe_token_dispatcher_type: "alltoall"
     moe_shared_expert_overlap: false
 
     peft:
diff --git a/examples/configs/sft_openmathinstruct2_megatron.yaml b/examples/configs/sft_openmathinstruct2_megatron.yaml
@@ -94,7 +94,7 @@ policy:
     bias_activation_fusion: True
     moe_per_layer_logging: False
     moe_enable_deepep: false
-    moe_token_dispatcher_type: "allgather"
+    moe_token_dispatcher_type: "alltoall"
     moe_shared_expert_overlap: false
     peft:
       enabled: false
diff --git a/examples/configs/vlm_grpo_3B.yaml b/examples/configs/vlm_grpo_3B.yaml
@@ -114,7 +114,7 @@ policy:
     defer_fp32_logits: False
     moe_per_layer_logging: False
     moe_enable_deepep: false
-    moe_token_dispatcher_type: "allgather"
+    moe_token_dispatcher_type: "alltoall"
     moe_shared_expert_overlap: false
 
     optimizer:
diff --git a/examples/configs/vlm_grpo_3B_megatron.yaml b/examples/configs/vlm_grpo_3B_megatron.yaml
@@ -156,7 +156,7 @@ policy:
     defer_fp32_logits: False
     moe_per_layer_logging: False
     moe_enable_deepep: false
-    moe_token_dispatcher_type: "allgather"
+    moe_token_dispatcher_type: "alltoall"
     moe_shared_expert_overlap: false
     peft:
       enabled: false
diff --git a/examples/nemo_gym/grpo_workplace_assistant_nemotron_nano_v2_9b.yaml b/examples/nemo_gym/grpo_workplace_assistant_nemotron_nano_v2_9b.yaml
@@ -118,7 +118,7 @@ policy:
     defer_fp32_logits: false
     moe_permute_fusion: false
     moe_enable_deepep: false
-    moe_token_dispatcher_type: "allgather"
+    moe_token_dispatcher_type: "alltoall"
     moe_shared_expert_overlap: false
 
     optimizer:
diff --git a/nemo_rl/models/policy/__init__.py b/nemo_rl/models/policy/__init__.py
@@ -205,7 +205,7 @@ class MegatronConfig(TypedDict):
     # Must set moe_token_dispatcher_type to 'flex'
     # Must set moe_shared_expert_overlap to False
     moe_enable_deepep: bool
-    # The type of token dispatcher to use. The default is 'allgather'.
+    # The type of token dispatcher to use. The default is 'alltoall'.
     # Options are 'allgather','alltoall' and 'flex'
     # Use 'flex' when using DeepEP
     moe_token_dispatcher_type: str
diff --git a/tests/unit/models/generation/test_vllm_generation.py b/tests/unit/models/generation/test_vllm_generation.py
@@ -194,7 +194,7 @@ def get_basic_megatron_test_config(
             "moe_router_bias_update_rate": 0.0,
             "moe_permute_fusion": False,
             "moe_enable_deepep": False,
-            "moe_token_dispatcher_type": "allgather",
+            "moe_token_dispatcher_type": "alltoall",
             "moe_shared_expert_overlap": False,
             "apply_rope_fusion": True,
             "bias_activation_fusion": True,
diff --git a/tests/unit/models/megatron/test_megatron_setup.py b/tests/unit/models/megatron/test_megatron_setup.py
@@ -187,7 +187,7 @@ def test_moe_configuration(self):
                 "moe_router_bias_update_rate": 0.0,
                 "moe_permute_fusion": True,
                 "moe_enable_deepep": False,
-                "moe_token_dispatcher_type": "allgather",
+                "moe_token_dispatcher_type": "alltoall",
                 "moe_shared_expert_overlap": True,
             }
         }
@@ -201,7 +201,7 @@ def test_moe_configuration(self):
         assert model_cfg.moe_router_bias_update_rate == 0.0
         assert model_cfg.moe_permute_fusion is True
         assert model_cfg.moe_enable_deepep is False
-        assert model_cfg.moe_token_dispatcher_type == "allgather"
+        assert model_cfg.moe_token_dispatcher_type == "alltoall"
         assert model_cfg.moe_shared_expert_overlap is True
 
 
diff --git a/tests/unit/models/policy/test_megatron_worker.py b/tests/unit/models/policy/test_megatron_worker.py
@@ -135,7 +135,7 @@ def create_megatron_test_config(
             "bias_activation_fusion": True,
             "moe_per_layer_logging": False,
             "moe_enable_deepep": False,
-            "moe_token_dispatcher_type": "allgather",
+            "moe_token_dispatcher_type": "alltoall",
             "moe_shared_expert_overlap": False,
             "defer_fp32_logits": defer_fp32_logits,
             "train_iters": 100,  # Required for Megatron training
