feat(nemotron): Add GTP + partial CG + cuDNN LN support for NT3 Ultra MXFP8 GB200/GB300/VRNVL72 recipes by ghadiaravi13 · Pull Request #6228 · NVIDIA-NeMo/Megatron-Bridge

ghadiaravi13 · 2026-09-28T20:08:23Z

What does this PR do ?

Adds GTP + partial CG support to NT3 Ultra MXFP8 GB200/GB300/VRNVL72 recipes
Also enables cuDNN Layernorm for both GB200 and GB300 recipes.

The GB200/VRNVL72 configs are also updated with NCCLEP support

Changelog

Partial CUDA graph for NT3 Ultra MXFP8 GB200/GB300/VRNVL72 recipes
Added GTP sharding for the moe_latent_proj module in Nemotron 3 Ultra GB200 configuration.
Implemented local CUDA graph handling in training and validation processes.
Updated tests to ensure proper validation of local per-layer CUDA graph modules.
Updated Rubin tests to match with the new GB300 config
Improved error handling for unsupported local whole-layer captures.

GitHub Actions CI

See the CI section in the Contributing doc for how to trigger the CI. A Nvidia developer will need to approve and trigger the CI for external contributors.

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 add or update any necessary documentation?
Does the PR affect components that are optional to install? (Ex: Numba, Pynini, Apex etc)
- Reviewer: Does the PR have correct import guards for all optional libraries?

If you haven't finished some of the above items you can still open "Draft" PR.

Additional Information

Related to # (issue)

…nd GB300 - Introduced `_apply_nemotron_3_ultra_gtp` function to configure GTP settings for dense weight sharding. - Updated `nemotron_3_ultra_pretrain_256gpu_gb200_fp8mx_config` and `nemotron_3_ultra_pretrain_256gpu_gb300_fp8mx_config` to apply GTP settings. - Added support for GTP sharding in the `moe_latent_proj` module to address memory pressure. Signed-off-by: Ravi Ghadia <rghadia@nvidia.com>

…into rghadia/nt3-ultra-gtp

Signed-off-by: Ravi Ghadia <rghadia@nvidia.com>

- Partial CUDA graph for NT3 Ultra MXFP8 GB200/GB300 recipes - Added GTP sharding for the `moe_latent_proj` module in Nemotron 3 Ultra GB200 configuration. - Implemented local CUDA graph handling in training and validation processes. - Updated tests to ensure proper validation of local per-layer CUDA graph modules. - Improved error handling for unsupported local whole-layer captures. Signed-off-by: Ravi Ghadia <rghadia@nvidia.com>

copy-pr-bot · 2026-09-28T20:08:27Z

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.

…into rghadia/nt3-ultra-gtp-cg

…into rghadia/nt3-ultra-gtp-cg Signed-off-by: rghadia <rghadia@nsc-lhr-slurm-1-vscode-01.nvidia.com>

- Changed references from HybridEP to NCCLEP in both GB200 and GB300 configurations. - Added support for cuDNN normalization kernels in both configurations. - Enabled NCCLEP and adjusted expert token handling in GB200. - Commented out the HSDP/FSDP dtype overrides in GB300 for clarity. These changes enhance the performance and compatibility of the Nemotron 3 Ultra pretraining setups. Signed-off-by: rghadia <rghadia@nvidia.com>

malay-nagda · 2026-10-10T06:48:55Z

/ok to test 24cb4f9

github-actions · 2026-10-10T06:49:11Z

Automatic Claude reviews have been retired. To request a pull-request review, post a comment containing:

/review

Add model=claude to use a Claude reviewer (the default is model=codex). mode=light|strict selects the review depth; for example, /review model=claude mode=strict. Comment /review help for all options.

Signed-off-by: rghadia <rghadia@nvidia.com>

…00 recipe Signed-off-by: rghadia <rghadia@nsc-lhr-slurm-1-vscode-01.nvidia.com>

ghadiaravi13 · 2026-10-11T04:34:01Z

/ok to test 39c25c6

copy-pr-bot · 2026-10-11T04:34:04Z

/ok to test 39c25c6

@ghadiaravi13, there was an error processing your request: E2

See the following link for more information: https://docs.gha-runners.nvidia.com/cpr/e/2/

ghadiaravi13 · 2026-10-11T04:36:15Z

/ok to test 39c25c6

…into rghadia/nt3-ultra-gtp-cg

- Replaced HybridEP references with NCCLEP in GB300 and Rubin configurations. - Moved NCCLEP to the base recipe constructor so that Rubin recipe can pick up the NCCLEP changes - Updated test cases to reflect changes in execution layout and environment variables. - Improved clarity in comments regarding NCCL EP settings. These modifications enhance the performance and compatibility of the Nemotron 3 Ultra pretraining setups. Signed-off-by: rghadia <rghadia@nvidia.com>

