Agent skill · software engineering · 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.
Why this skill is useful
Provides specific commands and configurations for activation recompute in Megatron Bridge that are not commonly known.
What it needs
Requires python, uvicorn installed locally. About 11k tokens when loaded. Last updated 2026-08-07. 2,816 stars on the source repository.
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