perf: DeepEP interface in megatron backend

[guyueh1]

What does this PR do ?

Add interface to configure deep_ep usage in megatron backend

closes #1396

Dup of #1645

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

Release Notes

• New Features

  • Added new Mixture-of-Experts (MOE) configuration options: enable/disable DeepEP optimization, select token dispatcher type (allgather, alltoall, flex), and control shared expert overlap behavior for enhanced training flexibility.

• Chores

  • Added DeepEP as an optional dependency to support advanced MOE optimizations.

_{✏️ Tip: You can customize this high-level summary in your review settings.}

[coderabbitai]

📝 Walkthrough

Walkthrough

Adds three new MOE-related configuration fields (moe_enable_deepep, moe_token_dispatcher_type, moe_shared_expert_overlap) across multiple configuration files and extends the MegatronConfig TypedDict schema. Also wires these fields into the Megatron policy worker and adds DeepEP as an optional dependency.

Changes

Cohort / File(s)	Summary
Configuration Files - Distillation & Math Recipes examples/configs/distillation_math.yaml, examples/configs/distillation_math_megatron.yaml, examples/configs/dpo.yaml, examples/configs/sft.yaml, examples/configs/sft_openmathinstruct2_megatron.yaml	Added three new MOE config fields with default values (moe_enable_deepep: false, moe_token_dispatcher_type: "allgather", moe_shared_expert_overlap: false) under policy.megatron_cfg
Configuration Files - GRPO Recipes examples/configs/grpo_math_1B.yaml, examples/configs/grpo_math_1B_megatron.yaml, examples/configs/recipes/llm/grpo-*.yaml	Added two to three new MOE config fields under policy.megatron_cfg; some recipes use moe_enable_deepep: true and moe_token_dispatcher_type: flex instead of defaults
Configuration Files - VLM & Assistant Recipes examples/configs/vlm_grpo_3B.yaml, examples/configs/vlm_grpo_3B_megatron.yaml, examples/nemo_gym/grpo_workplace_assistant_nemotron_nano_v2_9b.yaml	Added three new MOE config fields with default values under policy.megatron_cfg
Schema Definition nemo_rl/models/policy/__init__.py	Extended MegatronConfig TypedDict with three new fields: moe_enable_deepep (bool), moe_token_dispatcher_type (str with options "allgather"/"alltoall"/"flex"), moe_shared_expert_overlap (bool), including detailed docstring comments
Policy Worker Implementation nemo_rl/models/policy/workers/megatron_policy_worker.py	Reads three new MOE config fields from config.megatron_cfg and assigns them to modelCfg during initialization
Dependency Management pyproject.toml	Added DeepEP git dependency to both fsdp and mcore optional dependency groups
Test Configuration tests/unit/models/generation/test_vllm_generation.py, tests/unit/models/policy/test_megatron_worker.py	Added three new MOE config fields to test configurations with default values

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~12 minutes

Possibly related PRs

• feat: Add moe load balancing metrics #1520: Adds and wires up additional MOE configuration behavior into the Megatron policy surface by modifying MegatronConfig and megatron_policy_worker.
• feat: add Megatron support for on-policy distillation #1324: Performs large Megatron config work overlapping the same distillation YAML configuration files touched by this PR.
• feat: Deepseek migration to Megatron-Bridge + CP support #1059: Modifies Megatron policy worker and config surface to integrate Deepseek-related MOE features alongside other Megatron enhancements.

Suggested labels

CI:L1

Suggested reviewers

• terrykong
• yaoyu-33
• ashors1

🚥 Pre-merge checks | ✅ 4 | ❌ 2

❌ Failed checks (2 warnings)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 66.67% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.
Test Results For Major Changes	⚠️ Warning	PR adds DeepEP support feature but lacks test results, performance benchmarks, convergence validation, and has missing keys in exemplar configs raising runtime stability concerns.	Include test results demonstrating feature correctness, provide performance numbers for DeepSeekv3 training, mark new config keys as NotRequired, use .get() with defaults, and ensure all exemplar configs include these keys.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title 'perf: DeepEP interface in megatron backend' accurately summarizes the main change—adding DeepEP configuration interface to the Megatron backend.
Linked Issues check	✅ Passed	The PR addresses issue #1396 by adding DeepEP configuration support across Megatron backend through config keys, TypedDict updates, and policy worker integration.
Out of Scope Changes check	✅ Passed	All changes are directly related to DeepEP interface implementation: configuration keys, TypedDict schema updates, policy worker integration, and dependency addition.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing touches

• 📝 Generate docstrings

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

Actionable comments posted: 1

🤖 Fix all issues with AI agents

In `@nemo_rl/models/policy/__init__.py`:
- Around line 186-195: Update the TypedDict in nemo_rl/models/policy/__init__.py
to mark moe_enable_deepep, moe_token_dispatcher_type, and
moe_shared_expert_overlap as NotRequired and document recommended defaults
(e.g., False, 'allgather', False); then change the access in
megatron_policy_worker.py (around the logic at lines ~661–667) to use
config.get('moe_enable_deepep', False), config.get('moe_token_dispatcher_type',
'allgather'), and config.get('moe_shared_expert_overlap', False) so missing keys
won’t KeyError; finally add those three keys with the recommended default values
to the exemplar YAMLs (grpo_math_70B_megatron.yaml, grpo_math_8B_megatron.yaml,
grpo_math_qwen30ba3b_megatron.yaml).

[guyueh1]

@terrykong could you review? This is needed for a recent deepseek performance study urgently

[terrykong]

Lgtm. Can you please resolve the two comments?

Fyi: @yuki-97