perf(ncclep): right-size the NCCL EP receive budget

The NCCL EP receive buffer is sized at moe_expert_rank_capacity_factor x the
ideal token count, and every MoE activation saved for backward inherits that
padding. At 1.5 that accounted for 86.8 GB of the 114.1 GB step peak on
nemotron_3_nano_8gpu_gb300_bf16 (+45.9 GB mem-max-allocated vs the HybridEP
baseline) and 135.3 GB on qwen3_30b_a3b gb300 bf16, which OOMed.

These recipes set moe_router_force_load_balancing, so the routing is balanced by
construction. Measured against the ragged HybridEP receive buffers, the worst
case is 1.0072x the ideal token count over 2208 samples for nemotron_3_nano and
1.0050x over 930 samples for qwen3_30b_a3b. Drop the factor to 1.05, which still
leaves 7x (nemotron_3_nano) to 10x (qwen3_30b_a3b) the observed excess and
recovers ~26 GB and ~40 GB respectively. PagedStashRunner is wired on the
capacity factor rather than on paged stashing, so if a routing ever does exceed
the static budget it replays the step dropless and grows the budget: an
underestimate costs a rerun, not correctness.

Also gate moe_paged_stash on the quantized recipes. PagedStashManager only
stashes tensors carrying TE's grouped_tensor_scale_inv attribute, so in BF16
nothing is ever stashed and the run logs "Paged stash: max_tokens_dict is None,
skipping stash buffer allocation" - it was pure hook overhead. The MXFP8 recipes
keep it, where it allocates a real pool and does cut the transient peak.

Signed-off-by: Malay Nagda <malayn@nvidia.com>
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

Author: malay-nagda

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

@@ -87,8 +87,13 @@ def _enable_ncclep(cfg: ConfigContainer, *, mxfp8: bool, moe_a2a_overlap: bool)
     cfg.comm_overlap.delay_wgrad_compute = False
 
     cfg.model.offload_modules = []
-    cfg.model.moe_paged_stash = True
-    cfg.model.moe_expert_rank_capacity_factor = 1.5
+    # The static receive buffer is sized at capacity_factor x the ideal token count, and every MoE
+    # activation saved for backward inherits that padding. These recipes force-balance the router,
+    # so the measured worst case is only 1.008x ideal; 1.5 padded all of it by 50%. PagedStashRunner
+    # replays the step dropless and grows the budget if a routing ever exceeds this.
+    cfg.model.moe_expert_rank_capacity_factor = 1.05
+    # Paged stashing only captures TE's quantized grouped tensors, so it is a no-op outside MXFP8.
+    cfg.model.moe_paged_stash = mxfp8
     cfg.model.moe_paged_stash_buffer_size_factor_cuda = 1.2
     cfg.model.moe_paged_stash_buffer_size_factor_cpu = 1.0
 

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

@@ -70,7 +70,11 @@ def _enable_ncclep_full_iteration_mxfp8(cfg: ConfigContainer) -> None:
     cfg.model.offload_modules = []
     cfg.model.moe_pad_experts_for_cuda_graph_inference = True
     cfg.model.moe_paged_stash = True
-    cfg.model.moe_expert_rank_capacity_factor = 1.5
+    # The static receive buffer is sized at capacity_factor x the ideal token count, and every MoE
+    # activation saved for backward inherits that padding. These recipes force-balance the router,
+    # so the measured worst case is only ~1.005x ideal; 1.5 padded all of it by 50%.
+    # PagedStashRunner replays the step dropless and grows the budget if a routing exceeds this.
+    cfg.model.moe_expert_rank_capacity_factor = 1.05
     cfg.model.moe_paged_stash_buffer_size_factor_cuda = 1.2
     cfg.model.moe_paged_stash_buffer_size_factor_cpu = 1.0
 

src/megatron/bridge/perf_recipes/qwen/gb200/qwen3_moe.py

@@ -281,7 +281,11 @@ def qwen3_30b_a3b_pretrain_8gpu_gb200_bf16_ncclep_config() -> ConfigContainer:
 
     _benchmark_common(cfg, cross_entropy_impl="native")
 
-    cfg.model.moe_expert_rank_capacity_factor = 1.5
+    # The static receive buffer is sized at capacity_factor x the ideal token count, and every MoE
+    # activation saved for backward inherits that padding. This recipe force-balances the router, so
+    # the measured worst case is only ~1.005x ideal; 1.5 padded all of it by 50%. PagedStashRunner
+    # replays the step dropless and grows the budget if a routing ever exceeds this.
+    cfg.model.moe_expert_rank_capacity_factor = 1.05
     cfg.model.moe_grouped_gemm = True
     cfg.model.use_transformer_engine_op_fuser = True
 

src/megatron/bridge/perf_recipes/qwen/gb300/qwen3_moe.py

@@ -291,10 +291,13 @@ def qwen3_30b_a3b_pretrain_8gpu_gb300_bf16_ncclep_config() -> ConfigContainer:
     # cfg.model.moe_router_padding_for_quantization = True
 
     cfg.model.offload_modules = []
