Use Gradient (Activation) Checkpointing

Gradient checkpointing, also called activation checkpointing, trades a little extra compute for a large reduction in GPU memory by recomputing intermediate activations during the backward pass instead of storing them. It is especially powerful when combined with memory-efficient loss functions (for example, Linear-Cut Cross-Entropy) and parameter sharding using FSDP.

Enable Gradient Checkpointing

Configure in YAML

Add the activation_checkpointing: true flag under your distributed strategy.

Example (snippet):

# Add this field to the `distributed` section of your recipe config.
...
# FSDP2 (use strategy name; optional parallelism sizes)
distributed:
  strategy: fsdp2
  activation_checkpointing: true
  # dp_size: null
  # tp_size: 1
  # cp_size: 1
  ...

For FSDP2 and DDP, activation_checkpointing also accepts explicit policy strings:

distributed:
  strategy: fsdp2
  activation_checkpointing: selective

distributed:
  strategy: ddp
  activation_checkpointing: selective

Use true or full for full activation checkpointing. Use selective for PyTorch selective activation checkpointing on FSDP2 or DDP configs. Selective checkpointing saves expensive operations such as attention, collectives, and part of the matrix multiplications while recomputing cheaper operations during backward.

Scope Activation Checkpointing

For multimodal and VLM models, AutoModel can checkpoint only selected model parts. Use activation_checkpointing_scope with FSDP2 or DDP when you want to keep the memory/speed profile focused on part of the model:

distributed:
  strategy: fsdp2
  activation_checkpointing: true
  activation_checkpointing_scope: language

Valid scopes are all, language, vision, audio, and multimodal. all is the default. multimodal means the non-language towers currently identified by AutoModel, such as vision and audio. You can also pass a list, for example:

distributed:
  strategy: fsdp2
  activation_checkpointing: true
  activation_checkpointing_scope: [language, vision]

Use a narrower scope when only part of a multimodal model is memory-limited or when checkpointing trainable vision/audio towers costs more throughput than the memory savings are worth. Use all when you prefer the maximum activation-memory reduction across trainable model parts.

Expert-parallel Mixture-of-Experts (MoE) configs (ep_size > 1) use the MoE parallelizer’s activation checkpointing path, which respects the same activation_checkpointing_scope field with the same semantics: all (the default) checkpoints the text/MoE decoder blocks plus any trainable multimodal towers (vision, audio), language the decoder blocks only, and vision, audio, or multimodal only the selected trainable tower blocks. Tower checkpointing behaves the same as on the generic path and works with any attention backend; no additional configuration is needed. Frozen towers stay out of activation checkpointing on both paths.

If a vision, audio, or other tower is fully frozen, AutoModel leaves it out of activation checkpointing. This avoids extra recomputation for model parts that do not need parameter gradients. Layers with trainable adapters, such as LoRA, can still be checkpointed when their model part is selected by the scope because those adapter weights need gradients.

Programmatically, pass the scope on the strategy config:

from nemo_automodel.components.distributed.config import FSDP2Config
strategy = FSDP2Config(activation_checkpointing_scope="language")

selective is supported for FSDP2 and DDP. Megatron-FSDP raises an error when selective is requested. KV-sharing models (for example, Gemma4) automatically fall back to sub-module checkpointing, because attention cannot be recomputed through the KV cache.

Selective activation checkpointing (AC) only speeds things up when the model’s expensive operations are the ones being saved. To see the per-op save/recompute decisions for your model, set NEMO_SELECTIVE_AC_TRACE=1; each unique operation is logged once as SAVE, RECOMPUTE, or ALTERNATE. If an expensive op (for example, an expert grouped-GEMM) shows up as RECOMPUTE, selective AC will not beat full checkpointing for that model.

Full compared to selective: Selective AC saves the expensive operations (attention and part of the matmuls) and recomputes only the cheaper ones, so it does less recompute work than full AC while holding more activations in memory. Whether that nets out as faster, and at what memory cost, depends on the model, sequence length, and whether torch.compile is enabled, so benchmark full and selective checkpointing for your own setup. When you do, keep the torch.compile setting the same on both sides (compare full and selective both compiled, or both uncompiled). torch.compile is a large speed lever on its own and helps both modes, so mixing it in makes it hard to tell which gain came from the AC mode.