ghadiaravi13 · 2026-10-12T18:32:29Z

/ok to test 892a800

- Removed "moe_router" from the set_cuda_graph_modules function calls in both GB200 and GB300 configurations. - Added a TODO comment to indicate the need for this change due to MCore PR 7046, which affects CPU-to-GPU copy for per token loss. Signed-off-by: rghadia <rghadia@nvidia.com>

Signed-off-by: Ravi Ghadia <rghadia@nvidia.com>

malay-nagda · 2026-10-13T04:39:09Z

/ok to test b892a0e

malay-nagda · 2026-10-13T04:43:59Z

@cuichenx can you take a look too since it touches src code?

Signed-off-by: rghadia <rghadia@nvidia.com>

…avi13/NT3-Ultra-Megatron-Bridge into rghadia/nt3-ultra-gtp-cg Signed-off-by: rghadia <rghadia@nvidia.com>

ghadiaravi13 · 2026-10-13T05:22:54Z

/ok to test cbd1d28

…into rghadia/nt3-ultra-gtp-cg Signed-off-by: rghadia <rghadia@nvidia.com>

Signed-off-by: rghadia <rghadia@nvidia.com>

ghadiaravi13 · 2026-10-13T05:37:34Z

/ok to test 065e181

#6437) Signed-off-by: Ravi Ghadia <rghadia@nvidia.com> Signed-off-by: Malay Nagda <malayn@nvidia.com> Co-authored-by: Ravi Ghadia <40660742+ghadiaravi13@users.noreply.github.com>

rghadia added 5 commits September 26, 2026 11:10

Merge branch 'main' of https://github.com/NVIDIA-NeMo/Megatron-Bridge …

ca9baf2

…into rghadia/nt3-ultra-gtp

fix import statement

bdf372b

Signed-off-by: Ravi Ghadia <rghadia@nvidia.com>

Remove FLASHINFER_DISABLE_VERSION_CHECK=1 workarounds

d7e0fb8

Signed-off-by: Ravi Ghadia <rghadia@nvidia.com>

Merge branch 'main' of https://github.com/NVIDIA-NeMo/Megatron-Bridge …

bb02c3d

…into rghadia/nt3-ultra-gtp-cg

cuichenx added feature New capabilities, enhancements, or enablement work area:recipe Training recipes and launch configs labels Sep 28, 2026

Merge branch 'main' of https://github.com/NVIDIA-NeMo/Megatron-Bridge …

6e896da

…into rghadia/nt3-ultra-gtp-cg Signed-off-by: rghadia <rghadia@nsc-lhr-slurm-1-vscode-01.nvidia.com>

ghadiaravi13 changed the title ~~feat(nemotron): Add partial CG support for NT3 Ultra MXFP8 GB200/GB300 recipes~~ feat(nemotron): Add GTP + partial CG support for NT3 Ultra MXFP8 GB200/GB300 recipes Oct 9, 2026

ghadiaravi13 marked this pull request as ready for review October 9, 2026 23:04

ghadiaravi13 requested review from dingqingy-nv and malay-nagda as code owners October 9, 2026 23:04

cuichenx added the needs-review PR is ready for code review and waiting on a reviewer label Oct 9, 2026

ghadiaravi13 changed the title ~~feat(nemotron): Add GTP + partial CG support for NT3 Ultra MXFP8 GB200/GB300 recipes~~ feat(nemotron): Add GTP + partial CG + cuDNN LN support for NT3 Ultra MXFP8 GB200/GB300 recipes Oct 9, 2026

malay-nagda previously approved these changes Oct 10, 2026

View reviewed changes

Minor fix: Added _enable_ncclep import to GB200 perf-recipe

4419bec

Signed-off-by: rghadia <rghadia@nvidia.com>

ghadiaravi13 dismissed malay-nagda’s stale review via 4419bec October 11, 2026 00:05

fix(nemotron): remove unused hybrid_ep_ranks_per_nvlink_domain in GB2…

39c25c6

…00 recipe Signed-off-by: rghadia <rghadia@nsc-lhr-slurm-1-vscode-01.nvidia.com>

copy-pr-bot Bot deployed to test October 11, 2026 04:37 Active

ghadiaravi13 added 2 commits October 12, 2026 11:20

Merge branch 'main' of https://github.com/NVIDIA-NeMo/Megatron-Bridge …

212bdd7

…into rghadia/nt3-ultra-gtp-cg

ghadiaravi13 changed the title ~~feat(nemotron): Add GTP + partial CG + cuDNN LN support for NT3 Ultra MXFP8 GB200/GB300 recipes~~ feat(nemotron): Add GTP + partial CG + cuDNN LN support for NT3 Ultra MXFP8 GB200/GB300/VRNVL72 recipes Oct 12, 2026

copy-pr-bot Bot deployed to test October 12, 2026 18:33 Active

ghadiaravi13 added 2 commits October 12, 2026 16:51

Fix: Remove moe_preprocess from partial CG scope

b892a0e

Signed-off-by: Ravi Ghadia <rghadia@nvidia.com>

malay-nagda previously approved these changes Oct 13, 2026

View reviewed changes

malay-nagda requested a review from cuichenx October 13, 2026 04:43

ghadiaravi13 added 2 commits October 12, 2026 22:16

fix(nemotron): lint fixes

d94dc39

Signed-off-by: rghadia <rghadia@nvidia.com>

Merge branch 'rghadia/nt3-ultra-gtp-cg' of https://github.com/ghadiar…

cbd1d28

…avi13/NT3-Ultra-Megatron-Bridge into rghadia/nt3-ultra-gtp-cg Signed-off-by: rghadia <rghadia@nvidia.com>

ghadiaravi13 dismissed malay-nagda’s stale review via cbd1d28 October 13, 2026 05:21

copy-pr-bot Bot deployed to test October 13, 2026 05:24 Active

ghadiaravi13 added 2 commits October 12, 2026 22:34

Merge branch 'main' of https://github.com/NVIDIA-NeMo/Megatron-Bridge …

2d5410d

…into rghadia/nt3-ultra-gtp-cg Signed-off-by: rghadia <rghadia@nvidia.com>

fix: update comment to vrnvl72

065e181

Signed-off-by: rghadia <rghadia@nvidia.com>

copy-pr-bot Bot deployed to test October 13, 2026 05:38 Active

malay-nagda approved these changes Oct 13, 2026

View reviewed changes

malay-nagda added the r0.7.0 Auto-cherrypick to release branch. Apply before merge; cherrypick happens after merge. label Oct 13, 2026

malay-nagda merged commit 2eca468 into NVIDIA-NeMo:main Oct 13, 2026
85 checks passed

malay-nagda mentioned this pull request Oct 13, 2026

feat(recipe): backport NT3 Ultra performance updates to r0.7.0 (#6228) #6437

Merged

malay-nagda mentioned this pull request Oct 16, 2026

perf(recipe): default GB200 and VRNVL72 MoE perf recipes to NCCL EP and grouped tensors #6373

Draft

2 of 11 tasks

Navigation Menu

Repository navigation

feat(nemotron): Add GTP + partial CG + cuDNN LN support for NT3 Ultra MXFP8 GB200/GB300/VRNVL72 recipes - #6228

feat(nemotron): Add GTP + partial CG + cuDNN LN support for NT3 Ultra MXFP8 GB200/GB300/VRNVL72 recipes#6228
malay-nagda merged 18 commits into
NVIDIA-NeMo:mainfrom
ghadiaravi13:rghadia/nt3-ultra-gtp-cg

ghadiaravi13 commented Sep 28, 2026 •

edited

Loading

Uh oh!

copy-pr-bot Bot commented Sep 28, 2026

Uh oh!

malay-nagda commented Oct 10, 2026

Uh oh!

github-actions Bot commented Oct 10, 2026

Uh oh!

ghadiaravi13 commented Oct 11, 2026

Uh oh!

copy-pr-bot Bot commented Oct 11, 2026

Uh oh!

ghadiaravi13 commented Oct 11, 2026

Uh oh!

ghadiaravi13 commented Oct 12, 2026

Uh oh!

malay-nagda commented Oct 13, 2026

Uh oh!

malay-nagda commented Oct 13, 2026

Uh oh!

ghadiaravi13 commented Oct 13, 2026

Uh oh!

ghadiaravi13 commented Oct 13, 2026

Uh oh!

Uh oh!

This branch was successfully deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

ghadiaravi13 commented Sep 28, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do ?

Changelog

GitHub Actions CI

Before your PR is "Ready for review"

Additional Information

Uh oh!

copy-pr-bot Bot commented Sep 28, 2026

Uh oh!

malay-nagda commented Oct 10, 2026

Uh oh!

github-actions Bot commented Oct 10, 2026

Uh oh!

ghadiaravi13 commented Oct 11, 2026

Uh oh!

copy-pr-bot Bot commented Oct 11, 2026

Uh oh!

ghadiaravi13 commented Oct 11, 2026

Uh oh!

ghadiaravi13 commented Oct 12, 2026

Uh oh!

malay-nagda commented Oct 13, 2026

Uh oh!

malay-nagda commented Oct 13, 2026

Uh oh!

ghadiaravi13 commented Oct 13, 2026

Uh oh!

ghadiaravi13 commented Oct 13, 2026

Uh oh!

Uh oh!

This branch was successfully deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

ghadiaravi13 commented Sep 28, 2026 •

edited

Loading