nemo-mbridge-perf-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. For final library, performance, and verification-card recipes, no recompute and selective recompute on cheap supported modules are both appropriate starting choices. Choose by measured memory headroom and steady-state performance; a no-recompute OOM is not a prerequisite for selective recompute. Keep an explicit no-recompute control when it fits. Under selective granularity, recompute_modules=[] is valid and useful for this comparison.
3. For selective recompute, start with cheap supported boundaries selected 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. Full-layer recompute is allowed temporarily for diagnostic bring-up. Retain it in a final recipe only when target-scale measurements show that no recompute and supported selective boundaries still fail to fit, or matched runs show better end-to-end throughput with full recompute (for example, by enabling larger MBS or lower TP). Document the memory/OOM evidence, boundaries and layouts tested, and steady-state performance of the alternatives. Full recompute can severely reduce throughput and must not become the default from a bring-up success alone.
6. Treat CUDA graphs, FP8, context-parallel communication, and overlap features as compatibility constraints, not afterthoughts.