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.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

pyproject.toml

@@ -71,15 +71,17 @@ automodel = [
   "causal-conv1d",
   "nv-grouped-gemm",
   "transformer-engine[pytorch]>=2.9.0a0,<2.12.0",
-  "deep_ep @ git+https://github.com/deepseek-ai/DeepEP.git@bfded34800dfec415b71503f8205181de90b2480",
+  "deep_ep @ git+https://github.com/deepseek-ai/DeepEP.git@bfded34800dfec415b71503f8205181de90b2480 ; platform_machine == 'x86_64'",
+  "deep_ep @ git+https://github.com/deepseek-ai/DeepEP.git@7febc6e25660af0f54d95dd781ecdcd62265ecca ; platform_machine == 'aarch64'",
 ]
 vllm = [
   "cuda-python",
   "deep_gemm @ git+https://github.com/deepseek-ai/DeepGEMM.git@7b6b5563b9d4c1ae07ffbce7f78ad3ac9204827c",
   # deep_ep also needs libibverbs-dev
   # sudo apt-get update
   # sudo apt-get install libibverbs-dev
-  "deep_ep @ git+https://github.com/deepseek-ai/DeepEP.git@bfded34800dfec415b71503f8205181de90b2480",
+  "deep_ep @ git+https://github.com/deepseek-ai/DeepEP.git@bfded34800dfec415b71503f8205181de90b2480 ; platform_machine == 'x86_64'",
+  "deep_ep @ git+https://github.com/deepseek-ai/DeepEP.git@7febc6e25660af0f54d95dd781ecdcd62265ecca ; platform_machine == 'aarch64'",
   "vllm==0.17.1",
   "num2words>=0.5.14",
   "flashinfer-python==0.6.4",
@@ -107,7 +109,8 @@ mcore = [
   # https://github.com/facebookresearch/xformers/blob/8354497deb2c04c67fbb2e2ad911e86530da0e90/xformers/ops/fmha/flash.py#L76
   "flash-attn==2.8.1",
   "emerging-optimizers==0.1.0",
-  "deep_ep @ git+https://github.com/deepseek-ai/DeepEP.git@bfded34800dfec415b71503f8205181de90b2480",
+  "deep_ep @ git+https://github.com/deepseek-ai/DeepEP.git@bfded34800dfec415b71503f8205181de90b2480 ; platform_machine == 'x86_64'",
+  "deep_ep @ git+https://github.com/deepseek-ai/DeepEP.git@7febc6e25660af0f54d95dd781ecdcd62265ecca ; platform_machine == 'aarch64'",
 ]
 nemo_gym = ["nemo_gym"]
 
@@ -247,7 +250,8 @@ override-dependencies = [
   "llguidance>=1.3.0,<1.4.0",
   # Override setuptools range in other dependencies to address CVE GHSA-58pv-8j8x-9vj2
   "setuptools>=80.10.2",
-  "deep_ep @ git+https://github.com/deepseek-ai/DeepEP.git@bfded34800dfec415b71503f8205181de90b2480",
+  "deep_ep @ git+https://github.com/deepseek-ai/DeepEP.git@bfded34800dfec415b71503f8205181de90b2480 ; platform_machine == 'x86_64'",
+  "deep_ep @ git+https://github.com/deepseek-ai/DeepEP.git@7febc6e25660af0f54d95dd781ecdcd62265ecca ; platform_machine == 'aarch64'",
   # Pin flashinfer globally — flashinfer-python must match flashinfer-cubin at runtime, and
   # they're resolved independently by uv so ranges risk version mismatch.
   # When changing this version, check what each backend expects:

tests/test_suites/llm/performance/grpo-qwen3-235b-16n4g.sh

@@ -23,6 +23,7 @@ uv run examples/run_grpo.py \
     logger.wandb_enabled=True \
     logger.wandb.project=nemo-rl \
     logger.wandb.name=$EXP_NAME \
+    logger.tensorboard_enabled=True \
     logger.monitor_gpus=True \
     checkpointing.enabled=True \
     checkpointing.checkpoint_dir=$CKPT_DIR \