perf(recipe): enable Nemotron 3.5 MXFP8 full graphs on GB300 and VR200 (#6289)

Signed-off-by: Malay Nagda <malayn@nvidia.com>

Author: malay-nagda

src/megatron/bridge/perf_recipes/nemotronh/common.py

@@ -32,6 +32,7 @@
     nemotron_3_super_bf16_with_nvfp4_mixed,
     nemotron_3_ultra_bf16_with_nvfp4_mixed,
 )
+from megatron.bridge.utils.cuda_graph import set_full_iteration_cuda_graph
 
 
 _TE_QUANT_CFG_PATH = Path(__file__).with_name("te_quant.cfg")
@@ -142,6 +143,20 @@ def _enable_nemotron_3_super_full_iteration(cfg: ConfigContainer) -> None:
     cfg.model.delay_wgrad_compute = False
 
 
+def _enable_nemotron_3_5_lightning_full_iteration(cfg: ConfigContainer) -> None:
+    """Enable full-iteration capture with static MoE buffers for Lightning."""
+    set_full_iteration_cuda_graph(cfg.model)
+    cfg.rng.te_rng_tracker = True
+    cfg.model.use_te_rng_tracker = True
+
+    # Static receive capacity avoids dispatcher CPU synchronization during capture.
+    # The training loop retries capacity overflows without padding.
+    cfg.model.moe_expert_rank_capacity_factor = 1.5
+    cfg.model.moe_use_grouped_tensor = True
+    cfg.model.moe_paged_stash = False
+    cfg.model.offload_modules = []
+
+
 def _apply_nemotron_3_ultra_perf_defaults(cfg: ConfigContainer) -> None:
     """Apply shared Nemotron 3 Ultra perf defaults after recipe-specific overrides."""
 

src/megatron/bridge/perf_recipes/nemotronh/gb300/nemotronh.py

@@ -23,6 +23,7 @@
     _apply_nemotron_3_ultra_fsdp_hsdp,
     _apply_nemotron_3_ultra_perf_defaults,
     _benchmark_common,
+    _enable_nemotron_3_5_lightning_full_iteration,
     _enable_nemotron_3_super_full_iteration,
     _nemotron_3_super_nvfp4_precision,
     _nemotron_3_ultra_nvfp4_precision,
@@ -694,21 +695,19 @@ def _build_nemotron_3_5_lightning_gb300_mxfp8() -> ConfigContainer:
 
 
 def nemotron_3_5_lightning_pretrain_8gpu_gb300_fp8mx_config() -> ConfigContainer:
-    """Nemotron 3.5 Lightning pretrain: 8× GB300, MXFP8."""
+    """Nemotron 3.5 Lightning pretrain: 8× GB300, MXFP8 with full-iteration CUDA graph."""
     cfg = _build_nemotron_3_5_lightning_gb300_mxfp8()
     cfg.model.recompute_modules = []
     cfg.model.use_transformer_engine_op_fuser = True
     cfg.mixed_precision.fp8_dot_product_attention = True
     _enable_ncclep(cfg)
