Nemo Mbridge Perf Activation Recompute

NVIDIAOfficialPublished by NVIDIA

Skill
Guidelines
Software development
Apache-2.0
Medium risk

From the author

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.

— the description by NVIDIA

AutoGineer summary

Written by AutoGineer from the skill’s own files, not by the author

This skill provides a structured approach to managing GPU memory pressure in Megatron Bridge by implementing activation recomputation (checkpointing). By trading additional compute cycles for reduced memory footprint, users can resolve out-of-memory (OOM) errors. The guide emphasizes a data-driven, iterative process of selecting specific module boundaries based on model architecture—such as standard attention, MLA, or MoE—and validating results through matched measurements of peak memory and throughput.

What it does

Provides a decision framework for choosing between selective and full-layer activation recomputation.
Lists specific module boundaries for various architectures including MLA, MoE, and dense MLP.
Defines a rigorous measurement contract to ensure performance impact is accurately tracked.
Outlines compatibility constraints with features like CUDA graphs, FP8, and expert-parallel overlap.
Offers a diagnostic table to troubleshoot common OOM scenarios and recompute failures.
Provides configuration examples for both selective module lists and full-layer uniform/block methods.

When to use it

When encountering GPU memory OOM errors during model training.
When memory usage needs to be reduced to fit larger batch sizes or longer sequences.
When investigating performance regressions related to activation memory.
When evaluating the trade-off between memory capacity and compute throughput for a specific model architecture.

Advantages

Reduces peak GPU memory allocation by discarding intermediate activations.
Allows for fine-grained control over which model components are recomputed.
Provides architecture-specific guidance to minimize the compute penalty of recomputation.
Includes a structured validation process to avoid misleading performance conclusions.

Usage examples

“How do I enable selective recomputation for a standard attention model?”: The agent will provide the configuration snippet setting recompute_granularity to selective and suggest starting with core_attn as the module boundary.
“I am getting OOMs in an MoE model; which recompute modules should I test?”: The agent will recommend testing moe_act or layernorm first, and only moving to whole-moe recompute if narrower boundaries fail to resolve the memory pressure.

Install

Download the skill as a ZIP:
https://autogineer.com/agent-skills/nvidia/nemo-mbridge-perf-activation-recompute/download

Unzip it into the cross-client skills folder of your project:

mkdir -p .agents/skills && unzip nemo-mbridge-perf-activation-recompute.zip -d .agents/skills

Documentation by NVIDIA

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

Confirm the pressure is real allocation, not allocator fragmentation. Compare max_memory_allocated() with max_memory_reserved() on every rank.
Keep an explicit no-recompute control when the workload fits. Under selective granularity, recompute_modules=[] is valid and useful for this comparison.
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.
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.
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.
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"
cfg.model.recompute_num_layers = 1

uniform: checkpoint fixed groups of recompute_num_layers transformer layers.
block: checkpoint the first recompute_num_layers layers on each pipeline stage, with virtual-pipeline-aware distribution.

Selective Module Decision Table

The currently pinned Megatron Core accepts these labels. A development branch can add model-specific labels, so validate against the exact target revision rather than copying a list across branches.

Module	Checkpoint boundary	When to test it	Main cost or caveat
`core_attn`	Core attention	Standard attention, especially an unfused backend retaining attention intermediates	Replays attention. Incremental savings can be small with TE fused/Flash Attention; context parallelism can replay attention communication.
`mla_up_proj`	MLA Q/KV up-projection plus RoPE region	MLA models retaining expanded Q/K/V tensors	Replays the MLA expansion path. It is a distinct, potentially additive boundary from `core_attn`.
`layernorm`	Input and pre-MLP normalization outputs	Norm outputs contribute materially to the peak, often alongside MoE or MLA boundaries	Usually narrow, but savings depend on hidden size, sequence length, and which graph paths are active.
`moe_act`	Activation output between grouped expert FC1 and FC2	Grouped MoE expert-intermediate activations dominate	Narrow output-discard checkpoint. It does not replay dispatch, FC1, or FC2, but has FP8 delayed-scaling restrictions.
`mlp`	Whole dense MLP	Dense layers dominate after narrower boundaries are exhausted	Replays the complete dense MLP. It has no effect on layers whose MLP is MoE.
`moe`	Whole MoE forward	A broad MoE region must be discarded to make the workload fit	Replays routing, dispatch/combine communication, experts, and shared-expert work. It is incompatible with expert-parallel overlap.
`shared_experts`	Non-overlapped shared-expert MLP	Shared experts are a distinct material peak	Replays the shared-expert MLP and is invalid with shared-expert overlap. Outer `moe` already removes its original-forward saves, but nesting can still change the transient backward-replay peak.
`gdn_norm_out`	GDN gated-normalization output	GDN/hybrid models retain this output	Replays the normalization and its HP-to-CP all-to-all path.

