feat (recipe): add Nemotron 3 Ultra pretraining perf-recipe

[ghadiaravi13]

What does this PR do ?

Add Nemotron 3 Ultra pretraining configuration for 256 GPUs with MXFP8 support

• Introduced new configuration for Nemotron 3 Ultra pretraining on GB300 with 256 GPUs.
• Added model and performance settings tailored for MXFP8 mixed precision.
• Implemented necessary adjustments for Megatron-FSDP and HSDP compatibility.

This enhances the training capabilities for large-scale models using the latest precision techniques.

Changelog

• Add nemotron_3_ultra_pretrain_256gpu_gb300_fp8mx_v1_config to nemotronh/init.py
• Add helper functions to override config parameters (NT3 ultra specific: token_dispatcher = flex, dispatcher_backend = hybridep, TE op fuser (w/ CuteDSL fused grouped GEMM), selective activation CPU offloading for MoE sublayer, FSDP (and HSDP across nodes)

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)

[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.

[ghadiaravi13]

/ok to test 8bcdd80

[dingqingy-nv]

/claude review

[claude]

Review — feat(recipe): add Nemotron 3 Ultra 256-GPU GB300 MXFP8 perf-recipe

Scope: three files under perf_recipes/nemotronh/ — new shared builder + Ultra config helpers in common.py, the nemotron_3_ultra_pretrain_256gpu_gb300_fp8mx_v1_config entrypoint in gb300/nemotronh.py, and its init.py export. I verified every cfg.model.* / cfg.ddp.* / cfg.dist.* / cfg.scheduler.* attribute set in the diff resolves to a real config field (no phantom attributes), the base nemotron_3_ultra_pretrain_config alias exists, and the HSDP num_optim_instances = 256 // 64 = 4 -> outer_dp_sharding_strategy=optim logic is correct. Overall the change looks sound. A few small items:

Low — inaccurate comment on cross-entropy impl (common.py:203-204)
The comment reads: --cross-entropy-fusion-impl native (the common perf override forces te). But for the Ultra path the common override is _apply_nemotron_3_ultra_perf_defaults, which calls _benchmark_common(cfg, cross_entropy_impl=native) — it already sets native, not te. The re-assignment here is redundant and the parenthetical is misleading (appears copied from the generic perf path). Consider fixing the comment or dropping the redundant line.

Low — missing blank line between top-level functions (gb300/nemotronh.py:164-165)
Only one blank line separates the end of nemotron_3_ultra_pretrain_256gpu_gb300_fp8mx_v1_config from def nemotron_3_nano_pretrain_8gpu_gb300_bf16_config; every other top-level function in the file uses two. Run uv run pre-commit run --all-files (ruff format) before merge to fix.

Low — no test exercises the new config
The new config is exported but nothing instantiates it; test_perf_config_integration.py covers llama/deepseek/qwen but not Nemotron 3 Ultra. Config construction is pure (no GPUs needed), so a cheap instantiation test asserting key fields (EP=64, GBS=256, use_megatron_fsdp, outer_dp_sharding_strategy=optim, ckpt_format=fsdp_dtensor) would guard against silent regressions and phantom-attribute breakage on future config refactors.

Suggested test cases

• nemotron_3_ultra_pretrain_256gpu_gb300_fp8mx_v1_config — add a construction test (via get_perf_recipe_by_name with model_recipe_name=nemotron_3_ultra, task=pretrain, num_gpus=256, gpu=gb300, precision=fp8_mx, config_variant=v1) asserting the FSDP/HSDP + MoE fields resolve as intended.

(No functional/GPU perf test is impacted — the only touched recipe targets 256 GPUs, above the 2-GPU CI limit; coverage is limited to config-instantiation.)

[ghadiaravi13]

/claude review

[claude]

Review — nemotron_3_ultra 256xGB300 MXFP8 perf recipe

Light review. Well-structured; follows the Nemotron 3 Super GB300 pattern. HSDP/FSDP handling looks sound (guarding outer_dp_sharding_strategy=no_shard for a single NVLink domain, dtype overrides applied last). A few points:

1. Misplaced string literal (gb300/nemotronh.py:142-148): the triple-quoted block sits AFTER executable statements, so it is a no-op expression, not a docstring, and is never attached to doc. Move it to be the first statement in the function, or make it a comment.

2. Verify min_offloaded_tensor_size (gb300/nemotronh.py:161): this attribute appears nowhere else in the repo. If it is not a real field on the model provider / mcore TransformerConfig, the assignment silently creates a phantom attribute and the intended CPU-OOM guard becomes a no-op. Please confirm the name.

3. Test coverage: no test instantiates nemotron_3_ultra_pretrain_256gpu_gb300_fp8mx_config. A minimal instantiation test (asserting the config builds and EP=64, GBS=256, ckpt_format=fsdp_dtensor, outer_dp_sharding_strategy=optim at 256 GPUs) would guard against phantom-attribute regressions. The submodule is not initialized in this environment, so I could not build the config to confirm every cfg.model/cfg.ddp/cfg.dist attribute exists.

Minor: trailing whitespace on a few added lines; run 'uv run ruff format' before merge.

Suggested test cases:

• nemotron_3_ultra_pretrain_256gpu_gb300_fp8mx_perf
• nemotron_3_super_pretrain_64gpu_gb300_bf16_perf
• nemotron_3_super_pretrain_64gpu_gb300_fp8mx_perf
• nemotron_3_super_pretrain_64gpu_gb300_nvfp4_perf

[dingqingy-nv]

A generic feedback is that agent tends to over-comment. Some comment provide no value as the code itself is self-explanatory.

Also agents love to comment about irrelevant exploratory processes. e.g. we start this recipe from existing MLM jobs, and agent documents the argument translation to the comment, which provides no helpful information for a custom facing recipe. It shall comment on things that is helpful for custom to understand.

This seems to be a common problem and we shall fix with a skill/guideline.

[ghadiaravi13]

/ok to test d42338a

[ghadiaravi13]

/ok to test 186c7fc

[malay-nagda]

/ok to test 296eda1