Merge upstream main into PR 5524

Signed-off-by: Chen Cui <chcui@nvidia.com>

Author: cuichenx

.github/workflows/config/.secrets.baseline

@@ -17,9 +17,12 @@ As a rule of thumb:
 
 - start with **selective recomputation** before using full recomputation
 - use **full recomputation** only when selective recomputation still does not fit
-- for MoE and long-context training, prefer recomputing smaller modules such as
-  normalization, activation, MoE-side, or model-specific up-projection work
-  rather than recomputing the whole layer
+- choose selective modules from the model and measured peak: `core_attn` is a
+  common standard-attention candidate, MLA often benefits from `mla_up_proj`,
+  and grouped MoE often starts with `moe_act` or `layernorm` plus `moe_act`
+- with TE fused or Flash Attention, compare `core_attn` against an empty
+  selective module list because the backend already rematerializes attention
+  internals
 - revisit recomputation after enabling CUDA graphs, because TE-scoped graphs and
   full recomputation are not always compatible
 
@@ -86,7 +89,10 @@ When training with pipeline parallelism:
 
 ## Self-attention Recomputation
 
-Megatron Bridge supports selective self-attention recomputation that checkpoints the inputs of each self-attention block and recomputes the intermediate input activations. This cost-efficient method achieves high memory savings with minimal recomputation cost.
+Megatron Bridge supports selective self-attention recomputation that checkpoints
+the core-attention boundary and recomputes it during backward. This can be a
+cost-efficient choice for standard attention, but it is not the universal first
+boundary for MLA, MoE, or fused-attention workloads.
 
 The intermediate layers of the self-attention block account for a large share
 of activation memory because softmax, dropout, and QKV dot-product attention
@@ -105,36 +111,38 @@ from megatron.bridge.models import GPTModelProvider
 
 model_config = GPTModelProvider(
     recompute_granularity="selective",  # Enable selective recomputation
-    recompute_modules=["core_attn"],    # Recompute attention modules (default)
+    recompute_modules=["core_attn"],    # Common standard-attention candidate and MCore default
     # ... other model parameters
 )
 ```
 
 ### Recomputation Modules
 
-Megatron Bridge supports selective recomputation for various modules:
+The pinned Megatron Core accepts these selective labels. Do not combine them
+blindly: some are architecture-specific and some checkpoint overlapping
+regions.
 
-```python
-model_config = GPTModelProvider(
-    recompute_granularity="selective",
-    recompute_modules=[
-        "core_attn",      # Core attention computation (default)
-        "mlp",            # MLP layers
-        "layernorm",      # Layer normalization
-        "moe",            # Mixture of Experts layers
-        "moe_act",        # MoE activation functions
-        "shared_experts", # Shared expert layers
-        "mla_up_proj",    # Multi-Latent Attention up projection
-    ],
-)
-```
+| Label | Boundary and typical use |
+|---|---|
+| `core_attn` | Core attention; a common standard-attention candidate that can replay context-parallel communication |
+| `mla_up_proj` | Expanded MLA Q/KV projections and RoPE; often the first MLA candidate |
+| `moe_act` | Grouped-expert activation output without replaying dispatch or expert GEMMs |
+| `layernorm` | Input and pre-MLP normalization outputs, often paired with a MoE or MLA boundary |
+| `mlp` | Whole dense MLP; broader replay and no effect on pure-MoE layers |
+| `moe` | Whole MoE forward, including routing, expert compute, and communication |
+| `shared_experts` | Non-overlapped shared-expert MLP |
+| `gdn_norm_out` | GatedDeltaNet output norm and HP-to-CP communication |
+
+An empty list is valid under selective granularity and is useful as a matched
+no-recompute control.
 
 ### Flash Attention Integration
 
-Self-attention recomputation is automatically enabled when using Flash Attention
-through Transformer Engine. Flash Attention already recovers some memory by
-recomputing attention scores rather than storing them, so extra explicit
-attention recomputation is often less important than recomputing other modules.
+Flash Attention through Transformer Engine already recovers memory by
+rematerializing attention internals. That does not automatically make
+`recompute_modules=["core_attn"]` the right explicit setting. Compare it with an
+empty selective module list; MLA up projections, MoE activations, normalization,
+or another boundary may set the actual peak.
 
 ## Advanced Recomputation Configuration
 
@@ -169,29 +177,38 @@ model_config = GPTModelProvider(
     # MoE configuration
     num_moe_experts=8,
     expert_model_parallel_size=2,
+    moe_grouped_gemm=True,
 
     # MoE recomputation
     recompute_granularity="selective",
-    recompute_modules=["moe", "moe_act"],  # Recompute MoE-specific modules
+    recompute_modules=["moe_act", "layernorm"],  # Narrow MoE-side candidates
 )
 ```
 
