/docs/gradient_checkpointing.qmd

---
title: Gradient Checkpointing, Activation Offloading, and Layer Offloading
---

Gradient checkpointing and activation offloading are techniques used to optimize the performance of deep learning
models by reducing the memory footprint and improving computational efficiency.

### Enabling Gradient Checkpointing

```yaml
gradient_checkpointing: true
```

### Enabling Activation Offloading

```yaml
gradient_checkpointing: true  # required for activation offloading
activation_offloading: true
```

Activation offloading variants:

The default `activation_offloading: true` offloads activations to CPU and uses CUDA streams
to overlap the communications and computations when offloading.

The `activation_offloading: legacy` naively offloads activations to CPU and without additional optimizations.

For resource constrained environments with limited CPU memory, `activation_offloading: disk` offloads
activations to disk instead of CPU RAM so that much larger context lengths can be trained with minimal memory.

The `activation_offloading: hidden_states` mode keeps gradient checkpointing enabled and moves only
the decoder-layer checkpoint input (`hidden_states`) to CPU. With the default
`gradient_checkpointing_kwargs.use_reentrant: false`, Axolotl uses Transformers' non-reentrant
checkpointing plus a saved-tensor hook that only offloads marked hidden-state tensors. Other saved
tensors stay on GPU, which avoids the CPU memory and bandwidth cost of full activation offload.

Set `gradient_checkpointing_kwargs.use_reentrant: true` to use the older ALST-style implementation.
That path patches torch's reentrant `CheckpointFunction` and is DTensor-aware for sequence/context
parallelism.

### Choosing a mode

| Mode | What moves | Best for |
|------|-----------|----------|
| `true` | All activations → CPU, stream-overlapped | General use; LoRA/QLoRA (offloads instead of recomputing) |
| `legacy` | All activations → CPU, synchronous | Lowest resident GPU memory; when the streamed path's in-flight buffering inflates peak |
| `disk` | Activations → disk | Severely CPU-RAM-constrained hosts |
| `hidden_states` | Per-layer checkpoint input only → CPU, stream-overlapped | Long-context training where full activation offload is CPU-memory or bandwidth-bound |

`hidden_states` uses gradient checkpointing and recomputes each layer during backward. With LoRA/QLoRA,
`activation_offloading: true` can still be faster because adapter activations are smaller and pure offload
avoids the recompute cost.

### Selective Activation Checkpointing (SAC)

```yaml
gradient_checkpointing: true
selective_checkpointing: true
```

Plain gradient checkpointing recomputes every op in a decoder layer during backward. Selective
checkpointing keeps the recompute for cheap ops but *saves* the outputs of expensive ones — by
default the attention op — so backward skips re-running them. At long context, attention dominates
the recompute cost, so this recovers most of the checkpointing slowdown for one extra
hidden-state-sized tensor per layer.

This runs eagerly (no `torch.compile`) via torch's selective-checkpoint policy API. Ops are matched
at the dispatcher level: SDPA and flash-attention are matched out of the box, and custom kernels
that are not registered as torch ops are simply recomputed as usual — nothing breaks.

The explicit form lets you save additional ops, matched as substrings of the qualified op name:

```yaml
selective_checkpointing:
  save:
    - attention
    - aten::mm        # example: also save all matmuls
```

For hybrid models mixing full attention and sliding-window attention (SWA), only full-attention
calls are saved by default — SWA is cheap to recompute and not worth the memory. Hybrid layers are
detected two ways: the flash-attention window arguments, and the model's `layer_types` (published
per-layer via hooks, which also covers SDPA where the window lives inside the attention mask).
Set `selective_checkpointing.save_sliding_window: true` to save SWA calls too, or override which
layer types recompute with `selective_checkpointing.recompute_layer_types` (default
`[sliding_attention, chunked_attention]`; linear-attention layers never dispatch a matchable
attention op, so they need no entry).

To offload the saved tensors to pinned CPU memory instead of keeping them on GPU — asynchronous
side-stream copies during forward, prefetched back one region ahead during backward — set:

```yaml
selective_checkpointing:
  save: [attention]
  offload: true
```

This makes the saved attention outputs cost ~zero GPU memory at the price of CPU RAM and PCIe
traffic, and stacks with `activation_offloading: hidden_states`.

Requires non-reentrant checkpointing (the default). Composes with
`activation_offloading: hidden_states`: layer inputs go to CPU while attention outputs are saved
(on GPU, or on CPU with `offload: true`). It is incompatible with the TRL-offloader modes
(`true`/`legacy`/`disk`), which replace gradient checkpointing instead of running inside it.
Saving matmul ops (e.g. `aten::mm`) is rejected for LoRA/QLoRA runs: PEFT mutates the base linear
output in-place, which invalidates cached tensors.

### Enabling Layer Offloading

```yaml
layer_offloading: true
```

Layer offloading reduces GPU memory usage by moving frozen (non-trainable) decoder layer parameters to CPU
and streaming them back to GPU one layer at a time during the forward and backward passes. This is
particularly useful for LoRA/QLoRA training where most of the model's parameters are frozen — only the
trainable adapter weights stay on GPU permanently.

During training, forward and backward hooks on each decoder layer handle the transfer automatically:

- **Forward pass:** Before a layer executes, its frozen params are loaded to GPU. The next layer is
  prefetched asynchronously on a separate CUDA stream for overlap.
- **Backward pass:** Same pattern in reverse — the current layer's frozen params are loaded and the
  previous layer is prefetched.

After each layer finishes, its frozen params are offloaded back to CPU pinned memory.

This approach trades some CPU-GPU transfer overhead for significant GPU memory savings — the freed memory
is roughly equal to the size of all frozen parameters across all decoder layers, minus one layer's worth
that is kept on GPU at any given time.

**Requirements:**

- CUDA GPU (CPU-only training is not supported for this feature)
- Works with any HuggingFace model architecture that uses decoder layers (Llama, Mistral, Qwen, etc.)
- Best combined with LoRA/QLoRA where most parameters are frozen