-    # Device-side expert token counts: the legacy grouped MLP path syncs tokens_per_expert to the
-    # host every layer, which serializes the CPU behind the GPU when dispatch is fast.
-    cfg.model.moe_use_grouped_tensor = True
+    _enable_nemotron_3_5_lightning_full_iteration(cfg)
     cfg.env_vars = {
         **COMMON_PERF_ENV_VARS,
         "CUDA_DEVICE_MAX_CONNECTIONS": 32,
         "NCCL_GRAPH_REGISTER": 0,
-        "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True",
-        "TORCH_NCCL_AVOID_RECORD_STREAMS": 1,
+        "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True,graph_capture_record_stream_reuse:True",
+        "TORCH_NCCL_AVOID_RECORD_STREAMS": 0,
         "NCCL_NVLS_ENABLE": 0,
         # NCCL EP dispatcher mode and one GPU per rank.
         "NCCL_EP_HT_EM_PULL_PUSH": 1,

src/megatron/bridge/perf_recipes/nemotronh/vr200/nemotronh.py

@@ -16,6 +16,7 @@
 from megatron.bridge.perf_recipes.environment import COMMON_PERF_ENV_VARS
 from megatron.bridge.perf_recipes.nemotronh.common import (
     ConfigContainer,
+    _enable_nemotron_3_5_lightning_full_iteration,
 )
 from megatron.bridge.perf_recipes.nemotronh.gb300.nemotronh import (
     _build_nemotron_3_5_lightning_gb300_bf16,
@@ -320,19 +321,20 @@ def nemotron_3_5_lightning_pretrain_8gpu_vr200_bf16_config() -> ConfigContainer:
 
 
 def nemotron_3_5_lightning_pretrain_8gpu_vr200_fp8mx_config() -> ConfigContainer:
-    """Nemotron 3.5 Lightning pretrain: 8× VR200, MXFP8 (alias of GB300)."""
+    """Nemotron 3.5 Lightning pretrain: 8× VR200, MXFP8 with full-iteration CUDA graph."""
     cfg = _build_nemotron_3_5_lightning_gb300_mxfp8()
     cfg.model.recompute_modules = []
     cfg.model.use_transformer_engine_op_fuser = True
     cfg.mixed_precision.fp8_dot_product_attention = True
+    _enable_nemotron_3_5_lightning_full_iteration(cfg)
 
     # Keep the VR200 launch environment explicit instead of inheriting it from GB300.
     cfg.env_vars = {
         **COMMON_PERF_ENV_VARS,
         "CUDA_DEVICE_MAX_CONNECTIONS": 32,
         "NCCL_GRAPH_REGISTER": 0,
-        "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True",
-        "TORCH_NCCL_AVOID_RECORD_STREAMS": 1,
+        "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True,graph_capture_record_stream_reuse:True",
+        "TORCH_NCCL_AVOID_RECORD_STREAMS": 0,
         "NCCL_NVLS_ENABLE": 0,
         "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN": 8,
         "NUM_OF_TOKENS_PER_CHUNK_COMBINE_API": 128,

tests/unit_tests/recipes/test_nemotron_3_5_lightning_perf_recipes.py

@@ -57,6 +57,7 @@
     nemotron_3_nano_pretrain_16gpu_h100_fp8cs_config,
 )
 from megatron.bridge.training.config import ConfigContainer
+from megatron.bridge.utils.cuda_graph import cuda_graph_module_names, is_full_iteration_cuda_graph
 
 
 pytestmark = pytest.mark.unit
@@ -107,6 +108,10 @@
     nemotron_3_5_lightning_pretrain_8gpu_gb200_fp8mx_fsdp_config,
     nemotron_3_5_lightning_pretrain_8gpu_gb300_fp8mx_fsdp_config,
 )
+_FULL_ITERATION_MXFP8_RECIPES = (
+    nemotron_3_5_lightning_pretrain_8gpu_gb300_fp8mx_config,
+    nemotron_3_5_lightning_pretrain_8gpu_vr200_fp8mx_config,
+)
 _VR200_RECIPES = (
     nemotron_3_5_lightning_pretrain_8gpu_vr200_bf16_config,
     nemotron_3_5_lightning_pretrain_8gpu_vr200_fp8mx_config,
@@ -344,6 +349,12 @@ def test_nemotron_3_5_perf_recipes_inherit_nemotron_3_policy(
 
     if recipe_factory in _B_MXFP8_RECIPES:
         assert cfg.env_vars == {**base_cfg.env_vars, "NVTE_CUTEDSL_FUSED_GROUPED_MLP": 1}
+    elif recipe_factory in _FULL_ITERATION_MXFP8_RECIPES:
+        assert cfg.env_vars == {
+            **base_cfg.env_vars,
+            "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True,graph_capture_record_stream_reuse:True",
+            "TORCH_NCCL_AVOID_RECORD_STREAMS": 0,
+        }
     elif recipe_factory is not nemotron_3_5_lightning_pretrain_16gpu_h100_bf16_config:
         assert cfg.env_vars == base_cfg.env_vars
     assert cfg.model.calculate_per_token_loss == base_cfg.model.calculate_per_token_loss
@@ -394,6 +405,41 @@ def test_gb_mxfp8_enables_cutedsl_fusion(recipe_factory: Callable[[], ConfigCont
     assert cfg.comm_overlap.delay_wgrad_compute is False
 
 
+@pytest.mark.parametrize("recipe_factory", _FULL_ITERATION_MXFP8_RECIPES, ids=lambda recipe: recipe.__name__)
+def test_mxfp8_full_iteration_graph_config(recipe_factory: Callable[[], ConfigContainer]) -> None:
+    """Lightning full-iteration recipes satisfy graph and static MoE buffer constraints."""
+    cfg = recipe_factory()
+
+    assert cfg.model.cuda_graph_impl == "full_iteration"
+    assert cfg.model.cuda_graph_scope is None
+    assert cfg.model.cuda_graph_modules == []
+    assert is_full_iteration_cuda_graph(cfg.model)
+    assert cfg.model.use_te_rng_tracker is True
+    assert cfg.rng.te_rng_tracker is True
+    assert cfg.rerun_state_machine.check_for_nan_in_loss is False
+    assert cfg.ddp.check_for_nan_in_grad is False
+    assert cfg.model.recompute_granularity is None
+    assert cfg.model.offload_modules == []
+    assert cfg.model.moe_expert_rank_capacity_factor == 1.5
+    assert cfg.model.moe_use_grouped_tensor is True
+    assert cfg.model.moe_paged_stash is False
+    assert cfg.env_vars["TORCH_NCCL_AVOID_RECORD_STREAMS"] == 0
+    assert "graph_capture_record_stream_reuse:True" in cfg.env_vars["PYTORCH_CUDA_ALLOC_CONF"]
+    assert cfg.env_vars["NVTE_CUTEDSL_FUSED_GROUPED_MLP"] == 1
+    assert cfg.env_vars["CUDNNFE_CLUSTER_OVERLAP_MARGIN"] == 8
+    assert cfg.model.use_transformer_engine_op_fuser is True
+    assert cfg.model.moe_mlp_glu_interleave_size is None
+    assert cfg.model.high_priority_a2a_comm_stream is False
+    assert cfg.model.moe_hybridep_num_sms_preprocessing == 108
+    assert cfg.mixed_precision.fp8_dot_product_attention is True
+    assert cfg.comm_overlap.overlap_moe_expert_parallel_comm is False
+    assert cfg.comm_overlap.delay_wgrad_compute is False
+
+    # Validate against MCore as well as checking the unfinalized recipe settings.
+    cfg.model.finalize()
+    assert is_full_iteration_cuda_graph(cfg.model)
+
+
 @pytest.mark.parametrize("recipe_factory", _GB_FSDP_RECIPES, ids=lambda recipe: recipe.__name__)
 def test_gb_mxfp8_fsdp_skips_cutedsl_fusion(recipe_factory: Callable[[], ConfigContainer]) -> None:
     """The Lightning GB MXFP8 FSDP variants remain outside the CutDSL tuning scope."""
@@ -542,6 +588,14 @@ def test_gb200_perf_recipe_topology(recipe_factory: Callable[[], ConfigContainer
     assert cfg.env_vars["NVLINK_DOMAIN_SIZE"] == 72
     assert cfg.env_vars["USE_MNNVL"] == 1
 
+    if recipe_factory is nemotron_3_5_lightning_pretrain_8gpu_gb200_fp8mx_config:
+        assert cfg.model.cuda_graph_impl == "transformer_engine"
+        assert cuda_graph_module_names(cfg.model) == ["attn", "mamba", "moe_router", "moe_preprocess"]
+        assert cfg.model.moe_expert_rank_capacity_factor is None
+        assert cfg.model.moe_paged_stash is False
+        assert cfg.env_vars["PYTORCH_CUDA_ALLOC_CONF"] == "expandable_segments:True"
+        assert cfg.env_vars["TORCH_NCCL_AVOID_RECORD_STREAMS"] == 1
+
 
 @pytest.mark.parametrize("recipe_factory", _GB300_RECIPES, ids=lambda recipe: recipe.__name__)
 def test_gb300_perf_recipe_topology(recipe_factory: Callable[[], ConfigContainer]) -> None:
@@ -619,6 +673,8 @@ def test_gb_fsdp_perf_recipe_defaults(recipe_factory: Callable[[], ConfigContain
     assert cfg.model.cuda_graph_impl == "none"
     assert cfg.model.cuda_graph_scope is None
     assert cfg.model.cuda_graph_modules == []
+    assert cfg.model.moe_expert_rank_capacity_factor is None
+    assert cfg.model.moe_paged_stash is False
     assert cfg.model.init_model_with_meta_device is True
 
     assert cfg.mixed_precision.reuse_grad_buf_for_mxfp8_param_ag is False