For example, DeepSeek V4 configurations can use the model-specific mhc label only with their required Megatron Core development branch. It is not a portable label for the pinned revision and therefore is not included in the table above.

Common performance configurations consequently fall into several patterns rather than one universal list:

standard transformer recipes often use core_attn;
MLA recipes often use mla_up_proj, sometimes with mlp;
grouped-MoE recipes often use moe_act or layernorm plus moe_act;
higher-pressure MoE recipes sometimes use broader combinations such as moe plus layernorm.

These are candidate patterns, not an ordering guarantee. Peak attribution and matched measurements decide the final list.

Measurement Contract

For every candidate, capture:

exact Bridge and Megatron Core revisions;
model, sequence length, micro/global batch sizes, precision, attention backend, and parallelism;
the exact recompute_granularity, module list, method, and layer count;
per-rank max_memory_allocated() and max_memory_reserved();
steady-state step time or throughput after warmup;
a short convergence or numerical-sanity check appropriate to the task.

Use a matched no-recompute control and change one recompute choice at a time. Peak memory from different jobs, backends, or parallel layouts is not a module-ranking benchmark.

Do not call a candidate successful merely because it advances farther than the control. Run through optimizer-state initialization and multiple steady-state steps: selective recompute can move the memory wall from forward into gradient synchronization or the optimizer without making the workload viable.

Matched H100 Evidence: Moonlight 16B

A 2026-08-12 short-run study used the exact Bridge revision 600d069b824dd5ce50367a311a5a3244478faf22 and Megatron Core revision 24bad8e677d22625d86ef2a54c9506b6e4992c93. The Moonlight 16B BF16 pretraining recipe ran on 8 H100 80GB GPUs with sequence length 4096, MBS=1, GBS=4, TP=2, PP=1, CP=1, EP=8, mock data, and 20 steps. This model mixes one dense layer with 26 MLA+MoE layers. Each row changed only recompute_modules; all 20 losses were finite with zero skipped or NaN iterations.

Peak allocated memory is the maximum post-optimizer value reported after iteration 2. Time and throughput are means over iterations 11--20.

Selective modules	Peak allocated (GB)	Step time (ms)	TFLOP/s/GPU	Allocated vs `[]`	Time vs `[]`
`[]`	36.618	457.18	77.50	control	control
`core_attn`	36.614	474.80	74.44	-0.01%	+3.85%
`mla_up_proj`	35.902	480.89	73.72	-1.96%	+5.19%
`mla_up_proj`, `mlp`	35.917	496.73	71.50	-1.91%	+8.65%
`moe_act`	35.941	466.26	75.50	-1.85%	+1.99%
`layernorm`, `moe_act`	35.949	506.53	70.27	-1.83%	+10.79%

For this exact workload, moe_act is the best first boundary: it recovered nearly as much allocated memory as mla_up_proj for less replay cost. mla_up_proj is the next candidate if its roughly 39 MB additional reduction matters. Adding mlp to mla_up_proj or layernorm to moe_act did not improve the observed peak and made steps slower. Explicit core_attn added cost without material memory benefit under fused attention.

Maximum reserved memory stayed near 40 GB and did not fall monotonically. That is allocator caching, not contrary evidence: boundary selection in this study is based on allocated memory and successful end-to-end steps.

Matched H100 Evidence: Nemotron 3 Nano

The same 2026-08-12 study used the native 16-H100 BF16 performance recipe for the 52-layer hybrid Mamba/fused-attention MoE model. The matched short-run configuration used sequence length 8192, MBS=1, GBS=16, TP=1, PP=1, CP=1, EP=8, DP=16, expert-DP=2, HybridEP, grouped GEMM, TE CUDA graphs for attention and Mamba, mock data, and 12 steps. Each row changed only recompute_modules.