-For MoE training, it is often better to start with selective recomputation of
-MoE-side modules, normalization, activation functions, or model-specific
-up-projection modules than to enable blanket full recomputation immediately.
+`moe_act` requires grouped-GEMM experts. Whole-`moe` recompute is a broader
+alternative that replays routing, dispatch/combine communication, expert
+compute, and shared-expert work; it is incompatible with expert-parallel
+overlap. `shared_experts` recompute is likewise incompatible with shared-expert
+overlap. Prefer the narrow boundary that satisfies the memory target.
 
 ## Feature Interactions
 
-- TE-scoped CUDA graphs are usually paired with selective recomputation rather
-  than full recomputation.
+- Full recomputation with CUDA graphs requires
+  `cuda_graph_impl="full_iteration"` in the pinned Megatron Core. Otherwise use
+  selective recomputation or disable CUDA graphs.
+- A selective checkpoint boundary must lie wholly inside or wholly outside its
+  CUDA-graph scope.
 - MoE communication overlap paths often require recomputation settings that are
   more selective than "full."
 - At long context, recomputing SDPA-heavy attention internals can cost more than
   recomputing smaller supporting modules.
+- Advancing from a forward OOM to a gradient-synchronization or optimizer OOM
+  is useful diagnosis, but not a pass. Validate through optimizer-state
+  initialization and multiple steady-state iterations.
 
 ## Related Docs
 
 - [docs/training/cuda-graphs.md](cuda-graphs.md)
 - [docs/training/moe-optimization.md](moe-optimization.md)
-- [skills/nemo-mbridge-perf-activation-recompute/SKILL.md](../skills/nemo-mbridge-perf-activation-recompute/SKILL.md) — per-module cost/savings data, measured results
-- [skills/nemo-mbridge-perf-memory-tuning/SKILL.md](../skills/nemo-mbridge-perf-memory-tuning/SKILL.md) — expandable segments, parallelism resizing, and other memory reduction strategies
+- [skills/nemo-mbridge-perf-activation-recompute/SKILL.md](../../skills/nemo-mbridge-perf-activation-recompute/SKILL.md) — architecture-specific module selection, compatibility, and measurement guidance
+- [skills/nemo-mbridge-perf-memory-tuning/SKILL.md](../../skills/nemo-mbridge-perf-memory-tuning/SKILL.md) — expandable segments, parallelism resizing, and other memory reduction strategies

docs/training/activation-recomputation.md

@@ -172,6 +172,7 @@ items:
         final_loss: null
         last_10_steps_step_time_ms_avg: null
         last_10_steps_model_tflops_per_gpu_avg: null
+        last_10_steps_tokens_per_second_per_gpu_avg: null
       expected_result: >
         Megatron Bridge does not provide a public Gemma 3 VL 4B pretraining
         recipe. The available exact-model recipes cover SFT and PEFT, so a
@@ -197,10 +198,11 @@ items:
         final_loss: null
         last_10_steps_step_time_ms_avg: null
         last_10_steps_model_tflops_per_gpu_avg: null
