Add HybridEP support for GB200 NVL72

- Add deep_ep aarch64 dependency (7febc6e2, hybrid-ep branch)
- Add HybridEP setup in megatron setup.py (IMEX env vars, NVLink domain config)
- Add Qwen3-30B-A3B 4n4g config with moe_flex_dispatcher_backend=hybridep
- Add EP=4, EP=8, sms16 config variants for ablation testing
- Add test scripts for EP variants
- Update Dockerfile for aarch64 deep_ep build support
- Add HybridEP settings to 235B and DeepSeek-V3 performance configs

Signed-off-by: sna <sna@nvidia.com>

Author: seonjinn

docker/Dockerfile

@@ -112,6 +112,10 @@ ENV UV_LINK_MODE=copy
 # Ensure DeepEP is built for H100 and B200 (also mcore inference unified memory API now invokes a torch API that requires these to be set)
 ENV TORCH_CUDA_ARCH_LIST="9.0 10.0"
 
+# Ensure HybridEP/DeepEP JIT compilation can find nvcc at runtime
+ENV CUDA_HOME=/usr/local/cuda
+
+
 # First copy only the dependency files
 COPY --from=nemo-rl pyproject.toml uv.lock ./
 # Copy in the top level __init__.py/package_info.py since build-custom-vllm.sh needs the nemo_rl package to exist.

docker/Dockerfile.ngc_pytorch

@@ -90,6 +90,10 @@ ENV PATH="/opt/nemo_rl_venv/bin:$PATH"
 # Ensure DeepEP is built for H100 and B200
 ENV TORCH_CUDA_ARCH_LIST="9.0 10.0"
 
+# Ensure HybridEP/DeepEP JIT compilation can find nvcc at runtime
+ENV CUDA_HOME=/usr/local/cuda
+
+
 # First copy only the dependency files
 COPY --from=nemo-rl pyproject.toml uv.lock ./
 COPY --from=nemo-rl --link 3rdparty/ ./3rdparty/

examples/configs/recipes/llm/performance/grpo-deepseek-v3-32n4g.yaml

@@ -9,6 +9,10 @@ policy:
     expert_model_parallel_size: 16
     num_layers_in_first_pipeline_stage: 7
     num_layers_in_last_pipeline_stage: 6
+    # HybridEP settings
+    moe_token_dispatcher_type: flex
+    moe_flex_dispatcher_backend: hybridep
+    moe_hybridep_num_sms: 32
   generation:
     vllm_cfg:
       tensor_parallel_size: 32

examples/configs/recipes/llm/performance/grpo-qwen3-235b-16n4g.yaml

@@ -6,6 +6,10 @@ policy:
     pipeline_model_parallel_size: 4
     num_layers_in_first_pipeline_stage: 23
     num_layers_in_last_pipeline_stage: 23
+    # HybridEP settings
+    moe_token_dispatcher_type: flex
+    moe_flex_dispatcher_backend: hybridep
+    moe_hybridep_num_sms: 32
   generation:
     vllm_cfg:
       tensor_parallel_size: 8

examples/configs/recipes/llm/performance/grpo-qwen3-30ba3b-4n4g-ep4.yaml

@@ -0,0 +1,10 @@
+defaults: ./grpo-qwen3-30ba3b-4n4g.yaml
+checkpointing:
+  checkpoint_dir: results/grpo-qwen3-30ba3b-4n4g-ep4
+policy:
+  megatron_cfg:
+    expert_model_parallel_size: 4
+logger:
+  log_dir: logs/grpo-qwen3-30ba3b-4n4g-ep4
+  wandb:
+    name: grpo-qwen3-30ba3b-4n4g-ep4

examples/configs/recipes/llm/performance/grpo-qwen3-30ba3b-4n4g-ep8.yaml

@@ -0,0 +1,10 @@
+defaults: ./grpo-qwen3-30ba3b-4n4g.yaml
+checkpointing:
+  checkpoint_dir: results/grpo-qwen3-30ba3b-4n4g-ep8
+policy:
+  megatron_cfg:
+    expert_model_parallel_size: 8
+logger:
+  log_dir: logs/grpo-qwen3-30ba3b-4n4g-ep8
+  wandb:
+    name: grpo-qwen3-30ba3b-4n4g-ep8

examples/configs/recipes/llm/performance/grpo-qwen3-30ba3b-4n4g-sms16.yaml

@@ -0,0 +1,10 @@
+defaults: ./grpo-qwen3-30ba3b-4n4g.yaml
+checkpointing:
+  checkpoint_dir: results/grpo-qwen3-30ba3b-4n4g-sms16
+policy:
+  megatron_cfg:
+    moe_hybridep_num_sms: 16
+logger:
+  log_dir: logs/grpo-qwen3-30ba3b-4n4g-sms16
+  wandb:
+    name: grpo-qwen3-30ba3b-4n4g-sms16

