Phase 4: Training LLMs · ~70 minutes · Python

Gradient Checkpointing and Activation Recomputation

Backprop keeps every intermediate activation. At 70B parameters and 128K context that is 3 TB of activations per rank. Checkpointing trades FLOPs for memory: recompute instead of save. The question is which segments to drop, and the answer is not "all of them."

Hiring signal: Understanding of gradient checkpointing and activation recomputation internals

What you will learn

Implement gradient checkpointing and activation recomputation from scratch
Understand the math and intuition behind the algorithm
Use production libraries for the same task
Ship a reusable artifact

Introduction

Type: Build Languages: Python (with numpy, optional torch) Prerequisites: Phase 10 Lesson 04 (Pre-Training Mini-GPT), Phase 10 Lesson 05 (Scaling & Distributed) Time: ~70 minutes

The Problem

Training a transformer stores, for each layer, the inputs to every op that is differentiated in backward: the attention inputs, the Q/K/V projections, the softmax output, the FFN inputs, the norm outputs, and the residual stream. For a layer with hidden size d, sequence length L, batch B, this is on the order of 12 * B * L * d floats per layer.

For d=8192, L=8192, B=1, that's 800 MB/layer in BF16. A 64-layer model is 51 GB of activations — and that's before you multiply by microbatch size, before you add attention-softmax intermediates (L^2 per head), and before you factor tensor-parallel partial copies.

The two-sided bill: BF16 weights plus optimizer state might fit in 80GB, but activations push you past. Gradient checkpointing (aka activation recomputation) is the standard fix. Drop most activations; redo the forward during backward to get them back. Cost: extra FLOPs. Benefit: memory drops by the ratio of checkpoint segments to total layers.

Done naively, checkpointing costs roughly 33% more forward-pass FLOPs per step. Done well — selective checkpointing per the "smart selection" of Korthikanti et al. — you save 5x memory for under 5% FLOP overhead. And with FP8 matmuls, FSDP offload, and expert-parallel MoE this really matters: you can't afford either the memory or the wasted compute.

Unlock the full lesson

You've read the first 2 sections. The rest of this lesson covers The Concept, Build It, Use It, Ship It, Exercises, Key Terms, Further Reading — plus a hands-on lab, quiz, and project artifact.

Create a free account to unlock Phase 0 and Phase 1 of every course — no credit card.

Key terms in this lesson

Softmax