perf: enable MoE GroupedGEMM for MoE models by seonjinn · Pull Request #2278 · NVIDIA-NeMo/RL

seonjinn · 2026-04-17T04:19:34Z

Support moe_grouped_gemm (grouped GEMM for MoE experts) through the MegatronConfig TypedDict and _apply_moe_config(). Enables it in every root MoE performance recipes (Qwen3-30B-A3B 4n4g/4n8g/4n8g-40K, Qwen3-235B 16n8g, DeepSeek-V3 32n8g, DAPO DeepSeek-V3 64n8g); child recipes inherit.

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

...

Wires moe_grouped_gemm (CUTLASS grouped GEMM for MoE experts) through the MegatronConfig TypedDict and _apply_moe_config(). Enables it in every root MoE performance recipe (Qwen3-30B-A3B 4n4g/4n8g/4n8g-40K, Qwen3-235B 16n8g, DeepSeek-V3 32n8g, DAPO DeepSeek-V3 64n8g); child recipes inherit. Signed-off-by: sna <sna@nvidia.com>

copy-pr-bot · 2026-04-17T04:19:39Z

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.

seonjinn · 2026-04-17T06:14:32Z

/ok to test 1713776

terrykong

Summary

This PR adds moe_grouped_gemm support through the MegatronConfig TypedDict and _apply_moe_config(), enabling it in 6 root MoE performance recipes. The concept is sound and the feature is valuable, but the implementation has a critical issue with the config application pattern that could silently regress throughput for ~14 other MoE recipes.

Ship blocker

Silent performance regression: The .get("moe_grouped_gemm", False) pattern unconditionally overwrites Bridge's default of True with False for any config that omits the field. See inline comment for details and suggested fix.

Suggestions

Performance evidence: Could you share before/after throughput data (tokens/sec or step time) on at least one representative recipe? A short convergence sanity check would also be helpful.
PR description: The "What does this PR do?", "Issues", and "Usage" sections still have template placeholders — consider filling them in.

Generated by Claude Code

Co-authored-by: Terry Kong <terryk@nvidia.com> Signed-off-by: Seonjin <sna@nvidia.com>

seonjinn · 2026-04-20T08:00:05Z

/ok to test 2520299

seonjinn · 2026-04-23T21:41:35Z

/ok test 1cd8dd1

seonjinn · 2026-04-23T22:47:24Z

@terrykong I've made changes based on your suggestions. Does this PR look good to you?

seonjinn · 2026-05-14T21:03:42Z

/ok to test ab2dacb

…gemm _apply_moe_config requires expert_tensor_parallel_size and other keys via direct dict indexing. Minimal config triggered KeyError before reaching the moe_grouped_gemm branch. Base the new tests on the same megatron_cfg shape used by test_moe_configuration. Signed-off-by: sna <sna@nvidia.com>

seonjinn · 2026-05-16T23:29:31Z

/ok to test e71839a

seonjinn · 2026-05-19T01:09:25Z

@terrykong Can you also review this PR when you have a chance?

terrykong · 2026-05-19T01:25:21Z

lgtm. @seonjinn could you resolve the merge conflicts

Signed-off-by: sna <sna@nvidia.com> # Conflicts: # nemo_rl/models/policy/__init__.py # tests/unit/models/megatron/test_megatron_setup.py

seonjinn · 2026-05-19T02:50:14Z

/ok to test 5b19ee3

seonjinn · 2026-05-19T03:13:10Z

@terrykong Sure, I resolved the conflict issues

