Gradient (Activation) Checkpointing in NeMo AutoModel

Gradient checkpointing, also called activation checkpointing, trades a little extra compute for a large reduction in GPU memory by recomputing intermediate activations during the backwards 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
# examples/llm_finetune/llama_3_2_1b_my_finetune.yaml
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
  ...

Configure Programmatically

1
from nemo_automodel.components.distributed.config import FSDP2Config
2
from nemo_automodel.components.distributed.fsdp2 import FSDP2Manager
3
4
config = FSDP2Config(activation_checkpointing=True)
5
# device_mesh is created elsewhere (e.g. by the recipe via setup_distributed)
6
manager = FSDP2Manager(config, device_mesh=device_mesh, moe_mesh=moe_mesh)
7
model = manager.parallelize(model)

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)

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 your training script and inspect the peak memory:

$
# If running on 8x GPUs
$
automodel --nproc-per-node=8 examples/llm_finetune/llama3_2/llama_3_2_1b_my_finetune.yaml
$
$
# If running on 1x GPU
$
automodel examples/llm_finetune/llama3_2/llama_3_2_1b_my_finetune.yaml

If we run with the above settings (activation ckpt = on, lc-ce = on, fsdp = on), look for a log line similar to:

... | mem 7.30 GiB | ...