> For clean Markdown of any page, append .md to the page URL.
> For a complete documentation index, see https://docs.nvidia.com/nemo/automodel/llms.txt.
> For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/nemo/automodel/_mcp/server.

# 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):

```yaml
# 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:

```yaml
distributed:
  strategy: fsdp2
  activation_checkpointing: selective
```

```yaml
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:

```yaml
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:

```yaml
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 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:

```python
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 (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

```python
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, 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:

```yaml
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 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

1. **Extra compute**: Each checkpointed segment is recomputed once during the backward pass. In practice, the wall-clock overhead is ≈5-10% for transformer models.
2. **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:

```bash

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