Activation Recompute

Activation Recompute

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

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"