Selective modules	Outcome	Rank-0 measured peak	Failure or steady-state evidence
`[]`	OOM after iteration 1	66.297 GB after iteration 1	Iteration-2 MoE router allocation failed; hot ranks had about 72.9 GiB allocated.
`core_attn`	OOM in iteration 1	not comparable	Grouped-expert linear allocation failed; explicit attention recompute did not make the fused-attention workload fit.
`moe_act`	OOM after iteration 1	62.103 GB after iteration 1	4.194 GB (6.33%) below the control at the matched checkpoint, but the iteration-2 output projection still needed 2 GiB.
`layernorm`, `moe_act`	OOM in iteration 1	not comparable	Output projection still needed 2 GiB; CUDA-graph private pools were material.
`moe`	completed 12 steps	64.653 GB after iteration 2	657.42 ms and 277.72 TFLOP/s/GPU over iterations 7--12.
`moe`, `layernorm`	completed 12 steps	63.639 GB after iteration 2	677.62 ms and 270.62 TFLOP/s/GPU over iterations 7--12.

Both successful rows had finite losses and zero skipped or NaN iterations. For this exact capacity-limited recipe, whole-moe recompute is the smallest tested passing boundary. Adding layernorm recovered another 1.014 GB (1.57%) of rank-0 peak at 3.07% higher step time, so the recipe's broader combination is justified when that headroom is required. Narrow moe_act produced real activation relief but did not make the whole training step viable.

An exploratory native 8-H100 layout failed during FP32 optimizer-state initialization even at sequence length 4096. That is optimizer capacity, not a selective-boundary throughput baseline; no timing comparison from those runs is used here.

Cross-model conclusion

These measurements do not define one ranking. Moonlight fit with an empty control and favored narrow moe_act; Nemotron required broad whole-moe recompute; historical dense Llama evidence found whole-mlp replay costly and lacked an empty control. The correct first candidate is therefore the narrowest boundary implicated by the architecture and peak, followed by broader replay only when the narrow choice does not pass the complete step.

Compatibility and Validation

Configuration semantics

recompute_granularity="selective" uses recompute_modules; an empty list is accepted as an explicit control.
recompute_granularity="full" uses recompute_method and recompute_num_layers; selective labels do not apply.
Full granularity supersedes selective module choices rather than composing with them.
Unknown labels fail Megatron Core validation. Labels may differ on development branches, so use the exact revision's TransformerConfig validator as the source of truth.

Attention backend and context parallelism

TE fused and Flash Attention already use internal rematerialization. Explicit core_attn may still change retained inputs/outputs, but it must earn its place in a matched [] comparison.
Under context parallelism, an attention checkpoint can replay communication as well as compute. Include CP size and topology in the measurement record.

MoE restrictions

Whole-moe recompute is incompatible with expert-parallel overlap because backward replay would repeat the overlapped routing/communication region.
shared_experts recompute is incompatible with shared-expert overlap.
moe_act applies to grouped-GEMM experts and is the narrower choice when only the expert activation needs to be discarded.
mlp targets dense MLPs and is a no-op on MoE layers; mixed dense/MoE models can still benefit on their dense layers.

FP8 restrictions

moe_act and layernorm recompute are not supported with FP8 delayed scaling and require a compatible Transformer Engine version.
Absorbed MLA paths have additional FP8/FP4 restrictions. Validate the exact model/provider path before selecting mla_up_proj.

CUDA graphs

Selective recompute is valid only when a checkpointed module lies wholly inside or wholly outside the selected graph scope. A checkpoint boundary that straddles a graph boundary is invalid.
Capture/warmup can bypass checkpoint wrappers, so verify the final graph scope and replay path rather than assuming eager behavior carries over.
Full recompute with CUDA graphs requires cuda_graph_impl="full_iteration" in the pinned Megatron Core. Otherwise disable CUDA graphs; scoped/local graph capture is not a substitute for full-iteration capture here.

Historical Measurement: Context, Not a Module Ranking

Historical H100 measurements from Bridge PR #3107 used Llama 3 70B SFT on 32 H100 80GB GPUs with FP8 current scaling, sequence length 4096, micro-batch size 1, global batch size 32, TP=4, PP=4, VPP=5, and DP=2:

Configuration	TFLOP/s/GPU	Peak memory
`core_attn` baseline in that run	~704	58.8 GB (OOM on rank 0)
`mlp`	593.6	55.6 GB
`mlp` + `core_attn`	586.8	55.6 GB
`core_attn` + `layernorm`	~702	59.6 GB (OOM on rank 0)
Golden throughput recorded in the PR context	709.93	Not a paired memory measurement

