/nemo-mbridge-perf-activation-recompute

name: nemo-mbridge-perf-activation-recompute
description: Validate and use selective and full activation recompute in Megatron Bridge to reduce GPU memory usage at the cost of extra compute.
license: Apache-2.0
when_to_use: Reducing GPU memory via activation recompute, or investigating a commit that changed recompute settings and caused OOM or a regression; 'recompute_granularity', 'recompute_num_layers', 'recompute_modules', 'recompute_method', 'selective recompute', 'full recompute', 'activation memory OOM'.

Activation Recompute

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

<!-- NVSkills CI refresh: 2026-06-15. No instruction changes. -->

What It Is

Activation recompute trades GPU compute for memory by discarding intermediate activations during the forward pass and recomputing them during backward. Megatron Bridge supports two granularities:

| Granularity | What you specify | What gets recomputed | Memory savings | Compute cost | |---|---|---|---|---| | `selective` | `recompute_modules` list (e.g. `core_attn`, `mlp`) | specific submodules within each layer | moderate (module-dependent) | low to high | | `full` | `recompute_num_layers` + `recompute_method` | entire transformer layers (N layers) | strongest | highest |

Note: MCore names these "selective" (submodule-level) vs "full" (layer-level). "Full" means recomputing full layers, not the full model — you still choose how many layers via `recompute_num_layers`.

Quick Decision

1. Rule out allocator fragmentation first with `PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True`; see @skills/nemo-mbridge-perf-memory-tuning/SKILL.md. 2. For activation pressure, start with selective recompute: `recompute_granularity="selective"` and `recompute_modules=["core_attn"]`. 3. Add modules by cost: `"layernorm"` is cheap but saves little, while `"mlp"` saves much more memory at a clear throughput cost. 4. Use full-layer recompute only when selective recompute does not fit, and set all required fields: `recompute_granularity="full"`, `recompute_method`, and `recompute_num_layers`. 5. With FP8 or TE-scoped CUDA graphs, avoid full-layer recompute unless graph scope is `full_iteration`; otherwise use selective recompute or disable TE graph capture.

CPU offloading (`cpu_offloading=True`) is an alternative that avoids recompute cost entirely, but it is **incompatible with PP > 1**.

Enablement

Selective recompute

cfg.model.recompute_granularity = "selective"
cfg.model.recompute_modules = ["core_attn"]  # add "layernorm", "mlp", or other valid modules as needed

Full-layer recompute

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

Available recompute_modules

| Module | What it recomputes | Compute cost | Memory savings | |---|---|---|---| | `core_attn` | attention softmax/dropout/QKV dot product | low (Flash Attention already recomputes internally) | moderate | | `layernorm` | layer normalization | negligible (~0%) | negligible | | `mlp` | full FFN block | high (~16% on Llama3 70B, hidden=28672) | ~3 GB | | `moe` | MoE expert dispatch | varies | varies | | `moe_act` | MoE activation functions | low | small | | `shared_experts` | shared expert layers | moderate | moderate | | `mla_up_proj` | Multi-Latent Attention up projection | moderate | moderate |

Performance harness CLI

uv run python scripts/performance/run_script.py \
  -m llama \
  -mr llama3_8b \
  --task pretrain \
  -g h100 \
  -c bf16 \
  -ng 8 \
  --recompute_modules core_attn,layernorm \
  ...

Compatibility and Constraints

• `recompute_granularity=selective` requires a non-empty `recompute_modules` list
• `recompute_granularity=full` requires `recompute_method` and `recompute_num_layers`
• **Layer-level recompute (`recompute_granularity="full"` +

`recompute_num_layers`) is incompatible with TE-scoped CUDA graphs.** MCore calls this "full" granularity — the name refers to recomputing full transformer layers, not the full model. Even though you're selecting how many layers to recompute, MCore treats it differently from submodule recompute. Any TE-scoped scope (`attn`, `mlp`, `moe_router`, etc.) will assert. This commonly hits FP8 configs that enable TE-scoped graphs by default (e.g. `LLAMA3_70B_SFT_CONFIG_H100_FP8_CS_V1` sets `cuda_graph_impl="transformer_engine"`, `cuda_graph_scope="mlp"`). Options:

• use submodule recompute (`recompute_granularity="selective"` +

`recompute_modules`) — compatible with TE-scoped graphs

• disable CUDA graphs (`cuda_graph_impl="none"`) and use layer-level recompute
• switch to `cuda_graph_impl="local"`, `cuda_graph_scope="full_iteration"`
• `distribute_saved_activations=True` cannot be combined with `sequence_parallel=True`
• Combining `mlp` + `core_attn` recompute is slightly worse than `mlp` alone

due to double recompute overhead

Measured Results

Llama3 70B SFT on 32x H100 80GB, FP8 (Current Scaling):

• Baseline: TP=4, PP=4, VPP=5, DP=2, MBS=1, GBS=32, seq_len=4096
• Golden GPU utilization: 709.93 TFLOP/s/GPU
• Regression threshold: 5%

| Experiment | recompute_modules | TFLOP/s/GPU | vs Golden | Peak Mem (GB) | Result | |---|---|---|---|---|---| | Baseline | [core_attn] | ~704 | -0.8% | 58.8 (OOM rank0) | OOM | | Exp 1 | [mlp] | 593.6 | -16.4% | 55.6 | Perf regression | | Exp 2 | [mlp, core_attn] | 586.8 | -17.3% | 55.6 | Perf regression | | Exp 3 | [core_attn, layernorm] | ~702 | -1.1% | 59.6 (OOM rank0) | OOM |

Key takeaways:

• `layernorm` recompute is nearly free compute-wise but saves negligible memory
• `mlp` recompute saves ~3 GB peak but costs ~16% because the Llama3 70B FFN

(hidden=28672) is expensive to recompute

• Combining `mlp` + `core_attn` is slightly worse than `mlp` alone
• For this workload, the actual OOM fix was `PYTORCH_CUDA_ALLOC_CONF=expandable_segmen