perf: use HybridEP in AutoModel MoE recipes by seonjinn · Pull Request #4037 · NVIDIA-NeMo/RL

seonjinn · 2026-09-07T04:31:50Z

Summary

Use dispatcher: hybridep in AutoModel MoE recipes that previously selected DeepEP.
Replace the removed enable_deepep option with the explicit dispatcher setting.
Preserve existing expert backends and recipes that explicitly require the Torch dispatcher.

This is separate from #3438, which covers Megatron performance recipes.

Container requirement

Multi-node recipes require DeepEP to be built with HYBRID_EP_MULTINODE=1. The corresponding container build change is tracked separately.

Testing

18/18 resolved AutoModel recipe cases passed.
Verified inherited recipes preserve explicit dispatcher: torch overrides.
ruff check
ruff format --check
git diff --check

Runtime H100 validation will be added separately.

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

copy-pr-bot · 2026-09-07T04:31:54Z

Auto-sync is disabled for ready for review pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

seonjinn · 2026-09-08T01:12:32Z

/ok to test abfcffd

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

seonjinn · 2026-09-12T02:46:50Z

/ok to test 3051348

Plain "deepep" hangs with "DeepEP error: timeout (dispatch CPU)" on this EP=32, 4-node/32-rank layout under the deep_ep pin bumped by NVIDIA-NeMo#3436 (deep_ep 29d31c0 -> 17cfb81, "enable HybridEP on H100/B200"). The container's deep_ep is already built with HYBRID_EP_MULTINODE=1 (NVIDIA-NeMo#4038, merged), and "hybridep" is already a supported nemo_automodel BackendConfig.dispatcher value -- just not yet the repo-wide default for AutoModel MoE recipes (that's NVIDIA-NeMo#4037, still open). Switching these two recipes now rather than waiting on NVIDIA-NeMo#4037: verified with a live 4n8g eagle3 run past the exact point (policy.get_logprobs()'s first post-refit call) that failed twice under plain deepep. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Signed-off-by: Yuekai Zhang <zhangyuekai@foxmail.com>

seonjinn · 2026-09-18T06:50:12Z

/ok to test 27ac693

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

seonjinn · 2026-09-18T08:57:09Z

/ok to test 8cd048d

seonjinn · 2026-09-19T01:31:26Z

/ok to test 2eb68fa

sharonyu-115

