no token dispatcher

Signed-off-by: Malay Nagda <malayn@nvidia.com>

Author: malay-nagda

scripts/performance/argument_parser.py

@@ -683,21 +683,13 @@ def parse_cli_args():
         help="Comma separated list of modules to recompute. Defaults to None",
         required=False,
     )
-    performance_args.add_argument(
-        "--moe_token_dispatcher_type",
-        type=str,
-        help="MoE token dispatcher type. Options- alltoall, flex. Defaults to None",
-        choices=["alltoall", "flex", "allgather"],
-        required=False,
-        default=None,
-    )
     performance_args.add_argument(
         "--moe_flex_dispatcher_backend",
-        type=str,
+        type=lambda x: None if x == "None" else x,
         help="MoE flex dispatcher backend. Options- deepep, hybridep. Defaults to None",
-        choices=["deepep", "hybridep"],
+        choices=["deepep", "hybridep", None],
         required=False,
-        default=None,
+        default=-1,
     )
 
     # Logging

scripts/performance/configs/gpt_oss/gpt_oss_workload_base_configs.py

@@ -35,7 +35,6 @@
     expert_tensor_parallel_size=1,
     global_batch_size=512,
     micro_batch_size=1,
-    moe_token_dispatcher_type="alltoall",
 )
 
 

scripts/performance/configs/nemotronh/nemotronh_workload_base_configs.py

@@ -30,7 +30,6 @@
     num_gpus=64,
     global_batch_size=192,
     cuda_graph_impl="transformer_engine",
-    moe_token_dispatcher_type="alltoall",
 )
 
 

scripts/performance/configs/qwen/qwen3_workload_base_configs.py

@@ -116,7 +116,6 @@
     expert_model_parallel_size=8,
     global_batch_size=1024,
     moe_a2a_overlap=False,
-    moe_token_dispatcher_type="alltoall",
 )
 
 
@@ -127,7 +126,6 @@
     expert_model_parallel_size=8,
     global_batch_size=1024,
     moe_a2a_overlap=False,
-    moe_token_dispatcher_type="alltoall",
 )
 
 
@@ -143,7 +141,6 @@
     expert_model_parallel_size=8,
     global_batch_size=1024,
     moe_a2a_overlap=False,
-    moe_token_dispatcher_type="alltoall",
 )
 
 
@@ -154,7 +151,6 @@
     expert_model_parallel_size=8,
     global_batch_size=1024,
     moe_a2a_overlap=False,
-    moe_token_dispatcher_type="alltoall",
 )
 
 
@@ -171,7 +167,6 @@
     expert_model_parallel_size=32,
     global_batch_size=2048,
     moe_a2a_overlap=True,
-    moe_token_dispatcher_type="alltoall",
 )
 
 
@@ -184,7 +179,6 @@
     expert_model_parallel_size=32,
     global_batch_size=2048,
     moe_a2a_overlap=True,
-    moe_token_dispatcher_type="alltoall",
 )
 
 
@@ -466,39 +460,41 @@
     num_gpus=128,
     expert_model_parallel_size=128,
     micro_batch_size=1,
-    moe_token_dispatcher_type="alltoall",
 )
 
 QWEN3_NEXT_80B_A3B_PRETRAIN_CONFIG_H100_BF16_V1 = replace(
     BASE_QWEN3_NEXT_80B_A3B_CONFIG,
     num_gpus=128,
     expert_model_parallel_size=128,
     micro_batch_size=1,
-    moe_token_dispatcher_type="alltoall",
 )
 
 QWEN3_NEXT_80B_A3B_PRETRAIN_CONFIG_B300_FP8_MX_V1 = replace(
     BASE_QWEN3_NEXT_80B_A3B_CONFIG,
     num_gpus=64,
     micro_batch_size=2,
+    moe_flex_dispatcher_backend="deepep",
 )
 
 QWEN3_NEXT_80B_A3B_PRETRAIN_CONFIG_B300_BF16_V1 = replace(
     BASE_QWEN3_NEXT_80B_A3B_CONFIG,
     num_gpus=64,
     micro_batch_size=1,
+    moe_flex_dispatcher_backend="deepep",
 )
 
 QWEN3_NEXT_80B_A3B_PRETRAIN_CONFIG_B200_FP8_MX_V1 = replace(
     BASE_QWEN3_NEXT_80B_A3B_CONFIG,
     num_gpus=64,
     micro_batch_size=1,
+    moe_flex_dispatcher_backend="deepep",
 )
 
 QWEN3_NEXT_80B_A3B_PRETRAIN_CONFIG_B200_BF16_V1 = replace(
     BASE_QWEN3_NEXT_80B_A3B_CONFIG,
     num_gpus=64,
     micro_batch_size=1,
+    moe_flex_dispatcher_backend="deepep",
 )
 
 

scripts/performance/utils/overrides.py

@@ -222,10 +222,10 @@ def set_workload_base_configs(cfg: ConfigContainer, settings: WorkloadBaseConfig
     )
     _set_common_perf_overrides(cfg)
 
-    if settings.moe_token_dispatcher_type is not None:
-        cfg.model.moe_token_dispatcher_type = settings.moe_token_dispatcher_type
     if settings.moe_flex_dispatcher_backend is not None:
         apply_flex_dispatcher_backend(cfg.model, settings.moe_flex_dispatcher_backend)
+    else:
+        cfg.model.moe_token_dispatcher_type = "alltoall"
 
     return cfg
 
@@ -434,10 +434,10 @@ def set_user_overrides(recipe: ConfigContainer, args: argparse.Namespace) -> Con
     if args.pytorch_profiler:
         recipe.logger.tensorboard_dir = "/nemo_run/pytorch_profile"
 
-    if args.moe_token_dispatcher_type is not None:
-        recipe.model.moe_token_dispatcher_type = args.moe_token_dispatcher_type
     if args.moe_flex_dispatcher_backend is not None:
         apply_flex_dispatcher_backend(recipe.model, args.moe_flex_dispatcher_backend)
+    else:
+        recipe.model.moe_token_dispatcher_type = "alltoall"
 
     return recipe
 

scripts/performance/utils/utils.py

@@ -59,7 +59,6 @@ class WorkloadBaseConfig:
 
     # MoE configuration
     moe_flex_dispatcher_backend: Optional[str] = None
-    moe_token_dispatcher_type: Optional[str] = None
     moe_a2a_overlap: Optional[bool] = False
     peft: Optional[str] = None