perf: DeepEP interface in megatron backend

examples/configs/distillation_math.yaml

@@ -107,6 +107,9 @@ policy: &POLICY_BASE
         bias_activation_fusion: True
         defer_fp32_logits: False
         moe_per_layer_logging: False
+        moe_enable_deepep: false
+        moe_token_dispatcher_type: "allgather"
+        moe_shared_expert_overlap: false
 
         optimizer:
             optimizer: "adam"

examples/configs/distillation_math_megatron.yaml

@@ -59,6 +59,9 @@ policy: &POLICY_BASE
         bias_activation_fusion: True
         moe_per_layer_logging: False
         defer_fp32_logits: False
+        moe_enable_deepep: false
+        moe_token_dispatcher_type: "allgather"
+        moe_shared_expert_overlap: false
 
         optimizer:
             optimizer: "adam"

examples/configs/dpo.yaml

@@ -119,6 +119,9 @@ policy:
     bias_activation_fusion: True
     defer_fp32_logits: False
     moe_per_layer_logging: False
+    moe_enable_deepep: false
+    moe_token_dispatcher_type: "allgather"
+    moe_shared_expert_overlap: false
 
     optimizer:
       optimizer: "adam"

examples/configs/grpo_math_1B.yaml

@@ -116,6 +116,9 @@ policy:
     bias_activation_fusion: True
     defer_fp32_logits: False
     moe_per_layer_logging: False
+    moe_enable_deepep: false
+    moe_token_dispatcher_type: "allgather"
+    moe_shared_expert_overlap: false
 
     optimizer:
       optimizer: "adam"

examples/configs/grpo_math_1B_megatron.yaml

@@ -94,6 +94,9 @@ policy:
     moe_router_load_balancing_type: "none" # "seq_aux_loss" causes logprob error divergence for grpo
     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_shared_expert_overlap: false
     #gives ~20% training perf speedup with sequence packing 
     apply_rope_fusion: True
 

examples/configs/recipes/llm/grpo-dapomath17k-dsv3-megatron.yaml

@@ -29,6 +29,8 @@ policy:
     sequence_parallel: true
     moe_permute_fusion: true
     apply_rope_fusion: false
+    moe_enable_deepep: true
+    moe_token_dispatcher_type: flex
     optimizer:
       lr: 5.0e-07
       min_lr: 5.0e-08

examples/configs/recipes/llm/grpo-moonlight-16ba3b-4n8g-megatron.yaml

@@ -31,6 +31,8 @@ policy:
       lr: 1.0e-06
     scheduler:
       lr_warmup_iters: 50
+    moe_enable_deepep: true
+    moe_token_dispatcher_type: flex
 logger:
   monitor_gpus: false
   wandb:

examples/configs/sft.yaml

@@ -113,6 +113,9 @@ policy:
     bias_activation_fusion: True
     defer_fp32_logits: False
     moe_per_layer_logging: False
+    moe_enable_deepep: false
+    moe_token_dispatcher_type: "allgather"
+    moe_shared_expert_overlap: false
 
     peft:
       enabled: false

examples/configs/sft_openmathinstruct2_megatron.yaml

@@ -92,6 +92,9 @@ policy:
     # gives ~25% training perf speedup with sequence packing and apply_rope_fusion
     bias_activation_fusion: True
     moe_per_layer_logging: False
+    moe_enable_deepep: false
+    moe_token_dispatcher_type: "allgather"
+    moe_shared_expert_overlap: false
 
     env_vars:
       PYTORCH_CUDA_ALLOC_CONF: "expandable_segments:False"

examples/configs/vlm_grpo_3B.yaml

@@ -104,6 +104,9 @@ policy:
     bias_activation_fusion: True
     defer_fp32_logits: False
     moe_per_layer_logging: False
+    moe_enable_deepep: false
+    moe_token_dispatcher_type: "allgather"
+    moe_shared_expert_overlap: false
 
     optimizer:
       optimizer: "adam"

examples/configs/vlm_grpo_3B_megatron.yaml

@@ -146,6 +146,9 @@ policy:
     bias_activation_fusion: True
     defer_fp32_logits: False
     moe_per_layer_logging: False