Reviewed the AutoModel enable_deepep -> dispatcher: hybridep migration. The dispatcher swap itself checks out (cross-referenced against Automodel's own upstream reference configs, e.g. DeepSeek V4 Flash matches exactly), and the new regression test is a good addition.

Left inline notes (on the earlier review) on: 3 recipes left experts: unset rather than pinning it explicitly (unlike 9 of the other 11 recipes in this PR) -- not claiming a specific value is wrong, since Automodel's own examples use both gmm and torch_mm for these same model families, just flagging the inconsistency; 2 more AutoModel MoE recipes that still use the removed enable_deepep flag and appear to need the same migration; and a couple of test-coverage gaps tied to both.

[E2E-VALIDATION-EVIDENCE] This PR swaps the MoE token dispatcher (deepep -> hybridep, or removes enable_deepep) across all 11 changed recipes. 10 of them are gated by automated convergence checks, not just launch-smoke tests:

tests/test_suites/nightly.txt: dapo-nanov3.5-30BA3B-4n8g-automodel, dpo-nanov3-30B3AB-1n8g-fsdp8ep8-automodel.v2, grpo-deepseek-v4-flash-0731-16n8g-automodel-cp8ep128, grpo-moonlight-16b-automodel-1n8g-ep8, grpo-nemotron3-super-120BA12B-16n8g-automodel-ep8.v2, grpo-qwen3.5-35ba3b-2n8g-automodel-ep16, sft-gpt-oss-20b-1n8g-fsdp8ep8-automodel
tests/test_suites/release.txt: grpo-glm47-flash-4n8g-automodel, grpo-minimax-m27-dapo-8n8g-automodel, grpo-qwen3.5-35ba3b-dapo-4n8g-automodel

For example, grpo-moonlight-16b-automodel-1n8g-ep8.sh asserts concrete thresholds via check_metrics.py (gen_kl_error < 0.001, reward["30"] > 0.4, grad_norm["30"] < 0.2) -- a real convergence gate, not a smoke test. The PR description's Testing section only covers config-resolution-level checks (the new unit test + ruff/format) and explicitly notes "Runtime H100 validation will be added separately." Since a dispatcher swap changes the communication/kernel path for these recipes, could you link the nightly/release run(s) against this commit (or a manual wandb comparison) once available, to confirm no regression on these metrics before merge?

Generated by Claude Code

Make expert backends explicit, migrate the remaining Nano recipes, and preserve HybridEP alignment through dynamic batching and Gemma CP2. Cover resolved recipe inheritance and intentional backend choices. Signed-off-by: seonjinn <sna@nvidia.com>

seonjinn · 2026-09-21T20:52:36Z

Review follow-up at 5cb48a7 (merged main fd7112c374af23b3bcc392637fb737495f8a3f31):

Addressed all six inline configuration/test requests; replies above describe each fix.
Preserved expert kernel choices and explicit Torch overrides; migrated the two missing Nano recipes.
Fixed input alignment, including both padding and dynamic-batch rounding for Gemma CP2.
Local checks: 131 passed, 242 skipped (28 PR recipe checks plus 103 related checks). Skips are unrelated recipes outside enabled AutoModel EP scope. Ruff lint/format and git diff --check passed.

These were CPU configuration checks on macOS, not GPU execution or convergence tests. The requested new-head nightly/release convergence evidence remains pending, and Lfast does not satisfy that requirement by itself. No convergence thresholds were changed.

Container clarification: the merged main Dockerfiles already build DeepEP with HYBRID_EP_MULTINODE=1; older images may still lack that option. This supersedes the PR body's older statement that the container change is separately pending.

Requesting CI:Lfast for this exact head next.

seonjinn · 2026-09-21T20:52:39Z

/ok to test 5cb48a7

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

seonjinn · 2026-09-23T07:54:35Z

/ok to test d72a905

…60921 Signed-off-by: Seonjin Na <sna@nvidia.com>

seonjinn · 2026-09-25T05:28:48Z

/ok to test ffc8e30

yuki-97

Scope: the 14 recipe changes and the new test_hybridep_automodel_recipes.py.

Checked:

Pinning experts: torch_mm on GLM / MiniMax / Moonlight / Super keeps the backend those recipes already ran with: the pinned Automodel ignores enable_deepep (utils.py:303-314) and defaults experts to torch_mm (utils.py:243-245).
The 64-token alignment holds on every HybridEP recipe under dynamic batching: input padding and sequence_length_round are both multiples of 64, Gemma CP2 pads to 128, and DeepSeek-V4 CP8 is covered by the model's own per-rank padding to 128 (deepseek_v4/model.py:895).
The guard test finds EP recipes from the resolved config rather than the filename, so inherited variants like the DPO 1n4g torch override are covered.
Multi-node HybridEP in the container (HYBRID_EP_MULTINODE=1) already landed in #4038, and this branch includes it.

Remaining comments are on leftover DeepEP usages (three VLM expert-parallel recipes, the enable_deepep schema key, the DeepSeek-V4 guide and nightly labels) and on where the new recipe tests should live.

yuki-97 · 2026-09-25T04:46:37Z

+RECIPE_DIR = Path(__file__).parents[3] / "examples/configs/recipes/llm"
+RECIPES = tuple(sorted(RECIPE_DIR.rglob("*.yaml")))
+
+
+@pytest.mark.parametrize(
+    "recipe_path", RECIPES, ids=lambda path: str(path.relative_to(RECIPE_DIR))
+)


Three AutoModel expert-parallel VLM recipes still use dispatcher: deepep, and this test can't catch them because it only globs recipes/llm:

vlm_grpo-nemotron-omni-30ba3b-clevr-1n8g-automodel-ep8.v2.yaml:32 (EP8)

vlm_grpo-nemotron-omni-30ba3b-mmpr-4n8g-automodel-ep8.v1.yaml:61 (EP8)

vlm_grpo-qwen3.5-35ba3b-geo3k-2n8g-automodel-ep16.yaml:45 (EP16, the VLM twin of the LLM ep16 recipe migrated here)

vlm_grpo-gemma4-e4b-geo3k-1n8g-automodel.yaml:36 also says deepep; the model is dense so it is a no-op there, but worth flipping too so no AutoModel recipe is left on DeepEP.

Suggest fix:

Switch those recipes to dispatcher: hybridep and add make_sequence_length_divisible_by: 64 under policy: (they inherit 1 from grpo_math_1B.yaml today).

Move test_automodel_moe_recipes_use_hybridep_or_explicit_torch into test_config_validation.py:103 and parametrize it over the existing config_files with load_config_with_inheritance, as test_all_config_no_tp_size_accuracy_issues does; config_files already includes recipes/vlm/.

After switching them, rerun those recipes' nightly scripts under tests/test_suites/vlm/ to confirm they still pass with HybridEP.

yuki-97 · 2026-09-25T04:46:37Z

+        pytest.skip("Not an AutoModel expert-parallel recipe")
+
+    backend = dtensor_cfg.automodel_kwargs.backend
+    assert "enable_deepep" not in backend


enable_deepep is still declared in the schema even though this test now forbids it in every recipe:

init.py:137-139 keeps enable_deepep: NotRequired[bool], with a comment pointing users to dispatcher="deepep".

test_automodel_types.py:56 still passes "enable_deepep": True.

Nothing in nemo_rl reads it, and the pinned Automodel ignores it with a warning (utils.py:303-314).

Suggest fix:

Delete init.py:137-139 and add "hybridep" to the dispatcher comment on init.py:135.

Drop "enable_deepep": True from test_automodel_types.py:56 (optionally replace it with "dispatcher": "hybridep").

yuki-97 · 2026-09-25T04:46:37Z

        rms_norm: torch_fp32
        experts: torch_mm
-        dispatcher: deepep
+        dispatcher: hybridep


The DeepSeek-V4 guide still says this recipe uses DeepEP: deepseek-v4-flash.md:26-27 lists "the DeepEP expert dispatcher" under Validated Scope. The two gpt-oss nightly entries are also still labelled # gpt-oss 20b DeepEP test (nightly.txt:255, nightly_gb200.txt:74), though their recipe moves to HybridEP in this PR.

Suggest fix:

Change the guide line to "the HybridEP expert dispatcher" once the DSV4 HybridEP run is in.

Rename both nightly comments to # gpt-oss 20b HybridEP test.

yuki-97 · 2026-09-25T05:33:28Z

+        )
+
+
+def test_gemma4_cp_keeps_local_hybridep_inputs_aligned() -> None:


