Nemo Mbridge Perf Activation Recompute

Activation Recompute

Stable docs: @docs/training/activation-recomputation.md Card: @skills/nemo-mbridge-perf-activation-recompute/card.yaml

<!-- Guidance refreshed: 2026-08-12. -->

Activation recompute (activation checkpointing) trades additional forward work during backward for lower retained-activation memory. The useful checkpoint boundary depends on the model architecture, attention backend, parallelism, and the tensor that actually drives the per-rank peak.

Quick Decision Guide

1. Confirm the pressure is real allocation, not allocator fragmentation. Compare max_memory_allocated() with max_memory_reserved() on every rank.
2. Keep an explicit no-recompute control when the workload fits. Under selective granularity, recompute_modules=[] is valid and useful for this comparison.
3. Select the first boundary from the architecture and observed peak:
   • Standard attention: core_attn is the common first candidate. It is strongest when unfused attention materializes score/probability tensors. With Transformer Engine fused or Flash Attention, compare it against [] because those backends already rematerialize attention internals.
   • Multi-Latent Attention (MLA): start with mla_up_proj when expanded Q/K/V projections dominate. Add core_attn only when the attention-core state still matters.
   • Grouped MoE: start with moe_act when the expert intermediate activation dominates; add layernorm when norm outputs are material. Use whole moe recompute only after accounting for the extra expert compute and communication it replays.
   • Dense FFN: mlp can save the whole dense-MLP activation region, but it usually costs more compute than a narrow output-discard boundary.
4. Change one label at a time. Record per-rank allocated/reserved peaks plus steady-state step time or throughput; do not infer a global module ranking from one recipe.
5. Use full-layer recompute only when targeted selective boundaries do not make the workload fit. Full recompute has the broadest memory effect and the largest replay cost.
6. Treat CUDA graphs, FP8, context-parallel communication, and overlap features as compatibility constraints, not afterthoughts.

Megatron Core's cpu_offloading=True is an alternative when PCIe/NVLink transfer overhead is preferable to replayed compute. It cannot be combined with activation recompute and is not compatible with pipeline parallelism greater than one.

Enablement

Selective recompute

cfg.model.recompute_granularity = "selective"
cfg.model.recompute_modules = ["core_attn"]  # Common standard-attention candidate, not a universal default.

Use the decision table below to replace or extend that list for MLA, MoE, dense-MLP, or GDN workloads.

Full-layer recompute

cfg.model.recompute_granularity = "full"
cfg.model.recompute_method = "uniform"
cfg.model.recompute_num_layers = 1

• uniform: checkpoint fixed groups of recompute_num_layers transformer layers.
• block: checkpoint the first recompute_num_layers layers on each pipeline stage, with virtual-pipeline-aware distribution.

Selective Module Decision Table

The currently pinned Megatron Core accepts these labels. A development branch can add model-specific labels, so validate against the exact target revision rather than copying a list across branches.

For example, DeepSeek V4 configurations can use the model-specific mhc label only with their required Megatron Core development branch. It is not a portable label for the pinned revision and therefore is not included in the table above.

Common performance configurations consequently fall into several patterns rather than one universal list:

• standard transformer recipes often use core_attn;
• MLA recipes often use mla_up_proj, sometimes with mlp;
• grouped-MoE recipes often use moe_act or layernorm plus moe_act;
• higher-pressure MoE recipes sometimes use broader combinations such as moe plus layernorm.

These are candidate patterns, not an ordering guarantee. Peak attribution and matched measurements decide the final list.

Measurement Contract

For every candidate, capture:

• exact Bridge and Megatron Core revisions;
• model, sequence length, micro/global batch sizes, precision, attention backend, and parallelism;
• the exact recompute_granularity, module list, method, and layer count;
• per-rank max_memory_allocated() and max_memory_reserved();
• steady-state step time or throughput after warmup;
• a short convergence or numerical-sanity check appropriate to the task.

Use a matched no-recompute control and change one recompute choice at a time. Peak memory from different jobs, backends, or parallel layouts is not a module-ranking benchmark.

Do not call a candidate successful merely because it advances farther than the control. Run through optimizer-state initialization and multiple steady-state steps: selective recompute can move the memory wall from forward into gradient synchronization or the optimizer without making the workload viable.

Matched H100 Evidence: Moonlight 16B

A 2026-08-12 short-run study used the exact Bridge revision 600d069b824dd5ce50367a311a5a3244478faf22 and Megatron Core revision 24bad8e677d22625d86ef2a54c9506b6e4992c93. The Moonlight 16B BF16 pretraining recipe ran on 8 H100 80GB GPUs with sequence length 4096, MBS=1, GBS=4, TP=2, PP=1, CP=1, EP=8, mock data, and 20 steps. This model mixes one dense layer with 26 MLA+MoE layers. Each row changed only recompute_modules; all 20 losses were finite with zero skipped or NaN iterations.

Peak allocated memory is the maximum post-optimizer value reported after iteration 2. Time and throughput are means over iterations 11--20.

For this exact workload, moe_act is the best first boundary: it recovered nearly as much allocated memory as mla_up_proj for less replay cost. mla_up_proj is the next candidate if its roughly 39 MB additional reduction matters. Adding mlp to mla_up_proj or layernorm to moe_act did not improve the observed peak and made steps slower. Explicit core_attn added cost without material memory benefit under fused attention.

Maximum reserved memory stayed near 40 GB and did not fall monotonically. That is allocator caching, not contrary evidence: boundary selection in this study is based on allocated memory and successful end-to-end steps.

Matched H100 Evidence: Nemotron 3 Nano

The same 2026-08-12 study used the native 16-H100 BF16 performance recipe for the 52-layer hybrid Mamba/fused-attention MoE model. The matched short-run configuration used sequence length 8192, MBS=1, GBS=16, TP=1, PP=1, CP=1, EP=8, DP=16, expert-DP=2, HybridEP, grouped GEMM, TE CUDA graphs for attention and Mamba, mock data, and 12 steps. Each row changed only recompute_modules.

Both successful rows had finite losses and zero skipped or NaN iterations. For this exact capacity-limited recipe, whole-moe recompute is the smallest tested passing boundary. Adding layernorm recovered another 1.014 GB (1.57%) of rank-0 peak at 3.07% higher step time, so the recipe's broader combination is justified when that headroom is required. Narrow moe_act produced real activation relief but did not make the whole training step viable.

An exploratory native 8-H100 layout failed during FP32 optimizer-state initialization even at sequence length 4096. That is optimizer capacity, not a selective-boundary throughput baseline; no timing comparison from those runs is used here.