Limitations of this evidence:

it did not include a matched no-recompute row;
the golden row was not a paired module-only comparison;
the measurements cover one dense Llama workload, not MLA or MoE;
the table supports the local memory/throughput tradeoff only and must not be used to rank all recompute labels.

Code Anchors

Selective-label validation and cross-feature checks: 3rdparty/Megatron-LM/megatron/core/transformer/transformer_config.py
Checkpoint implementations: 3rdparty/Megatron-LM/megatron/core/tensor_parallel/random.py
Standard-attention checkpoint boundary: 3rdparty/Megatron-LM/megatron/core/transformer/attention.py
MLA up-projection boundary: 3rdparty/Megatron-LM/megatron/core/transformer/multi_latent_attention.py
Layernorm, dense-MLP, and outer-MoE placement: 3rdparty/Megatron-LM/megatron/core/transformer/transformer_layer.py
Grouped expert activation boundary: 3rdparty/Megatron-LM/megatron/core/transformer/moe/experts.py
Shared-expert and whole-MoE paths: 3rdparty/Megatron-LM/megatron/core/transformer/moe/moe_layer.py
GDN normalization boundary: 3rdparty/Megatron-LM/megatron/core/ssm/gated_delta_net/gdn.py

Failure Diagnosis

Symptom	Likely cause	Next action
`core_attn` gives little or no peak reduction	Fused/Flash attention already rematerializes the expensive internals, or the peak is elsewhere	Compare with `[]`, attribute the peak, then test the architecture-specific boundary such as `mla_up_proj` or `moe_act`.
MLA still OOMs after `core_attn`	Expanded Q/K/V projection tensors, not attention-core tensors, dominate	Test `mla_up_proj`; add `core_attn` only if matched evidence supports it.
MoE peak remains high	Expert intermediate or norm outputs dominate	Test `moe_act`, then `layernorm`; reserve whole `moe` for broader pressure.
Expert-overlap validation fails	Whole-`moe` or `shared_experts` recompute conflicts with overlap	Keep overlap and use a compatible inner boundary, or disable overlap and remeasure the entire configuration.
A selected label has no measurable effect	That module is absent or inactive on the measured layers, or graph capture bypassed the wrapper	Inspect the provider/layer mix and final graph scope; for example, `mlp` is ineffective on pure-MoE layers.
Full recompute plus CUDA graphs asserts	Graph implementation is not full-iteration	Set `cuda_graph_impl="full_iteration"` or disable CUDA graphs.
Reserved memory is high but allocated memory is stable	Allocator fragmentation or caching	Try `PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True` before adding recompute.
OOM moves to a different rank after enabling recompute	Pipeline/virtual-pipeline layer distribution changed the bottleneck	Compare per-rank peaks and tune full block/uniform placement or selective boundaries for the actual hot stage.
A candidate gets farther but still OOMs	Recompute moved the peak into gradient synchronization or optimizer-state initialization	Record the changed failure stage as diagnostic evidence, but require optimizer initialization and multiple steady steps before calling it a pass.

Known Limitations

A module list is not portable across model families, attention backends, parallel layouts, or Megatron Core revisions.
Memory savings are nonlinear when boundaries overlap or nest; additive arithmetic is unreliable.
Full recompute changes RNG execution paths; dropout workloads need a numerical/convergence check.
Activation recompute does not address parameter, optimizer-state, or allocator-fragmentation pressure.
The correct result is the smallest measured replay cost that satisfies the per-rank memory target, not the longest module list.

Files

ATTRIBUTION.mdAdded by AutoGineer875 B
BENCHMARK.md4.8 KiB
LICENSELicence file added by AutoGineer11.2 KiB
SKILL.md18.5 KiB
card.yaml16.5 KiB
evals/evals.json5.3 KiB
skill-card.md4.2 KiB
skill.oms.sig4.7 KiB

card.yaml

title: activation_recompute
validated_on: "2026-08-12"
summary: >
  Selective activation recompute trades replayed forward work for lower retained
  activation memory. The correct checkpoint boundary is architecture- and
  backend-specific: core_attn is common for standard attention, mla_up_proj is
  often more relevant for MLA, and moe_act or layernorm plus moe_act are common
  narrow choices for grouped MoE. The pinned Megatron Core also accepts mlp,
  moe, shared_experts, and gdn_norm_out. Full-layer recompute has the broadest
  memory effect and replay cost. Development branches may expose additional
  model-specific labels, so the exact target revision is authoritative.