+        last_10_steps_tokens_per_second_per_gpu_avg: null
       expected_result: >
         Pending exact-revision CORD-v2 full-SFT verification with the public
         Gemma 3 VL 4B H100 recipe. The run must complete at least ten optimizer
-        steps with finite loss, no skipped or NaN iterations, all four metrics,
+        steps with finite loss, no skipped or NaN iterations, all five metrics,
         and a reloadable final full-model checkpoint.
 
   sft_export_inference:
@@ -228,6 +230,7 @@ items:
         final_loss: null
         last_10_steps_step_time_ms_avg: null
         last_10_steps_model_tflops_per_gpu_avg: null
+        last_10_steps_tokens_per_second_per_gpu_avg: null
       expected_result: >
         Gemma 3 VL image-bidirectional attention uses dense local/global
         additive biases with FusedAttention. The provider rejects context
@@ -255,10 +258,11 @@ items:
         final_loss: null
         last_10_steps_step_time_ms_avg: null
         last_10_steps_model_tflops_per_gpu_avg: null
+        last_10_steps_tokens_per_second_per_gpu_avg: null
       expected_result: >
         Pending exact-revision CORD-v2 LoRA verification with the public Gemma 3
         VL 4B H100 recipe. The run must complete at least ten optimizer steps
-        with finite loss, no skipped or NaN iterations, all four metrics, and a
+        with finite loss, no skipped or NaN iterations, all five metrics, and a
         reloadable adapter checkpoint.
 
   checkpoint_resume:
@@ -273,6 +277,7 @@ items:
         final_loss: null
         last_10_steps_step_time_ms_avg: null
         last_10_steps_model_tflops_per_gpu_avg: null
+        last_10_steps_tokens_per_second_per_gpu_avg: null
       resume_comparison: null
       expected_result: >
         Checkpoint-resume verification depends on a supported pretraining

examples/model_verification_cards/gemma-3-4b-it/card.yaml

@@ -172,6 +172,7 @@ items:
         final_loss: null
         last_10_steps_step_time_ms_avg: null
         last_10_steps_model_tflops_per_gpu_avg: null
+        last_10_steps_tokens_per_second_per_gpu_avg: null
       expected_result: >
         Megatron Bridge does not publish a pretraining recipe for this Gemma 4
         conditional-generation variant; the current model package provides SFT
@@ -199,6 +200,7 @@ items:
         final_loss: 0.09127883
         last_10_steps_step_time_ms_avg: 15936.25
         last_10_steps_model_tflops_per_gpu_avg: 32.20
+        last_10_steps_tokens_per_second_per_gpu_avg: 1028.096
         peak_allocated_memory_gib: 65.223
         peak_reserved_memory_gib: 72.356
       expected_result: >
@@ -244,11 +246,12 @@ items:
         final_loss: null
         last_10_steps_step_time_ms_avg: null
         last_10_steps_model_tflops_per_gpu_avg: null
+        last_10_steps_tokens_per_second_per_gpu_avg: null
       expected_result: >
         No Gemma 4 VL long-context recipe or run was verified. A future
         workload must establish sequence-packing and context-parallel settings,
         complete at least 10 optimizer steps with finite loss and no skipped or
-        NaN iterations, report all four required metrics, and save a reloadable
+        NaN iterations, report all five required metrics, and save a reloadable
         final checkpoint.
 
   peft:
@@ -273,6 +276,7 @@ items:
         final_loss: 0.1072377
         last_10_steps_step_time_ms_avg: 14163.25
         last_10_steps_model_tflops_per_gpu_avg: 65.01
+        last_10_steps_tokens_per_second_per_gpu_avg: 2313.593
         peak_allocated_memory_gib: 44.252
         peak_reserved_memory_gib: 45.637
       expected_result: >
@@ -303,6 +307,7 @@ items:
         final_loss: null
         last_10_steps_step_time_ms_avg: null
         last_10_steps_model_tflops_per_gpu_avg: null
