Gradient Checkpointing

Gradient checkpointing in Axolotl is configured and documented as activation checkpointing.

NoteThis page has moved

Gradient checkpointing is now documented under Activation Checkpointing & Offloading.

Gradient checkpointing and activation checkpointing are two names for the same technique. During the forward pass each decoder layer keeps only its input and drops its intermediate activations; during the backward pass it runs its forward again to rebuild them. Activation memory drops to about one hidden-state tensor per layer, at the cost of roughly one extra forward pass per step.

PyTorch, DeepSpeed, accelerate and transformers (from 5.23) call it activation checkpointing, and Axolotl’s config follows them: activation_checkpointing: true replaces gradient_checkpointing: true, which still works but is deprecated, and the same block configures selective checkpointing and activation offloading.

activation_checkpointing: true

On the new page:

• Enabling activation checkpointing, including reentrant and the extra torch.utils.checkpoint arguments
• Selective activation checkpointing (SAC)
• Activation offloading
• Layer offloading
• Migrating from gradient_checkpointing and the other legacy keys