LinearizedRow.routed_experts arrives with the Gym-side R3 change (Gym PR #2278 follow-up); until the submodule pin carries it, read via getattr (capture-R3 degrades to the sentinel/group-drop path) and skip the two routed-publish tests. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Pranav Thombre <pthombre@oci-hsg-cs-001-vscode-02.cm.cluster>

seonjinn requested review from a team as code owners April 17, 2026 04:19

seonjinn requested review from a team as code owners April 17, 2026 05:10

seonjinn force-pushed the sj/grouped-gemm branch from f7bdb19 to 1713776 Compare April 17, 2026 06:12

seonjinn self-assigned this Apr 17, 2026

terrykong reviewed Apr 19, 2026

View reviewed changes

Comment thread nemo_rl/models/megatron/setup.py Outdated

Comment thread nemo_rl/models/policy/__init__.py

seonjinn and others added 3 commits April 20, 2026 00:15

Update nemo_rl/models/policy/__init__.py

bd94eaf

Co-authored-by: Terry Kong <terryk@nvidia.com> Signed-off-by: Seonjin <sna@nvidia.com>

Update nemo_rl/models/megatron/setup.py

6cd7278

Co-authored-by: Terry Kong <terryk@nvidia.com> Signed-off-by: Seonjin <sna@nvidia.com>

Merge branch 'main' into sj/grouped-gemm

95c7169

seonjinn added the CI:Lfast Runs a fast test suite and re-use nightly `main` container (but sync dependencies to PRs version) label Apr 23, 2026

Merge branch 'main' into sj/grouped-gemm

6fa8f74

seonjinn added 2 commits May 5, 2026 23:32

Merge branch 'main' into sj/grouped-gemm

7f9a0d9

Merge remote-tracking branch 'origin/main' into sj/grouped-gemm

3e16b02

seonjinn requested a review from a team May 14, 2026 21:03

copy-pr-bot Bot temporarily deployed to public May 14, 2026 21:04 Inactive

copy-pr-bot Bot temporarily deployed to nemo-ci May 14, 2026 21:04 Inactive

copy-pr-bot Bot temporarily deployed to public May 14, 2026 21:04 Inactive

seonjinn requested review from a team May 14, 2026 21:04

copy-pr-bot Bot temporarily deployed to public May 16, 2026 22:28 Inactive

copy-pr-bot Bot temporarily deployed to public May 16, 2026 23:29 Inactive

copy-pr-bot Bot temporarily deployed to nemo-ci May 16, 2026 23:29 Inactive

copy-pr-bot Bot temporarily deployed to public May 16, 2026 23:30 Inactive

copy-pr-bot Bot temporarily deployed to public May 16, 2026 23:34 Inactive

guyueh1 previously approved these changes May 18, 2026

View reviewed changes

terrykong previously approved these changes May 19, 2026

View reviewed changes

Merge main into sj/grouped-gemm

5b19ee3

Signed-off-by: sna <sna@nvidia.com> # Conflicts: # nemo_rl/models/policy/__init__.py # tests/unit/models/megatron/test_megatron_setup.py

seonjinn dismissed stale reviews from terrykong and guyueh1 via 5b19ee3 May 19, 2026 02:49

copy-pr-bot Bot temporarily deployed to public May 19, 2026 02:50 Inactive

copy-pr-bot Bot temporarily deployed to nemo-ci May 19, 2026 02:50 Inactive

copy-pr-bot Bot temporarily deployed to public May 19, 2026 02:50 Inactive

copy-pr-bot Bot temporarily deployed to public May 19, 2026 02:54 Inactive

terrykong approved these changes May 19, 2026

View reviewed changes

terrykong merged commit 9053154 into main May 19, 2026
37 checks passed

terrykong deleted the sj/grouped-gemm branch May 19, 2026 04:58

terrykong mentioned this pull request Sep 4, 2026

feat(sc): gate-authoritative token capture via an external TransferQueue sink #3837

Open

Uh oh!

Conversation

seonjinn commented Apr 17, 2026

What does this PR do ?

Issues

Usage

Before your PR is "Ready for review"

Additional Information

Uh oh!

copy-pr-bot Bot commented Apr 17, 2026

Uh oh!

seonjinn commented Apr 17, 2026

Uh oh!

terrykong left a comment

Choose a reason for hiding this comment

Summary

Ship blocker

Suggestions

Uh oh!

Uh oh!

Uh oh!

seonjinn commented Apr 20, 2026

Uh oh!

seonjinn commented Apr 23, 2026

Uh oh!

seonjinn commented Apr 23, 2026

Uh oh!

seonjinn commented May 14, 2026

Uh oh!

seonjinn commented May 16, 2026

Uh oh!

seonjinn commented May 19, 2026

Uh oh!

terrykong commented May 19, 2026

Uh oh!

seonjinn commented May 19, 2026

Uh oh!

seonjinn commented May 19, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants