feat: Add HybridEP support for MoE expert parallelism

[seonjinn]

• Update DeepEP dependency to hybrid-ep branch for HybridEP support
  • automodel, vllm, mcore dependency groups updated
• Add HybridEP configuration options in _apply_moe_config():
  • moe_flex_dispatcher_backend: Flex dispatcher backend (e.g., 'hybridep')
  • moe_hybridep_num_sms: Number of SMs for HybridEP operations

Usage in config:
policy.megatron_cfg.moe_token_dispatcher_type=flex
policy.megatron_cfg.moe_flex_dispatcher_backend=hybridep
policy.megatron_cfg.moe_hybridep_num_sms=32

See: https://github.com/deepseek-ai/DeepEP/tree/hybrid-ep

What does this PR do ?

Add a one line overview of what this PR aims to accomplish.

Issues

List issues that this PR closes (syntax):

Usage

• You can potentially add a usage example below

# Add a code snippet demonstrating how to use this

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you run the unit tests and functional tests locally? Visit our Testing Guide for how to run tests
• Did you add or update any necessary documentation? Visit our Document Development Guide for how to write, build and test the docs.

Additional Information

• ...

Summary by CodeRabbit

• New Features

  • Added configuration support for additional Mixture of Experts (MoE) model parameters, including dispatcher backend and HybridEP settings.

• Dependencies

  • Updated DeepEP dependency reference to use the hybrid-ep branch across multiple dependency groups.

[coderabbitai]

Walkthrough

The changes add two conditional configuration hooks to the MoE setup function for optional dispatcher and HybridEP parameters, and update the DeepEP dependency reference from a specific commit to the hybrid-ep branch across multiple dependency groups in the project configuration.

Changes

Cohort / File(s)	Summary
MoE Configuration Hooks nemo_rl/models/megatron/setup.py	Adds conditional assignments for moe_flex_dispatcher_backend and moe_hybridep_num_sms configuration parameters in the _apply_moe_config function.
Dependency Updates pyproject.toml	Updates DeepEP git dependency reference from commit bfded34800dfec415b71503f8205181de90b2480 to the hybrid-ep branch across automodel, vllm, and mcore dependency groups with explanatory comments.

Estimated code review effort

🎯 1 (Trivial) | ⏱️ ~5 minutes

✨ Finishing Touches

📝 Generate docstrings

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch sj/hybridep-support

[coderabbitai]

Actionable comments posted: 1

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

pyproject.toml (1)

324-327: ⚠️ Potential issue | 🟠 Major

Stale dependency-metadata version for deep_ep.

The dependency is pinned to the hybrid-ep branch (which dynamically generates its version from the current commit hash via git rev-parse --short HEAD), but the dependency-metadata version is statically set to v1.2.1+bfded34. This means the metadata version will become stale whenever the branch advances, potentially causing uv resolver failures.

Either:

1. Update to pin to a specific commit hash instead of a branch, or
2. Update the metadata version to match the current HEAD of hybrid-ep and regenerate it whenever the dependency updates

🤖 Fix all issues with AI agents

In `@nemo_rl/models/megatron/setup.py`:
- Around line 405-412: The new runtime keys moe_flex_dispatcher_backend and
moe_hybridep_num_sms are missing from the MegatronConfig TypedDict and from
example configs; add both to the MegatronConfig definition in
nemo_rl/models/policy/__init__.py as NotRequired entries (use the exact symbol
name MegatronConfig) with short docstrings: "Backend type for MoE flex
dispatcher (HybridEP)" for moe_flex_dispatcher_backend and "Number of SMs for
HybridEP" for moe_hybridep_num_sms, and then update at least one exemplar YAML
in examples/configs (e.g., a megatron MoE config) to include these keys with
sensible defaults (recommended defaults) so they are documented and visible to
users.

🧹 Nitpick comments (1)

pyproject.toml (1)

70-72: Branch ref instead of pinned commit reduces build reproducibility.

All three dependency groups now point to @hybrid-ep (a moving branch) instead of a fixed commit hash. This means builds are not reproducible — a force-push or new commit on that branch silently changes what gets installed. Consider pinning to a specific commit on the hybrid-ep branch once it stabilizes.

[terrykong]

mostly lgtm @seonjinn , but the main thing i'm concerned about is the introduction of conda, we can continue discussing on the thread is tarted

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[github-actions]

❌ Submodule Fast-Forward Check Failed

Check based on commit: 0349e83 (PR #1942 from sj/hybridep-support)

❌ Submodules that need attention:

Megatron-Bridge: ❌ PR branch is BEHIND main branch
TARGET (main branch): https://github.com/NVIDIA-NeMo/Megatron-Bridge/commits/7110a964272a5c74dcb6b680b691087e190c220c/
CURRENT (PR #1942 from sj/hybridep-support): https://github.com/NVIDIA-NeMo/Megatron-Bridge/commits/a2bb70b91b827bd6b085a77442c7cf60cfdb59fe/

Megatron-LM: ❌ PR branch is BEHIND main branch
TARGET (main branch): https://github.com/NVIDIA/Megatron-LM/commits/17a67b9a97fb11a75933fd7f76ad76e1ac98a53d/
CURRENT (PR #1942 from sj/hybridep-support): https://github.com/NVIDIA/Megatron-LM/commits/9e2810417315a7ee93b41d4e234454abd3c16af5/

Please ensure all submodule commits are fast-forwards of the main branch before merging.

[github-actions]

❌ Submodule Fast-Forward Check Failed

Check based on commit: ec53cb6 (PR #1942 from sj/hybridep-support)

❌ Submodules that need attention:

Megatron-Bridge: ❌ PR branch is BEHIND main branch
TARGET (main branch): https://github.com/NVIDIA-NeMo/Megatron-Bridge/commits/7110a964272a5c74dcb6b680b691087e190c220c/
CURRENT (PR #1942 from sj/hybridep-support): https://github.com/NVIDIA-NeMo/Megatron-Bridge/commits/a2bb70b91b827bd6b085a77442c7cf60cfdb59fe/

Megatron-LM: ❌ PR branch is BEHIND main branch
TARGET (main branch): https://github.com/NVIDIA/Megatron-LM/commits/17a67b9a97fb11a75933fd7f76ad76e1ac98a53d/
CURRENT (PR #1942 from sj/hybridep-support): https://github.com/NVIDIA/Megatron-LM/commits/9e2810417315a7ee93b41d4e234454abd3c16af5/

Please ensure all submodule commits are fast-forwards of the main branch before merging.

[seonjinn]

/ok to test 18f3318

[seonjinn]

/ok to test 9bb1778

[seonjinn]

/ok to test a630e9e

[seonjinn]

/ok to test c2fe685

[seonjinn]

/ok to test 37794a5

[seonjinn]

/ok to test 37794a5

[seonjinn]

/ok to test 1ed0c10

[seonjinn]

/ok to test b9517fa

[seonjinn]

/ok to test a4f0ead

[seonjinn]

Empty retry commit to clear known fp8 numeric flake on L0_Unit_Tests_Policy::test_megatron_policy_training[2gpu_tp2_llama_fp8].

Flake evidence (same PR branch, alternating pass/fail):

SHA	L0_Policy
7008449 (05-07)	PASS
c8f0121 (05-11)	PASS
9bb1778 (05-14)	FAIL
a630e9e (05-15)	PASS
b9517fa (05-15)	FAIL — losses[0]=1.17204e-4 vs losses[-1]=1.17239e-4 (3e-8 abs, fp8 noise floor)

Same code, alternating results → tolerance flake, not a regression. The PR touches DeepEP submodule + monkey-patches, not Megatron policy training paths.

[guyueh1]

@terrykong to take a look, i think this PR is ready

[terrykong]

lgtm, with the exception of that differing deep EP commit. asked @seonjinn if that was intentional (if not we should document how to find each commit)

[terrykong]

Filed #2522 to track unifying the duplicated deep_ep commit pins in pyproject.toml — the same x86_64/aarch64 hashes are now repeated across 4 dependency groups (automodel, vllm, mcore, override-dependencies). That cleanup can land separately.