+        last_10_steps_tokens_per_second_per_gpu_avg: null
       resume_comparison: null
       expected_result: >
         Megatron Bridge does not provide a Gemma 4 VLM pretraining recipe, so

examples/model_verification_cards/gemma-4-26b-a4b-it/card.yaml

@@ -177,6 +177,7 @@ items:
         final_loss: 6.385243
         last_10_steps_step_time_ms_avg: 31204.920
         last_10_steps_model_tflops_per_gpu_avg: 54.595
+        last_10_steps_tokens_per_second_per_gpu_avg: 190.926
       expected_result: >
         The uninterrupted 352-H100 TP1/PP11/CP1/EP32/ETP1 run completed 100
         real-data steps at GBS/MBS 1024/1 with cuDNN DSA, sparse indexer loss
@@ -216,6 +217,7 @@ items:
         final_loss: 6.523522
         last_10_steps_step_time_ms_avg: 64786.540
         last_10_steps_model_tflops_per_gpu_avg: 102.520
+        last_10_steps_tokens_per_second_per_gpu_avg: 337.189
       expected_result: >
         The 192-GB200 TP1/PP6/CP1/EP32/ETP1 run completed exactly 100
         RP2 head_01 pretraining steps at GBS/MBS 1024/1 with cuDNN DSA, sparse
@@ -254,6 +256,7 @@ items:
         final_loss: 0.6634203
         last_10_steps_step_time_ms_avg: 4436.290
         last_10_steps_model_tflops_per_gpu_avg: 10.154
+        last_10_steps_tokens_per_second_per_gpu_avg: 35.511
       expected_result: >
         The immutable-revision 416-H100 run completes exactly 100 full-SFT
         steps at TP1/PP13/CP2/EP32/ETP1, GBS/MBS 32/1, with cuDNN DSA, sparse
@@ -286,6 +289,7 @@ items:
         final_loss: 0.3682597
         last_10_steps_step_time_ms_avg: 8593.990
         last_10_steps_model_tflops_per_gpu_avg: 12.446
+        last_10_steps_tokens_per_second_per_gpu_avg: 39.718
       expected_result: >
         The immutable-revision 192-GB200 run completed exactly 100 full-SFT
         steps at TP1/PP6/CP4/EP32/ETP1, DP8, and GBS/MBS 8/1 without gradient
@@ -363,6 +367,7 @@ items:
         final_loss: 2.557671
         last_10_steps_step_time_ms_avg: 36603.370
         last_10_steps_model_tflops_per_gpu_avg: 39.648
+        last_10_steps_tokens_per_second_per_gpu_avg: 116.828
       expected_result: >
         The 608-H100 TP1/PP19/CP32/EP32/ETP1 run completed exactly 20 optimizer
         steps over 200,000-token packed rows with cuDNN DSA, sparse indexer loss
@@ -394,6 +399,7 @@ items:
         final_loss: 0.000511673
         last_10_steps_step_time_ms_avg: 143248.820
         last_10_steps_model_tflops_per_gpu_avg: 88.848
+        last_10_steps_tokens_per_second_per_gpu_avg: 266.874
       expected_result: >
         The immutable-revision 192-GB200 run completed exactly 20 optimizer
         steps at TP1/PP6/CP32/EP32/ETP1, DP1, and GBS/MBS 56/1 with cuDNN DSA,
@@ -427,6 +433,7 @@ items:
         final_loss: 0.8835775
         last_10_steps_step_time_ms_avg: 4638.550
         last_10_steps_model_tflops_per_gpu_avg: 19.423
+        last_10_steps_tokens_per_second_per_gpu_avg: 67.926
       expected_result: >
         The immutable-revision 208-H100 run completes exactly 100 LoRA steps at
         TP1/PP13/CP1/EP16/ETP1, GBS/MBS 32/1, with cuDNN DSA, sparse indexer
