As we saw in autodiff, the backward pass needs the saved activations from the forward pass. These intermediate activations must stay in GPU memory, taking up substantial space and frequently causing out-of-memory errors for deep models. Storing and reloading them from GPU memory also generates heavy memory traffic, especially for memory-bound elementwise operations.

To reduce memory consumption, we can trade compute for memory using activation checkpointing. Instead of saving intermediate activations at every op, we only keep them at selected points along the chain, as shown below.

This reduces memory consumption during the forward pass and decreases memory traffic when loading saved activations during the backward pass.

However, when the backward pass reaches an op whose input activation was not saved, such as op2 needing x1, the framework must recompute it on the fly. It does this by running op1 forward again starting from the nearest saved checkpoint x0, as shown below.

Activation checkpointing is essentially trading compute for more memory space: rerunning forward ops consumes additional compute cycles, but in return, it frees up substantial GPU memory.

So it is not suitable for compute-bound operations, since adding extra compute directly worsens the bottleneck and slows down training. Fortunately, modern frameworks handle this selection properly, choosing which operations to save and which to recompute to minimize memory usage without significantly hurting performance.¹

Throughout this chapter, we have seen how frameworks translate high-level Python code into efficient execution on hardware, chaining forward and backward passes together and managing activation memory through checkpointing. So far, however, we have assumed that our models and their computation graphs fit comfortably on a single GPU. As modern models scale to tens or hundreds of billions of parameters, their weights, activations, and optimizer states rapidly outgrow what any individual GPU can hold. In the next chapter, we scale our systems outward, exploring how clusters of GPUs connect and how to distribute both computation and memory across them.