validation_status:
  selective_recompute_config:
    - code_verified  # transformer_config.py and module-local checkpoint sites
  full_layer_recompute_config:
    - code_verified  # transformer_config.py and transformer_block.py
  recompute_modules_enum:
    - code_verified  # transformer_config.py
  llama3_70b_sft_fp8_cs_experiment:
    - measured  # historical PR #3107; contextual, not a cross-module ranking
  moonlight_16b_bf16_experiment:
    - measured  # matched 8-H100 module sweep with an empty-list control
  nemotron_3_nano_bf16_experiment:
    - measured  # matched 16-H100 capacity sweep through complete training steps
training_dimensions:
  speed:
    effect: workload-dependent replay overhead
    confidence: medium
    rationale: >
      Narrow output-discard boundaries such as moe_act can replay much less work
      than whole mlp or moe boundaries. Attention cost also depends on the
      backend and context-parallel communication. The historical Llama 3 70B
      run measured mlp only relative to a core_attn baseline, not to a matched
      no-recompute control. The matched Moonlight run found 1.99% overhead for
      moe_act, 5.19% for mla_up_proj, and no useful core_attn memory reduction
      in that exact fused-attention recipe. In the matched Nemotron run,
      whole-moe recompute was the smallest tested passing boundary; adding
      layernorm cost 3.07% more step time for 1.014 GB more rank-0 headroom.
  memory:
    effect: workload-dependent reduction in retained activations
    confidence: medium
    rationale: >
      Savings depend on which tensors drive the per-rank peak. Standard
      attention, MLA expansions, dense FFNs, grouped expert activations, and
      norm outputs require different boundaries; nested boundaries are not
      additive. The matched Moonlight sweep directly confirmed comparable
      savings from moe_act and mla_up_proj, no material core_attn reduction,
      and no incremental peak benefit from the tested nested additions. The
      Nemotron sweep showed that a narrow boundary can materially reduce an
      intermediate peak yet remain insufficient for the complete step.
  scale:
    effect: can replay communication at some boundaries
    confidence: high
    rationale: >
      core_attn under context parallelism, whole moe, and gdn_norm_out can
      replay communication during backward. Scale and topology therefore form
      part of the measurement contract.
  convergence:
    effect: no semantic change expected when RNG replay is correct; verify
    confidence: medium
    rationale: >
      Recomputation should reproduce the forward region during backward, but
      dropout, graph capture, precision, and branch-specific behavior warrant a
      numerical or short convergence check.
  stability:
    effect: conditional on feature compatibility
    confidence: high
    rationale: >
      Megatron Core validates restrictions involving CUDA graphs, FP8 modes,
      overlap features, and the selected module labels.
enable_when:
  - allocated activation memory, rather than allocator fragmentation, drives OOM or inadequate headroom
  - a matched architecture-specific checkpoint boundary reduces the hot rank's peak at acceptable throughput
  - full-layer recompute is needed after narrower selective boundaries are exhausted
avoid_when:
  - the model fits with acceptable headroom and recompute does not improve another measured objective
  - max reserved memory is high but max allocated memory does not support an activation-capacity diagnosis
  - the selected module is absent or inactive for the model's measured layers
  - a required CUDA-graph, FP8, or overlap feature is incompatible with the boundary
interactions:
  required: []
  conditional:
    - selective granularity accepts an empty recompute_modules list as a matched no-recompute control
    - full granularity requires recompute_method and recompute_num_layers
    - full recompute with CUDA graphs requires cuda_graph_impl=full_iteration
    - selective graph capture requires each checkpoint boundary to lie wholly inside or wholly outside the graph scope
    - core_attn can replay context-parallel communication and may add little with fused or Flash Attention
    - moe_act and layernorm have FP8 delayed-scaling restrictions
    - whole moe recompute is incompatible with expert-parallel overlap
    - shared_experts recompute is incompatible with shared-expert overlap
  incompatible:
    - full recompute with scoped or local non-full-iteration CUDA graphs
    - whole moe recompute with expert-parallel overlap
    - shared_experts recompute with shared-expert overlap
