Gradient Checkpointing, Activation Offloading, and Layer Offloading

Gradient checkpointing, activation offloading and layer offloading reduce GPU memory during training. Each trades it for something else: extra compute, host memory, or PCIe bandwidth.

How it works

In the forward pass, each layer keeps some of its intermediate tensors (activations), because the backward pass needs them to compute gradients. With long sequences these activations, not the weights, take most of the GPU memory: they grow with batch size, sequence length, hidden size and the number of layers.

• Gradient checkpointing (gradient_checkpointing: true) keeps only each decoder layer’s input during forward and drops the rest. In backward, each layer runs its forward again (recompute) to rebuild what it needs. Memory drops to about one hidden-state tensor per layer, at the cost of roughly one extra forward pass per step.
• Activation offloading (activation_offloading) keeps activations, but in CPU RAM until backward needs them: it spends PCIe bandwidth and host memory instead of compute.
• Selective checkpointing (selective_checkpointing) recomputes each layer except for the ops it saves, attention by default, which is usually the most expensive part to recompute.

How the settings combine (decoder layers; the output head always stays on GPU):

gradient_checkpointing	activation_offloading	Recomputed in backward	Moved to CPU
false	false	nothing	nothing: every activation stays on GPU
true	false	every layer	nothing: each layer’s input stays on GPU
true	true	every layer	each layer’s input, and what is saved outside the layers
false	true	nothing	every saved activation
true	hidden_states	every layer	each layer’s input only

Adding selective_checkpointing to any row with gradient_checkpointing: true skips recomputing the saved ops. Their outputs stay on GPU, or move to CPU with selective_checkpointing.offload: true. For example, to offload each layer’s input and not recompute attention:

gradient_checkpointing: true
activation_offloading: hidden_states
selective_checkpointing:
  save: [attention]
  offload: true  # optional: keep the saved attention outputs in CPU RAM too

Enabling Gradient Checkpointing

gradient_checkpointing: true

Unless gradient_checkpointing_kwargs.use_reentrant is set, Axolotl uses PyTorch’s recommended non-reentrant checkpointing (use_reentrant: false). SFT runs on DeepSpeed ZeRO-3 default to use_reentrant: true, since non-reentrant recompute raises a CheckpointError there; rl: runs on ZeRO-3 need it set explicitly.

Enabling Activation Offloading

gradient_checkpointing: true
activation_offloading: true

With activation_offloading: true, gradient_checkpointing controls recompute:

• true: decoder layers are recomputed in backward; only tensors saved outside them are offloaded.
• false: nothing is recomputed and every saved activation is offloaded, bounded by host RAM and PCIe bandwidth. Can be faster for LoRA/QLoRA, whose activations are small. Not supported for ORPO, SimPO, EBFT or process reward models.

Checkpointing runs exactly as it would without offloading, for adapters and full finetunes alike, honoring gradient_checkpointing_kwargs; the non-reentrant default above applies to rl: runs as well (DeepSpeed ZeRO-3 stays reentrant). hidden_states requires gradient_checkpointing: 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.

activation_offloading: legacy is a deprecated alias of true. activation_offloading: disk is disabled: it ran the same CPU offloader as true and never offloaded to disk.

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	Saved activations → CPU, stream-overlapped	General use; LoRA/QLoRA without recompute (gradient_checkpointing: false)
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 always recomputes each layer. For LoRA/QLoRA, activation_offloading: true with gradient_checkpointing: false can be faster: adapter activations are small and pure offload skips recompute.

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.

Flash Attention 4 (attn_implementation: flash_attention_4, or FA2/FA3 upgraded to it) is savable too. FA4 (flash_attn.cute) registers no torch op, so inside checkpointed layers Axolotl runs it through axolotl::flash_attn_4_fwd / axolotl::flash_attn_4_varlen_fwd. These call FA4’s own forward, and their backward calls FA4’s own backward on the saved output and log-sum-exp, so the kernels and numerics are unchanged and the backward does not re-run the attention forward. Sliding-window FA4 calls are recognized from their window and recomputed like any other (save_sliding_window applies), and offload: true moves the saved FA4 outputs to CPU. FA4 calls with attention sinks (learnable_sink) or mask/score mods, calls outside checkpointed layers and calls under torch.compile go to FA4 directly and are recomputed as before. The kernels-community/flash-attn4 hub kernel is not covered: use flash_attention_4 with flash-attn-4 installed.

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

Requires non-reentrant checkpointing (the default). Composes with activation_offloading: hidden_states or true: layer inputs go to CPU (with true, so does what is saved outside the layers) while the saved ops’ outputs stay on GPU, or go to CPU with offload: true. 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