MoE/expert parallelism: Selective AC is designed for dense transformers and generally does not help Mixture-of-Experts models with expert parallelism. In an MoE block the experts dominate the cost (they are cheap to recompute but expensive to store), and the expert-parallel dispatch/communication is opaque to the selective policy, so it is recomputed regardless. As a result, selective AC tends to add activation memory without a corresponding speedup for MoE, matching what reference implementations such as TorchTitan observe. Prefer full activation checkpointing (true/full) for MoE; selective remains available as an opt-in for FSDP2, including MoE configs, and for DDP.

Configure Programmatically

from nemo_automodel import NeMoAutoModelForCausalLM
from nemo_automodel.components.distributed.config import DistributedSetup
# Use activation_checkpointing="selective" for FSDP2 or DDP selective checkpointing.
distributed_setup = DistributedSetup.build(
    strategy="fsdp2",
    activation_checkpointing="selective",
)
model = NeMoAutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3.2-1B",
    distributed_setup=distributed_setup,
)

Combine with Linear-Cut Cross-Entropy (LC-CE)

Linear-Cut Cross-Entropy (LC-CE) reduces the hidden-state memory required to compute the loss by calculating the softmax on the fly, which avoids the need to allocate memory for the logits. It is already available using nemo_automodel.components.loss.linear_ce.FusedLinearCrossEntropy and can be enabled in recipes by using the following:

model:
  ...
  output_hidden_states: true
loss_fn:
  _target_: nemo_automodel.components.loss.linear_ce.FusedLinearCrossEntropy

LC-CE and gradient checkpointing target different memory hot-spots (output layer compared to transformer blocks), so their benefits stack almost linearly.

Example Memory Savings (H100-80GB, Llama-3.2-1B)

Technique	Max GPU Mem (GB)	Δ vs Baseline
Baseline	53.03	-
+ FSDP (dp_size=8)	47.59	↓ 10 %
+ Gradient Checkpointing	33.06	↓ 38 %
+ LC-CE	7.30	↓ 86 %
FSDP + LC-CE + Checkpointing	7.30	↓ 86 %

Measurements taken with local batch size = 8, sequence len = 2048, AdamW, PyTorch 2.8.
Peak memory reported by torch.cuda.max_memory_allocated() averaged across DP ranks.
Expect ±5 % variance depending on exact model, sequence length, and GPU architecture.

Performance Considerations

Extra compute: Each checkpointed segment is recomputed once during the backward pass. In practice, the wall-clock overhead is ≈5-10% for transformer models.
Throughput compared to batch size: The goal is usually to increase batch size or sequence length while keeping throughput constant.
MoE recompute cache (HybridEP dispatcher): With block checkpointing, the backward re-executes every MoE layer’s forward, including the HybridEP combine (unpermute, expert-parallel barrier, all-to-all) and the expert down projection whose output only feeds that combine. model.backend.dispatcher_recompute_cache keeps each layer’s combine output from the checkpoint forward (one [tokens, hidden] activation per MoE layer per in-flight microbatch) and serves it in the recompute:

model:
  backend:
    dispatcher: hybridep
    experts: torch_mm
    dispatcher_recompute_cache: moe_tail   # none (default) | combine | moe_tail
distributed:
  activation_checkpointing: true

combine skips the recomputed unpermute, barrier, and combine all-to-all; moe_tail additionally builds the down projection’s autograd node without running its GEMM (torch grouped-GEMM experts without expert bias, router weight applied inside the activation). Gradients are unchanged. On DeepSeek-V4.1-Flash at 16 GB200 nodes (EP64, 4k sequences, local batch 2) the step time dropped by 1.4 % (combine) and 3.9 % (moe_tail) for +3.2 GiB of peak allocated memory (40 MoE layers × 80 MiB).

Verify It Works

Run a checked-in recipe config with activation checkpointing enabled and inspect the peak memory:

# If running on 8x GPUs
uv run automodel examples/llm_finetune/llama3_2/llama3_2_1b_squad.yaml \
  --nproc-per-node 8 \
  --distributed.activation_checkpointing true
# If running on 1x GPU
uv run automodel examples/llm_finetune/llama3_2/llama3_2_1b_squad.yaml \
  --distributed.activation_checkpointing true

Compare the peak memory with the same command run without the override. The exact value depends on the model, batch size, sequence length, loss, and hardware. Look for the memory field in a log line similar to:

... | mem <peak-memory> GiB | ...