gradient checkpointing
In plain wordsA memory-saving trick: instead of storing every intermediate result needed for learning, the computer keeps only a few and recalculates the rest when needed, trading extra computing time for memory.
A technique that reduces the memory needed for the backward passbackward pass, n.: The step in training where the network works backward from its error, layer by layer, to calculate how each internal setting should change. by saving only a subset of activations (the checkpoints) during the forward passforward pass, n.: The step where a neural network takes an input and runs it through its layers, one after another, to produce an output — its answer or prediction. and recomputing the discarded ones from the nearest checkpoint when backpropagation reaches them. For a chain of n layers, checkpointing every √n layers cuts activation memory from O(n) to O(√n) for roughly one extra forward pass, often 20–40% more compute in practice; more aggressive schemes reach O(log n) memory with O(n log n) recomputation. It is standard for training large transformerstransformer, n.: The design behind today's chatbots and much of modern AI. It handles text by letting every word look at every other word and decide which of them matter for understanding it. and long-context models, usually applied per transformer block, and selective variants recompute only operations that are cheap to redo but memory-hungry, such as attention scores. Compilers can choose automatically what to rematerialize. It sits in the memory-saving toolkit alongside mixed-precision trainingmixed-precision training, n.: Training a model with shorter, less exact numbers for most of the arithmetic and full-length numbers only where accuracy really matters, which saves a great deal of memory and lets modern chips train far faster., gradient accumulationgradient accumulation, n.: A way to train with large batches on limited hardware: the computer processes several small batches one after another, adds up what each says, and only then makes a single update., ZeROZeRO, prop. n.: A memory-saving technique for training big models on many chips: instead of every chip keeping its own complete copy of all the training bookkeeping, the chips divide it among themselves and share pieces only when needed. and offloadingoffloading, n.: Moving part of an AI model's data or work out of a chip's scarce fast memory onto ordinary memory, a disk or another computer, so that models too big for the chip can still run..
Turning on activation checkpointing let us double the micro-batch at a 30% speed cost.
Origin
Checkpointing originated in automatic differentiation, where Andreas Griewank's 1992 binomial scheme (later the 'revolve' algorithm with Andrea Walther) bounded the memory of reverse mode. Tianqi Chen, Bing Xu, Chiyuan Zhang and Carlos Guestrin popularized it in deep learning with 'Training Deep Nets with Sublinear Memory Cost' (2016), cutting a 1,000-layer residual network's memory from 48 GB to 7 GB.
Usage note
Unrelated to model checkpointscheckpoint, n.: A saved snapshot of a model partway through or at the end of training, like a save point in a video game: training can resume from it after a crash, and finished snapshots are what people download and run. saved to disk. Not to be confused with gradient accumulationgradient accumulation, n.: A way to train with large batches on limited hardware: the computer processes several small batches one after another, adds up what each says, and only then makes a single update., which saves memory by splitting the batch instead.
Related
- backward passbackward pass, n.: The step in training where the network works backward from its error, layer by layer, to calculate how each internal setting should change.
- activationactivation, n.: The numbers a neural network computes inside itself as it processes an input — a reading of which of its artificial neurons are firing, and how strongly, at that moment.
- gradient accumulationgradient accumulation, n.: A way to train with large batches on limited hardware: the computer processes several small batches one after another, adds up what each says, and only then makes a single update.
- mixed-precision trainingmixed-precision training, n.: Training a model with shorter, less exact numbers for most of the arithmetic and full-length numbers only where accuracy really matters, which saves a great deal of memory and lets modern chips train far faster.
- ZeROZeRO, prop. n.: A memory-saving technique for training big models on many chips: instead of every chip keeping its own complete copy of all the training bookkeeping, the chips divide it among themselves and share pieces only when needed.
- offloadingoffloading, n.: Moving part of an AI model's data or work out of a chip's scarce fast memory onto ordinary memory, a disk or another computer, so that models too big for the chip can still run.
- checkpointcheckpoint, n.: A saved snapshot of a model partway through or at the end of training, like a save point in a video game: training can resume from it after a crash, and finished snapshots are what people download and run.
- gradient accumulationgradient accumulation, n.: A way to train with large batches on limited hardware: the computer processes several small batches one after another, adds up what each says, and only then makes a single update.
Sources
- Chen, T., Xu, B., Zhang, C. & Guestrin, C. (2016). Training Deep Nets with Sublinear Memory Cost. arXiv:1604.06174. arxiv.org/abs/1604.06174
- Griewank, A. (1992). Achieving logarithmic growth of temporal and spatial complexity in reverse automatic differentiation. Optimization Methods and Software 1(1), 35–54.
Mentioned in
backpropagation through timebackpropagation through time, n.: The way recurrent networks are trained: the network is unrolled into a long chain, one copy per time step, and errors are traced back through every step to adjust the shared weights., backward passbackward pass, n.: The step in training where the network works backward from its error, layer by layer, to calculate how each internal setting should change., checkpointcheckpoint, n.: A saved snapshot of a model partway through or at the end of training, like a save point in a video game: training can resume from it after a crash, and finished snapshots are what people download and run., computational graphcomputational graph, n.: A map of a calculation drawn as boxes and arrows: each box is a simple operation and each arrow carries a number to the next step. AI software uses it to trace errors backward., full fine-tuningfull fine-tuning, n.: Fine-tuning in which every one of a model's internal numbers is allowed to change — the most flexible option, but also the most expensive in memory and storage., gradient accumulationgradient accumulation, n.: A way to train with large batches on limited hardware: the computer processes several small batches one after another, adds up what each says, and only then makes a single update., model FLOPs utilizationmodel FLOPs utilization, n.: The share of a computer's top calculating speed that actually goes into useful AI training work. Real systems often reach only around a third to a half., sequence parallelismsequence parallelism, n.: Splitting the long string of words a model is working on into chunks handled by different chips, which saves memory and makes it possible to train on very long documents.