@@ -455,6 +462,7 @@ items:
         final_loss: 0.8957523
         last_10_steps_step_time_ms_avg: 68588.270
         last_10_steps_model_tflops_per_gpu_avg: 1.423
+        last_10_steps_tokens_per_second_per_gpu_avg: 4.977
       expected_result: >
         The immutable-revision 192-GB200 run completed exactly 100 LoRA steps
         at TP1/PP6/CP1/EP32/ETP1, GBS/MBS 32/1, with offline-packed Tulu 3
@@ -491,6 +499,7 @@ items:
         final_loss: 6.372624
         last_10_steps_step_time_ms_avg: 32002.290
         last_10_steps_model_tflops_per_gpu_avg: 53.235
+        last_10_steps_tokens_per_second_per_gpu_avg: 186.168
       resume_comparison:
         reference_item: pretrain
         sentinel_steps: [51, 100]
@@ -536,6 +545,7 @@ items:
         final_loss: 6.503141
         last_10_steps_step_time_ms_avg: 60079.390
         last_10_steps_model_tflops_per_gpu_avg: 110.552
+        last_10_steps_tokens_per_second_per_gpu_avg: 363.608
       resume_comparison:
         reference_item: pretrain
         sentinel_steps: [51, 100]

examples/model_verification_cards/glm5-2/card.yaml

@@ -139,9 +139,10 @@ items:
         final_loss: null
         last_10_steps_step_time_ms_avg: null
         last_10_steps_model_tflops_per_gpu_avg: null
+        last_10_steps_tokens_per_second_per_gpu_avg: null
       expected_result: >
         A public GLM-5 H100 recipe must complete a bounded 100-step run with
-        finite loss, no skipped or NaN iterations, all four metrics, and a
+        finite loss, no skipped or NaN iterations, all five metrics, and a
         reloadable final checkpoint. Training is deferred by this card.
 
   sft:
@@ -156,9 +157,10 @@ items:
         final_loss: null
         last_10_steps_step_time_ms_avg: null
         last_10_steps_model_tflops_per_gpu_avg: null
+        last_10_steps_tokens_per_second_per_gpu_avg: null
       expected_result: >
         A pinned-data 100-step full-SFT run must finish with finite loss, no
-        skipped or NaN iterations, all four metrics, and a reloadable final
+        skipped or NaN iterations, all five metrics, and a reloadable final
         checkpoint. Training is deferred by this card.
 
   sft_export_inference:
@@ -185,9 +187,10 @@ items:
         final_loss: null
         last_10_steps_step_time_ms_avg: null
         last_10_steps_model_tflops_per_gpu_avg: null
+        last_10_steps_tokens_per_second_per_gpu_avg: null
       expected_result: >
         A dedicated packed long-context SFT run must complete a bounded run
-        with finite loss, no skipped or NaN iterations, all four metrics, and
+        with finite loss, no skipped or NaN iterations, all five metrics, and
         a reloadable checkpoint. Training is deferred by this card.
 
   peft:
@@ -202,9 +205,10 @@ items:
         final_loss: null
         last_10_steps_step_time_ms_avg: null
         last_10_steps_model_tflops_per_gpu_avg: null
+        last_10_steps_tokens_per_second_per_gpu_avg: null
       expected_result: >
         A GLM-5 PEFT recipe with an audited adapter target set must complete a
-        bounded run with finite loss, all four metrics, and a reloadable
+        bounded run with finite loss, all five metrics, and a reloadable
         adapter checkpoint. Training is deferred by this card.
 
   checkpoint_resume:
@@ -219,6 +223,7 @@ items:
         final_loss: null
         last_10_steps_step_time_ms_avg: null
         last_10_steps_model_tflops_per_gpu_avg: null
+        last_10_steps_tokens_per_second_per_gpu_avg: null
       resume_comparison:
         reference_item: pretrain
         sentinel_steps: [51, 100]