feature_meaning:
  recompute_granularity: >
    Controls scope: null (no recompute), selective (per-module), or full
    (complete transformer-layer regions).
  recompute_modules: >
    Pinned-revision labels are core_attn, mla_up_proj, layernorm, moe_act, mlp,
    moe, shared_experts, and gdn_norm_out. Development revisions can add labels;
    for example, DeepSeek V4's mhc label requires its development MCore branch.
  recompute_method: >
    For full granularity: uniform groups fixed-size layer blocks; block selects
    a fixed number of layers per pipeline stage with virtual-pipeline-aware
    distribution.
  recompute_num_layers: >
    For full granularity: the layer group size or number of selected layers,
    depending on recompute_method.
config_keys:
  - model.recompute_granularity
  - model.recompute_modules
  - model.recompute_method
  - model.recompute_num_layers
  - model.distribute_saved_activations
recommended_path:
  first_step: compare per-rank allocated and reserved peaks and keep a matched no-recompute control
  standard_attention: test core_attn, but compare it with an empty list under fused or Flash Attention
  mla: test mla_up_proj first when expanded Q/K/V tensors dominate; add core_attn only with evidence
  grouped_moe: test moe_act and then layernorm when their outputs drive the peak
  dense_ffn: test mlp only when its broader replay cost is acceptable
  broader_moe: reserve whole moe for cases where narrower MoE boundaries do not fit
  last_resort: use full recompute with recompute_method and recompute_num_layers
  alternative: see skills/nemo-mbridge-perf-memory-tuning for allocator, parallelism, and offloading options
expected_metric_change:
  - metric: peak_memory
    direction: down
    magnitude: must be measured for the exact boundary and hot rank
    conditions: matched model, backend, parallelism, precision, batch, and sequence configuration
    evidence: module checkpoint implementations, matched Moonlight and Nemotron studies, and historical measured_pr_3107
  - metric: steady_state_step_time
    direction: up
    magnitude: depends on replayed compute and communication
    conditions: same matched configuration and warmup
    evidence: module checkpoint implementations, matched Moonlight and Nemotron studies, and historical measured_pr_3107
measured_results:
  - model: Moonlight 16B
    task: pretrain_short_run
    revisions:
      bridge: 600d069b824dd5ce50367a311a5a3244478faf22
      megatron_core: 24bad8e677d22625d86ef2a54c9506b6e4992c93
    parallelism: TP2_PP1_CP1_EP8_DP4
    gpus: 8
    gpu: H100_80GB
    precision: BF16
    seq_length: 4096
    mbs: 1
    gbs: 4
    steps: 20
    measurement: peak allocated after iteration 2; mean time and throughput over iterations 11-20
    numerical_sanity: all losses finite; zero skipped and NaN iterations
    experiments:
      - name: empty_control
        recompute_modules: []
        peak_mem_gb: 36.618
        step_time_ms: 457.18
        tflops_per_gpu: 77.50
      - name: core_attn
        recompute_modules: ["core_attn"]
        peak_mem_gb: 36.614
        peak_mem_vs_control_pct: -0.01
        step_time_ms: 474.80
        step_time_vs_control_pct: 3.85
      - name: mla_up_proj
        recompute_modules: ["mla_up_proj"]
        peak_mem_gb: 35.902
        peak_mem_vs_control_pct: -1.96
        step_time_ms: 480.89
        step_time_vs_control_pct: 5.19
      - name: mla_up_proj_plus_mlp
        recompute_modules: ["mla_up_proj", "mlp"]
        peak_mem_gb: 35.917
        peak_mem_vs_control_pct: -1.91
        step_time_ms: 496.73
        step_time_vs_control_pct: 8.65
      - name: moe_act
        recompute_modules: ["moe_act"]
        peak_mem_gb: 35.941
        peak_mem_vs_control_pct: -1.85
        step_time_ms: 466.26
        step_time_vs_control_pct: 1.99
      - name: layernorm_plus_moe_act
        recompute_modules: ["layernorm", "moe_act"]
        peak_mem_gb: 35.949
        peak_mem_vs_control_pct: -1.83
        step_time_ms: 506.53
        step_time_vs_control_pct: 10.79
    conclusion: >
      For this exact mixed dense/MLA/MoE recipe, moe_act is the efficient first
      boundary and mla_up_proj is the next candidate when its roughly 39 MB
      additional saving matters. Adding mlp or layernorm did not reduce the
      observed peak, and core_attn added cost without material memory benefit.
  - model: Nemotron 3 Nano
    task: pretrain_capacity_short_run
    revisions:
      bridge: 600d069b824dd5ce50367a311a5a3244478faf22
      megatron_core: 24bad8e677d22625d86ef2a54c9506b6e4992c93
    architecture: 52-layer hybrid Mamba/fused-attention MoE
    parallelism: TP1_PP1_CP1_EP8_DP16
    expert_data_parallel_size: 2
    gpus: 16
    gpu: H100_80GB
    precision: BF16
    seq_length: 8192
    mbs: 1
    gbs: 16
    steps: 12
    execution: HybridEP, grouped GEMM, TE CUDA graphs for attention and Mamba
    measurement: rank-0 reported peak; mean time and throughput over iterations 7-12 for passing rows
    experiments:
      - name: empty_control
        recompute_modules: []
        peak_mem_gb_after_iteration_1: 66.297
        status: OOM_iteration_2_moe_router
      - name: core_attn
        recompute_modules: ["core_attn"]
        status: OOM_iteration_1_grouped_expert_linear
      - name: moe_act
        recompute_modules: ["moe_act"]
        peak_mem_gb_after_iteration_1: 62.103
        peak_mem_vs_control_pct_at_iteration_1: -6.33
        status: OOM_iteration_2_output_projection
      - name: layernorm_plus_moe_act
        recompute_modules: ["layernorm", "moe_act"]
        status: OOM_iteration_1_output_projection
      - name: moe
        recompute_modules: ["moe"]
        peak_mem_gb: 64.653
        step_time_ms: 657.42
        tflops_per_gpu: 277.72
        status: completed_12_steps
      - name: moe_plus_layernorm
        recompute_modules: ["moe", "layernorm"]
        peak_mem_gb: 63.639
        peak_mem_vs_moe_pct: -1.57
        step_time_ms: 677.62
        step_time_vs_moe_pct: 3.07
        tflops_per_gpu: 270.62
        status: completed_12_steps
    numerical_sanity: both passing rows had finite losses and zero skipped or NaN iterations
    comparison_limitations: >
      Failed rows do not provide steady-state timing, and peaks from different
      failure phases are not directly rankable. A native 8-H100 exploration
      hit FP32 optimizer-state capacity and is excluded from throughput claims.
    conclusion: >
      Whole-moe recompute was the smallest tested boundary that completed the
      exact capacity-limited workload. Adding layernorm bought 1.014 GB more
      rank-0 headroom at 3.07% higher step time. Moe-act provided real relief
      but did not make the complete step viable; core-attn did not make the
      workload fit and yielded no positive memory evidence.
  - model: Llama3 70B
    task: sft
    parallelism: TP4_PP4_VPP5_DP2
    gpus: 32
    gpu: H100_80GB
    precision: FP8_CS
    seq_length: 4096
    mbs: 1
    gbs: 32
    golden_tflops: 709.93
    comparison_limitations: >
      No matched no-recompute row was recorded; golden throughput was not a
      paired module-only comparison; results do not cover MLA or MoE.
    experiments:
      - name: core_attn_baseline
        recompute_modules: ["core_attn"]
        tflops: 704
        vs_golden_pct: -0.8
        peak_mem_gb: 58.8
        status: OOM_on_rank0
      - name: mlp_only
        recompute_modules: ["mlp"]
        tflops: 593.6
        vs_golden_pct: -16.4
        peak_mem_gb: 55.6
        status: perf_regression
      - name: mlp_plus_core_attn
        recompute_modules: ["mlp", "core_attn"]
        tflops: 586.8
        vs_golden_pct: -17.3
        peak_mem_gb: 55.6
        status: perf_regression
      - name: core_attn_plus_layernorm
        recompute_modules: ["core_attn", "layernorm"]
        tflops: 702
        vs_golden_pct: -1.1
        peak_mem_gb: 59.6
        status: OOM_on_rank0