test_gemma4_cp_keeps_local_hybridep_inputs_aligned (line 55) and test_automodel_recipes_preserve_backend_choices pin the settings of specific recipes, and tests of this kind are already spread across tests/unit/tools/ and tests/unit/models/generation/.

Suggest fix:

Move these two tests into a new tests/unit/test_recipe_settings.py.

Optional for this PR, since they predate it: move the existing recipe-setting tests there too.

test_hybridep_default_8g_recipes.py

test_vllm_qwen35_bf16_trtllm_recipe.py

test_video_grpo_recipes.py

test_swe1_dynamo_config.py

#3724, #3733 and #4077 add similar tests as well.

perf: use HybridEP in AutoModel MoE recipes

629b972

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

seonjinn requested review from a team as code owners September 7, 2026 04:31

seonjinn added the CI:L1 Run doctests, unit tests, and functional tests label Sep 7, 2026

seonjinn requested a review from sharonyu-115 September 7, 2026 04:33

seonjinn mentioned this pull request Sep 7, 2026

build: enable multi-node HybridEP in container images #4038

Merged

Merge branch 'main' into sna/hybridep-automodel-recipes-20260906

abfcffd

Merge branch 'main' into sna/hybridep-automodel-recipes-20260906

3051348

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

seonjinn added CI:Lfast Runs a fast test suite and re-use nightly `main` container (but sync dependencies to PRs version) and removed CI:L1 Run doctests, unit tests, and functional tests labels Sep 17, 2026

Merge branch 'main' into sna/hybridep-automodel-recipes-20260906

27ac693

seonjinn requested a review from terrykong September 18, 2026 06:59

seonjinn added 3 commits September 18, 2026 01:03

Merge branch 'main' into sna/hybridep-automodel-recipes-20260906

cf1446a

fix: include DeepSeek V4 in AutoModel HybridEP recipe migration

5aaa7af

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

Merge branch 'main' into sna/hybridep-automodel-recipes-20260906

8cd048d

Merge branch 'main' into sna/hybridep-automodel-recipes-20260906

2eb68fa

sharonyu-115 requested changes Sep 19, 2026

View reviewed changes

sharonyu-115 reviewed Sep 19, 2026

View reviewed changes

zpqiu reviewed Sep 20, 2026

View reviewed changes

Comment thread examples/configs/recipes/llm/grpo-glm47-flash-4n8g-automodel.yaml

Merge upstream main into HybridEP AutoModel recipes

d72a905

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

Merge remote-tracking branch 'origin/main' into fix/pr4037-review-202…

ffc8e30

…60921 Signed-off-by: Seonjin Na <sna@nvidia.com>

yuki-97 reviewed Sep 25, 2026

View reviewed changes

		)


		def test_gemma4_cp_keeps_local_hybridep_inputs_aligned() -> None:

Conversation

seonjinn commented Sep 7, 2026

Summary

Container requirement

Testing

Uh oh!

copy-pr-bot Bot commented Sep 7, 2026

Uh oh!

seonjinn commented Sep 8, 2026

Uh oh!

seonjinn commented Sep 12, 2026

Uh oh!

seonjinn commented Sep 18, 2026

Uh oh!

seonjinn commented Sep 18, 2026

Uh oh!

seonjinn commented Sep 19, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

sharonyu-115 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

seonjinn commented Sep 21, 2026

Uh oh!

seonjinn commented Sep 21, 2026

Uh oh!

seonjinn commented Sep 23, 2026

Uh oh!

seonjinn commented Sep 25, 2026

Uh oh!

yuki-97 left a comment

Choose a reason for hiding this comment

Uh oh!

yuki-97 Sep 25, 2026

Choose a reason for hiding this comment

Uh oh!

yuki-97 Sep 25, 2026

Choose a reason for hiding this comment

Uh oh!

yuki-97 Sep 25, 2026

Choose a reason for hiding this comment

Uh oh!

yuki-97 Sep 25, 2026

Choose a reason for hiding this comment

Uh oh!

This branch has not been deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants