2121 _apply_nemotron_3_nano_perf_defaults ,
2222 _apply_nemotron_3_super_perf_defaults ,
2323 _apply_nemotron_3_ultra_fsdp_hsdp ,
24+ _apply_nemotron_3_ultra_gtp ,
2425 _apply_nemotron_3_ultra_perf_defaults ,
2526 _benchmark_common ,
2627 _enable_nemotron_3_5_lightning_full_iteration ,
@@ -209,13 +210,15 @@ def nemotron_3_super_pretrain_64gpu_gb300_nvfp4_config() -> ConfigContainer:
209210def _nemotron_3_ultra_gb300_fp8mx_config (
210211 * , num_gpus : int , expert_model_parallel_size : int , global_batch_size : int
211212) -> ConfigContainer :
212- """Shared builder for Nemotron 3 Ultra GB300 MXFP8 Megatron-FSDP perf recipes."""
213+ """Shared builder for Nemotron 3 Ultra GB300 MXFP8 GTP perf recipes."""
213214 cfg = nemotron_3_ultra_pretrain_config ()
214215 cfg .mixed_precision = _perf_precision ("fp8_mx" )
215216
217+ _apply_nemotron_3_ultra_perf_defaults (cfg )
218+
216219 """
217220 Uses TP1 / PP1 / CP1 / EP64 / ETP1, GBS 256 / MBS 1,
218- seq 8192, HybridEP flex dispatcher, CuteDSL fused grouped MLP, selective
221+ seq 8192, NCCLEP flex dispatcher, CuteDSL fused grouped MLP, selective
219222 recompute + activation offload of the expert MLP, MTP=2. The MoE architecture
220223 (512 experts, latent MoE, MTP, squared-relu, hybrid Mamba/attention pattern,
221224 ...) is inherited from the base recipe via ``AutoBridge``.
@@ -252,20 +255,30 @@ def _nemotron_3_ultra_gb300_fp8mx_config(
252255 cfg .model .recompute_granularity = "selective"
253256 cfg .model .recompute_modules = ["moe_act" ]
254257
255- _apply_nemotron_3_ultra_perf_defaults (cfg )
258+ # Shard dense weights with GTP within each 64-GPU rack. GTP currently
259+ # requires a single distributed-optimizer instance.
260+ _apply_nemotron_3_ultra_gtp (cfg )
256261
257- # Apply HSDP / FSDP dtype overrides last so they win over the generic defaults.
258- _apply_nemotron_3_ultra_fsdp_hsdp (cfg , num_gpus = num_gpus )
262+ cfg .model .cuda_graph_impl = "local"
263+ set_cuda_graph_modules (cfg .model , ["attn" , "mamba" ])
264+ # TODO: "moe_preprocess", "moe_router" needs to be disabled due to MCore PR 7046, performing a CPU-to-GPU copy for per token loss
265+ # TE attention requires TE's RNG tracker while MCore local graphs are capturing.
266+ cfg .rng .te_rng_tracker = cfg .model .use_te_rng_tracker = True
267+
268+ _enable_ncclep (cfg )
269+ # Device-side expert token counts: the legacy grouped MLP path syncs tokens_per_expert to the
270+ # host every layer, which serializes the CPU behind the GPU when dispatch is fast.
271+ cfg .model .moe_use_grouped_tensor = True
259272
260273 return cfg
261274
262275
263276def nemotron_3_ultra_pretrain_256gpu_gb300_fp8mx_config (* , num_gpus : int = 256 ) -> ConfigContainer :
264- """Nemotron 3 Ultra (550B-A55B LatentMoE) pretrain: 256× GB300, MXFP8, Megatron-FSDP (HSDP) .
277+ """Nemotron 3 Ultra (550B-A55B LatentMoE) pretrain: 256× GB300, MXFP8, GTP .
265278
266279 TP1 / PP1 / CP1 / EP64 / ETP1, GBS 256 / MBS 1, seq 8192, BF16 + MXFP8 mixed
267- precision, HybridEP flex dispatcher, CuteDSL fused grouped MLP, selective
268- recompute + fine-grained activation offload of the expert MLP, MTP=2.
280+ precision, dense GTP64, NCCLEP flex dispatcher, CuteDSL fused grouped MLP,
281+ selective recompute + fine-grained activation offload of the expert MLP, MTP=2.
269282
270283 Args:
271284 num_gpus: Launch GPU count used to size HSDP optimizer instances.
@@ -276,10 +289,6 @@ def nemotron_3_ultra_pretrain_256gpu_gb300_fp8mx_config(*, num_gpus: int = 256)
276289 expert_model_parallel_size = 64 ,
277290 global_batch_size = 256 ,
278291 )
279- _enable_ncclep (cfg )
280- # Device-side expert token counts: the legacy grouped MLP path syncs tokens_per_expert to the
281- # host every layer, which serializes the CPU behind the GPU when dispatch is fast.
282- cfg .model .moe_use_grouped_tensor = True
283292 # Keep process settings next to the recipe so users can see the exact benchmark environment.
284293 cfg .env_vars = {
285294 ** COMMON_PERF_ENV_VARS ,
@@ -296,6 +305,9 @@ def nemotron_3_ultra_pretrain_256gpu_gb300_fp8mx_config(*, num_gpus: int = 256)
296305 # Transformer Engine overlap settings for this model.
297306 "NVTE_BWD_LAYERNORM_SM_MARGIN" : 20 ,
298307 "NVTE_FWD_LAYERNORM_SM_MARGIN" : 20 ,
308+ # Use cuDNN normalization kernels for the VR200 performance candidate.
309+ "NVTE_NORM_BWD_USE_CUDNN" : 1 ,
310+ "NVTE_NORM_FWD_USE_CUDNN" : 1 ,
299311 # Required by fine_grained_activation_offloading (TE >= 2.10.0) to avoid
300312 # offloading weights;
301313 "NVTE_CPU_OFFLOAD_V1" : 1 ,
0 commit comments