failure_modes:
  - name: core_attn_has_little_incremental_effect
    symptom: no material peak reduction against the empty-list control
    likely_cause: fused attention already rematerializes the expensive internals or another architecture region sets the peak
    fix: attribute the peak and test mla_up_proj, moe_act, layernorm, or another applicable boundary
  - name: module_is_inactive
    symptom: selected recompute label causes no measurable change
    likely_cause: the module is absent on the measured layers, such as mlp on pure-MoE layers, or graph capture bypasses the wrapper
    fix: inspect the provider and final graph scope before choosing another boundary
  - name: overlap_validation_failure
    symptom: Megatron Core rejects whole-moe or shared-expert recompute
    likely_cause: expert-parallel or shared-expert overlap conflicts with backward replay
    fix: retain overlap and use a compatible inner boundary, or disable overlap and remeasure the complete configuration
  - name: later_failure_is_misclassified_as_success
    symptom: recompute advances beyond the control's failure point but still OOMs
    likely_cause: the boundary moved the peak into a later forward region, gradient synchronization, or optimizer initialization
    fix: record the changed failure stage, but require optimizer initialization and multiple steady-state steps before calling it a pass
  - name: full_recompute_with_scoped_cuda_graphs
    symptom: "AssertionError: full recompute is only supported with full iteration CUDA graph"
    likely_cause: full granularity is combined with a non-full-iteration graph implementation
    fix: use selective recompute, set cuda_graph_impl=none, or use cuda_graph_impl=full_iteration
known_constraints:
  - recompute_granularity=selective accepts an empty recompute_modules list
  - recompute_granularity=full requires recompute_method and recompute_num_layers
  - distribute_saved_activations cannot be used with sequence_parallel=True
  - valid module labels are revision-specific; use the exact target Megatron Core validator
  - module boundaries can overlap or nest, so memory and replay costs are not additive
known_limitations:
  - no automatic selection of the optimal recompute_modules list
  - no single module ordering applies across standard attention, MLA, MoE, and GDN models
  - matched Moonlight evidence covers one short BF16 MLA+MoE configuration and is not a universal ranking
  - matched Nemotron evidence covers one short BF16 hybrid MoE capacity configuration and is not a universal ranking
  - historical Llama evidence lacks a matched empty-list control
evidence:
  - docs/training/activation-recomputation.md
  - "2026-08-12 matched H100 Moonlight 16B activation-recompute study"
  - "2026-08-12 matched H100 Nemotron 3 Nano activation-recompute study"
  - "PR #3107 (historical Llama3 70B SFT experiment)"
  - 3rdparty/Megatron-LM/megatron/core/transformer/transformer_config.py
  - 3rdparty/Megatron-LM/megatron/core/transformer/attention.py
  - 3rdparty/Megatron-LM/megatron/core/transformer/multi_latent_attention.py
  - 3rdparty/Megatron-LM/megatron/core/transformer/transformer_layer.py
  - 3rdparty/Megatron-LM/megatron/core/transformer/moe/experts.py
  - 3rdparty/Megatron-LM/megatron/core/transformer/moe/moe_layer.py
  - 3rdparty/Megatron-LM/megatron/core/ssm/gated_delta_net/gdn.py
follow_up_validation:
  - Add matched GDN measurements.
  - Measure context-parallel communication replay for attention and GDN boundaries.
  - Validate representative selective boundaries under the supported CUDA-graph scopes and FP8 recipes.

Requirements and compatibility

Works with: Claude Code, Claude apps, OpenAI Codex, GitHub Copilot, Cursor, Gemini CLI, Kiro, OpenCode, Goose, Amp, Roo Code, JetBrains Junie, Factory Droid, Devin (ex-Windsurf), Mistral Vibe, OpenHands, Letta

“Works with” means the tool reads the open Agent Skills format; a skill written for one tool may use features others lack.

Needs: Needs Python

Security

What our scanner noticed· Medium risk

Large base64-encoded blobs

A static scan is not a guarantee. Skills are third-party content: review a skill before you let an agent use it.

Attribution and licence

Published by

NVIDIA

Publisher

NVIDIA (Official)

Source

NVIDIA/skills · skills/nemo-mbridge-perf-activation-recompute · commit fce3cca

Licence

Apache License 2.0 (redistribution allowed)

LICENSE (added by AutoGineer from LICENSE-APACHE)
ATTRIBUTION.md (written by AutoGineer)

Synced

2026-11-09

AutoGineer is not affiliated with or endorsed by NVIDIA or NVIDIA.

Download ZIP

The ZIP holds the skill’s files, its licence files and an ATTRIBUTION.md. Download nemo-mbridge-perf-activation-recompute.zip