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

gradient_checkpointing: true

Enabling Activation Offloading

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)

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:

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:

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 (via save, save_modules or save_matmul_min_k) is rejected while a fused LoRA kernel is enabled, because the kernel adds the adapter delta into the base matmul output in place, which invalidates the cached tensor.

Saving specific projections

All projections dispatch the same aten::mm/aten::addmm, so save: [aten::mm] saves every one of them. Two rules target individual projections instead; both also work with offload: true.

save_modules saves the matmuls that run inside the named modules:

selective_checkpointing:
  save: [attention]
  save_modules: [o_proj]

An entry matches a module name exactly, as a trailing dotted suffix (o_proj, self_attn.o_proj), or as a glob over the full name when it contains *, ? or [. Names are those of the base model, before PEFT wraps it, and a PEFT-wrapped match resolves to its base_layer, so only the frozen base matmul is saved, not lora_A/lora_B. Modules that enclose the decoder layers, or run outside them (lm_head, embed_tokens), are never in scope.

save_matmul_min_k saves every matmul whose contraction dim K (a linear’s in_features) is at least the threshold, including MoE expert (aten::_grouped_mm) and bitsandbytes 4-bit gemms:

selective_checkpointing:
  save: [attention]
  save_matmul_min_k: 8192

Saving only helps for ops that recompute would otherwise re-run. Recompute stops at the last tensor backward needs, so a layer’s final projection (down_proj in Llama/Qwen3-style layers, whose output only feeds the residual add) is never replayed; saving it costs memory for no speedup. Prefer o_proj, gate_proj or up_proj.

Validation rejects entries covered by an enabled fused LoRA kernel (lora_mlp_kernel, lora_qkv_kernel, lora_o_kernel, auto-enabled for LoRA/QLoRA), and the error names the flag to turn off. At runtime a warning is logged for any rule that matched no module, whose hooks never fired, whose saved tensors were never replayed, or whose threshold no matmul reached.

save: [] saves nothing through the list, so save: [] with save_modules saves only the named projections and recomputes attention.

Plugin-registered saves

Plugins can register ops that every selective-checkpoint policy saves, ahead of save and the rules, through axolotl.monkeypatch.selective_checkpointing. register_mandatory_save is for correctness: ops whose recompute would diverge from the forward, such as a router’s topk whose near-ties would re-route tokens, or a device-to-host copy that sizes a collective and would re-sync. Their outputs are cloned for the forward so an in-place edit (a router normalising its topk weights) cannot corrupt the cached copy. register_preferred_save only trades memory for skipping an expensive recompute. FSDP2 activation_checkpointing applies registered saves on its own; under plain gradient_checkpointing a plugin installs them with apply_selective_checkpointing(model, save=[]).

Enabling Layer Offloading

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