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 (e.g., 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):

1 # Add this field to the `distributed` section of your recipe config.
2 ...
3 
4 # FSDP2 (use strategy name; optional parallelism sizes)
5 distributed:
6   strategy: fsdp2
7   activation_checkpointing: true
8   # dp_size: null
9   # tp_size: 1
10   # cp_size: 1
11   ...

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

1 distributed:
2   strategy: fsdp2
3   activation_checkpointing: selective

1 distributed:
2   strategy: ddp
3   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:

1 distributed:
2   strategy: fsdp2
3   activation_checkpointing: true
4   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:

1 distributed:
2   strategy: fsdp2
3   activation_checkpointing: true
4   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 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:

1 from nemo_automodel.components.distributed.config import FSDP2Config
2 
3 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 (e.g., Gemma4) automatically fall back to sub-module checkpointing, because attention cannot be recomputed through the KV cache.

Selective 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 (e.g., an expert grouped-GEMM) shows up as RECOMPUTE, selective AC will not beat full checkpointing for that model.

Full vs. 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 vs. selective 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

1 from nemo_automodel import NeMoAutoModelForCausalLM
2 from nemo_automodel.components.distributed.config import DistributedSetup
3 
4 # Use activation_checkpointing="selective" for FSDP2 or DDP selective checkpointing.
5 distributed_setup = DistributedSetup.build(
6     strategy="fsdp2",
7     activation_checkpointing="selective",
8 )
9 
10 model = NeMoAutoModelForCausalLM.from_pretrained(
11     "meta-llama/Llama-3.2-1B",
12     distributed_setup=distributed_setup,
13 )

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, thus avoiding 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:

1 model:
2   ...
3   output_hidden_states: true
4 
5 loss_fn:
6   _target_: nemo_automodel.components.loss.linear_ce.FusedLinearCrossEntropy

LC-CE and gradient checkpointing target different memory hot-spots (output layer vs. 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 vs. Batch Size: The goal is usually to increase batch size or sequence length while keeping throughput constant.

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 | ...