+    moe_enable_deepep: false
+    moe_token_dispatcher_type: "allgather"
+    moe_shared_expert_overlap: false
     optimizer:
       optimizer: adam
       lr: 2.0e-07

examples/nemo_gym/grpo_workplace_assistant_nemotron_nano_v2_9b.yaml

@@ -100,6 +100,9 @@ policy:
     apply_rope_fusion: True
     defer_fp32_logits: false
     moe_permute_fusion: false
+    moe_enable_deepep: false
+    moe_token_dispatcher_type: "allgather"
+    moe_shared_expert_overlap: false
 
     optimizer:
       optimizer: "adam"

nemo_rl/models/policy/__init__.py

@@ -183,6 +183,16 @@ class MegatronConfig(TypedDict):
     # Force overwrite of the initial checkpoint even if it exists (default: False)
     force_overwrite_initial_ckpt: NotRequired[bool]
     moe_per_layer_logging: bool
+    # Set to true to enable DeepEP for expert parallel communication
+    # 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'.
+    # Options are 'allgather','alltoall' and 'flex'
+    # Use 'flex' when using DeepEP
+    moe_token_dispatcher_type: str
+    # Can be used only with 'alltoall' token dispatcher
+    moe_shared_expert_overlap: bool
     optimizer: MegatronOptimizerConfig
     scheduler: MegatronSchedulerConfig
     distributed_data_parallel_config: MegatronDDPConfig

nemo_rl/models/policy/workers/megatron_policy_worker.py

@@ -658,6 +658,13 @@ def __init__(
         model_cfg.moe_router_bias_update_rate = self.cfg["megatron_cfg"][
             "moe_router_bias_update_rate"
         ]
+        model_cfg.moe_enable_deepep = self.cfg["megatron_cfg"]["moe_enable_deepep"]
+        model_cfg.moe_token_dispatcher_type = self.cfg["megatron_cfg"][
+            "moe_token_dispatcher_type"
+        ]
+        model_cfg.moe_shared_expert_overlap = self.cfg["megatron_cfg"][
+            "moe_shared_expert_overlap"
+        ]
 
         model_cfg.moe_permute_fusion = self.cfg["megatron_cfg"]["moe_permute_fusion"]
         if "layernorm_epsilon" in self.cfg["megatron_cfg"]:

pyproject.toml

@@ -100,6 +100,7 @@ mcore = [
   # https://github.com/NVIDIA/TransformerEngine/blob/v2.3/transformer_engine/pytorch/attention/dot_product_attention/utils.py#L108
   # https://github.com/facebookresearch/xformers/blob/8354497deb2c04c67fbb2e2ad911e86530da0e90/xformers/ops/fmha/flash.py#L76
   "flash-attn==2.8.1",
+  "deep_ep @ git+https://github.com/deepseek-ai/DeepEP.git@bfded34800dfec415b71503f8205181de90b2480",
   # Remove this once https://github.com/NVIDIA-NeMo/RL/issues/501 resolved
   "vllm==0.11.2",
 ]

tests/unit/models/generation/test_vllm_generation.py

@@ -178,6 +178,9 @@ def get_basic_megatron_test_config(
             "moe_router_load_balancing_type": "none",
             "moe_router_bias_update_rate": 0.0,
             "moe_permute_fusion": False,
+            "moe_enable_deepep": False,
+            "moe_token_dispatcher_type": "allgather",
+            "moe_shared_expert_overlap": False,
             "apply_rope_fusion": True,
             "bias_activation_fusion": True,
             "moe_per_layer_logging": False,

tests/unit/models/policy/test_megatron_worker.py

@@ -135,6 +135,9 @@ def create_megatron_test_config(
             "apply_rope_fusion": True,
             "bias_activation_fusion": True,
             "moe_per_layer_logging": False,
+            "moe_enable_deepep": False,
+            "moe_token_dispatcher_type": "allgather",
+            "moe_shared_expert_overlap": False,
             "defer_fp32_logits": defer_fp32_logits,
             "train_iters": 100,  # Required for Megatron training
             "optimizer": {