examples/configs/recipes/llm/performance/grpo-qwen3-30ba3b-4n4g.yaml

@@ -21,6 +21,10 @@ policy:
     pipeline_model_parallel_size: 1
     expert_model_parallel_size: 16
     sequence_parallel: false
+    # HybridEP settings
+    moe_token_dispatcher_type: flex
+    moe_flex_dispatcher_backend: hybridep
+    moe_hybridep_num_sms: 32
     optimizer:
       lr: 3.0e-07
       min_lr: 3.0e-08
@@ -42,4 +46,3 @@ logger:
 cluster:
   gpus_per_node: 4
   num_nodes: 4
-

nemo_rl/models/megatron/setup.py

@@ -477,6 +477,49 @@ def _apply_moe_config(model_cfg: Any, config: PolicyConfig) -> None:
         "moe_shared_expert_overlap"
     ]
 
+    # HybridEP settings for MoE expert parallelism
+    # See: https://github.com/deepseek-ai/DeepEP/tree/hybrid-ep
+    if "moe_flex_dispatcher_backend" in config["megatron_cfg"]:
+        model_cfg.moe_flex_dispatcher_backend = config["megatron_cfg"][
+            "moe_flex_dispatcher_backend"
+        ]
+    if "moe_hybridep_num_sms" in config["megatron_cfg"]:
+        model_cfg.moe_hybridep_num_sms = config["megatron_cfg"]["moe_hybridep_num_sms"]
+
+    # HybridEP environment variables
+    # These are required by DeepEP's hybrid-ep branch for NVLink domain configuration.
+    # Users can set them explicitly via config, or they will be auto-computed with a warning.
+    if config["megatron_cfg"].get("moe_flex_dispatcher_backend") == "hybridep":
+        ep_size = model_cfg.expert_model_parallel_size
+
+        # NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN
+        if "hybridep_num_ranks_per_nvlink_domain" in config["megatron_cfg"]:
+            val = config["megatron_cfg"]["hybridep_num_ranks_per_nvlink_domain"]
+            os.environ["NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN"] = str(val)
+        elif "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN" not in os.environ:
+            default_val = min(ep_size, 64)
+            os.environ["NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN"] = str(default_val)
+            warnings.warn(
+                f"HybridEP: NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN not configured. "
+                f"Auto-setting to min(expert_model_parallel_size={ep_size}, 64) = {default_val}. "
+                f"Set 'hybridep_num_ranks_per_nvlink_domain' in megatron_cfg to override.",
+                stacklevel=2,
+            )
+
+        # USE_MNNVL
+        if "hybridep_use_mnnvl" in config["megatron_cfg"]:
+            val = config["megatron_cfg"]["hybridep_use_mnnvl"]
+            os.environ["USE_MNNVL"] = str(int(val))
+        elif "USE_MNNVL" not in os.environ:
+            default_val = int(ep_size > 4)
+            os.environ["USE_MNNVL"] = str(default_val)
+            warnings.warn(
+                f"HybridEP: USE_MNNVL not configured. "
+                f"Auto-setting to int(expert_model_parallel_size={ep_size} > 4) = {default_val}. "
+                f"Set 'hybridep_use_mnnvl' in megatron_cfg to override.",
+                stacklevel=2,
+            )
+
     model_cfg.moe_permute_fusion = config["megatron_cfg"]["moe_permute_fusion"]
 
 

nemo_rl/models/policy/__init__.py

@@ -236,6 +236,14 @@ class MegatronConfig(TypedDict):
     moe_token_dispatcher_type: str
     # Can be used only with 'alltoall' token dispatcher
     moe_shared_expert_overlap: bool
+    # HybridEP settings for MoE expert parallelism (requires moe_token_dispatcher_type='flex')
+    # See: https://github.com/deepseek-ai/DeepEP/tree/hybrid-ep
+    moe_flex_dispatcher_backend: NotRequired[str]
+    moe_hybridep_num_sms: NotRequired[int]
+    # Number of HybridEP ranks per NVLink domain (default: min(expert_model_parallel_size, 64))
+    hybridep_num_ranks_per_nvlink_domain: NotRequired[int]
+    # Enable multi-node NVLink support (default: expert_model_parallel_size > 4)
+    hybridep_use_mnnvl: NotRequired[bool]
     peft: NotRequired[MegatronPeftConfig | MegatronPeftConfigDisabled]
     optimizer: MegatronOptimizerConfig
     scheduler: MegatronSchedulerConfig