Gradient Checkpointing - PyTorch Basics

Gradient Checkpointing

What you'll build: the two halves of activation checkpointing: a forward pass that keeps only selected activations, and a recompute function that rebuilds any activation from the nearest saved one.

Background

Backpropagation needs the activations of the forward pass, so memory grows linearly with depth: O(n) for n layers. Checkpointing stores only some of them and recomputes the rest during the backward pass. With sqrt(n) evenly spaced checkpoints, memory falls to O(sqrt(n)) at the cost of roughly one extra forward pass. It is how large transformers fit long sequences on a GPU (torch.utils.checkpoint).

Conventions

• layers is a list of callables; activation i is the output of layers[i].
• The input x is always saved, under the key -1.
• checkpoint_indices lists which layer outputs to keep.

Algorithm

checkpoint_forward(layers, x, checkpoint_indices)

1. saved = {-1: copy of x}.
2. Run the layers in order. After layer i, if i is in checkpoint_indices, store a copy of its output under key i.
3. Return (final_output, saved).

checkpoint_recompute(layers, saved_activations, checkpoint_indices, layer_idx)

1. Find the nearest saved key strictly below layer_idx: start = max(k for k in saved if k < layer_idx) (-1 always qualifies).
2. Starting from saved[start], run layers[start + 1], ..., layers[layer_idx].
3. Return the result, the output of layer layer_idx.

Inputs

• layers: list of functions mapping an array to an array.
• x: input array; checkpoint_indices: list of layer indices (may be empty).
• layer_idx: index of the activation to rebuild, 0 <= layer_idx < len(layers).

Expected Return Format

• checkpoint_forward returns a tuple (output, saved): the final activation and a dict {index: activation} containing exactly -1 and the checkpoint indices.
• checkpoint_recompute returns the activation of layer layer_idx.

Edge Cases and Conventions

• Store copies, so later in-place changes can't corrupt the checkpoints.
• Recompute runs only the layers after the nearest checkpoint, never the whole network.

Constraints

• PyTorch track: torch.utils.checkpoint isn't available in the browser; run the forward under torch.no_grad() (no graph, so no stored intermediates) and recompute with autograd enabled.

Follow-up

• Why is sqrt(n) the best number of evenly spaced checkpoints?
• Which layers are cheapest to recompute and most expensive to store in a transformer block?

Input: layers = [x*2, x+1, x*3, x-1], checkpoint_forward(layers, np.array([1.0]), [1, 3])
Output: (array([8.]), {-1: array([1.]), 1: array([3.]), 3: array([8.])})

Input: checkpoint_recompute(layers, saved, [1, 3], 2)
Output: array([9.])

import numpy as np

def checkpoint_forward(layers, x, checkpoint_indices):
    """
    Forward pass that saves only the input and the checkpoint activations.

    Args:
        layers: List of layer functions
        x: Input array
        checkpoint_indices: Indices of the layers whose outputs to save

    Returns:
        (output, saved): the final output and a dict index -> activation
        (key -1 holds the input)
    """
    # Your code here
    pass

def checkpoint_recompute(layers, saved_activations, checkpoint_indices, layer_idx):
    """
    Recompute the output of layer layer_idx from the nearest saved activation below it.

    Args:
        layers: List of layer functions
        saved_activations: Dict index -> activation from checkpoint_forward
        checkpoint_indices: Indices that were saved
        layer_idx: Layer whose output to rebuild

    Returns:
        Activation (output) of layer layer_idx
    """
    # Your code here
    pass