Agent skill · NVIDIA
nemo-mbridge-perf-activation-recompute
Validate and use selective and full activation recompute in Megatron Bridge to reduce GPU memory usage at the cost of extra compute. Use for activation memory OOMs or regressions involving recompute_granularity, recompute_num_layers, recompute_modules, recompute_method, selective recompute, full recompute, or activation checkpointing.
What it needs
About 11k tokens when loaded.
What this skill does
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 maxmemoryallocated() with maxmemoryreserved() on every rank. 2. Keep an explicit no-recompute control when the workload fits. Under selective granularity, recomputemodules=[] is valid and useful for this comparison. 3. Select the first boundary from the architecture and observed peak: Standard attention: coreattn 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 mlaupproj when expanded Q/K/V projections dominate. Add coreattn only when the attention-core state still matters. Grouped MoE: start with moeact 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. …
How to use it
Reference it in AdaL, Claude Code, Cursor or any coding agent — nothing to install:
@skills NVIDIA/nemo-mbridge-perf-activation-recompute