-    cfg.model.moe_paged_stash = True
-    cfg.model.moe_expert_rank_capacity_factor = 1.5
-    cfg.model.moe_paged_stash_buffer_size_factor_cuda = 1.2
-    cfg.model.moe_paged_stash_buffer_size_factor_cpu = 1.0
+    # Paged stashing only captures TE's quantized grouped tensors, so it is a no-op in BF16.
+    cfg.model.moe_paged_stash = False
+    # The static receive buffer is sized at capacity_factor x the ideal token count, and every MoE
+    # activation saved for backward inherits that padding. This recipe force-balances the router, so
+    # the measured worst case is only ~1.005x ideal; 1.5 padded all of it by 50%. PagedStashRunner
+    # replays the step dropless and grows the budget if a routing ever exceeds this.
+    cfg.model.moe_expert_rank_capacity_factor = 1.05
 
     # Keep process settings next to the recipe so users can see the exact benchmark environment.
     cfg.env_vars = {

tests/unit_tests/recipes/test_nemotron_3_5_lightning_perf_recipes.py

@@ -226,8 +226,7 @@ def test_ncclep_perf_recipe_defaults(recipe_factory: Callable[[], ConfigContaine
     assert cfg.model.moe_mlp_glu_interleave_size == 32
 
     assert cfg.model.offload_modules == []
-    assert cfg.model.moe_paged_stash is True
-    assert cfg.model.moe_expert_rank_capacity_factor == 1.5
+    assert cfg.model.moe_expert_rank_capacity_factor == 1.05
     assert cfg.model.moe_paged_stash_buffer_size_factor_cuda == 1.2
     assert cfg.model.moe_paged_stash_buffer_size_factor_cpu == 1.0
 
@@ -248,8 +247,11 @@ def test_ncclep_perf_recipe_defaults(recipe_factory: Callable[[], ConfigContaine
 
     if "fp8mx" in recipe_factory.__name__:
         assert cfg.model.moe_router_padding_for_quantization is True
+        assert cfg.model.moe_paged_stash is True
     else:
         assert cfg.model.moe_router_padding_for_quantization is False
+        # Paged stashing only captures TE's quantized grouped tensors, so it is a no-op in BF16.
+        assert cfg.model.moe_paged_stash is False
 
 
 @pytest.mark.parametrize(

tests/unit_tests/recipes/test_qwen_perf_recipes.py

@@ -68,7 +68,7 @@ def test_qwen3_30b_gb200_bf16_ncclep_config():
 
     assert cfg.model.moe_token_dispatcher_type == "flex"
     assert cfg.model.moe_flex_dispatcher_backend == "ncclep"
-    assert cfg.model.moe_expert_rank_capacity_factor == 1.5
+    assert cfg.model.moe_expert_rank_capacity_factor == 1.05
     assert cfg.model.moe_grouped_gemm is True
     assert cfg.model.use_transformer_engine_op_fuser is True
     assert cfg.model.moe_ncclep_zero_copy is False
@@ -83,7 +83,7 @@ def test_qwen3_30b_gb200_mxfp8_ncclep_config():
 
     assert cfg.model.moe_token_dispatcher_type == "flex"
     assert cfg.model.moe_flex_dispatcher_backend == "ncclep"
-    assert cfg.model.moe_expert_rank_capacity_factor == 1.5
+    assert cfg.model.moe_expert_rank_capacity_factor == 1.05
     assert cfg.model.moe_ncclep_zero_copy is False
     assert cfg.model.moe_paged_stash is True
     assert cfg.model.moe_grouped_gemm is True
@@ -101,7 +101,9 @@ def test_qwen3_30b_gb300_bf16_ncclep_config():
 
     assert cfg.model.moe_token_dispatcher_type == "flex"
     assert cfg.model.moe_flex_dispatcher_backend == "ncclep"
-    assert cfg.model.moe_expert_rank_capacity_factor == 1.5
+    assert cfg.model.moe_expert_rank_capacity_factor == 1.05
+    # Paged stashing only captures TE's quantized grouped tensors, so it is a no-op in BF16.
+    assert cfg.model.moe_paged_stash is False
     assert cfg.model.moe_grouped_gemm is True
     assert cfg.model.use_transformer_engine_op_fuser is True
     assert cfg.model.moe_ncclep_zero_copy is False
@@ -119,7 +121,7 @@ def test_qwen3_30b_gb300_mxfp8_ncclep_config():
 
     assert cfg.model.moe_token_dispatcher_type == "flex"
     assert cfg.model.moe_flex_dispatcher_backend == "ncclep"
-    assert cfg.model.moe_expert_rank_capacity_factor == 1.5
+    assert cfg.model.moe_expert_rank_capacity_factor == 1.05
     assert cfg.model.moe_ncclep_zero_copy is False
     assert cfg.model.moe_paged_stash is True
     assert cfg.model.moe_grouped_gemm is True