The Memory Wall: Why Training Needs More Than One GPU

The first surprise in distributed training is that the model weights are a minority of the memory.

Standard mixed-precision training with the Adam optimizer keeps five things per parameter. A 16-bit copy of the parameter used in the forward and backward pass, at 2 bytes. A 16-bit gradient, at 2 bytes. Then, in 32-bit because they accumulate small updates and would underflow otherwise: a master copy of the parameter at 4 bytes, the first moment estimate at 4 bytes, and the second moment estimate at 4 bytes.

That totals 16 bytes per parameter, and only 2 of them are the weights you would ship for inference.

Run it on a 70-billion-parameter model. Parameters in 16-bit: 140 GB. Gradients: 140 GB. The 32-bit optimizer trio: 840 GB. Total model state: 1120 GB.

An 80 GB accelerator holds one seventieth of that. Fourteen of them, perfectly filled, would hold the state and have nothing left for activations. This is why training is a distributed systems problem and inference often is not.

A reasonable objection: if the forward pass is happy in 16 bits, why keep 32-bit copies at all?

Because updates are tiny relative to weights. A typical weight might be around 0.01 and a typical update around 0.0000001. In 16-bit floating point, adding a number that small to a number that large returns the large number unchanged. The update is silently discarded, and training stalls without any error.

The master weights in 32-bit exist to accumulate those updates faithfully. The optimizer applies the update to the 32-bit copy, then casts down to 16 bits for the next forward pass, so the small increments accumulate across steps until they matter.

The two moment estimates are simply what Adam is. It maintains a running average of the gradient and of the squared gradient, each the same shape as the parameters, which is what makes Adam cost three times more memory than plain stochastic gradient descent.

Switching optimizer is therefore a real memory lever, not a detail. It is also a change to training dynamics, which is why it is rarely the first thing tried.

Model state is fixed once you choose a model and optimizer. Activation memory is not, and it is the term you actually control.

The backward pass needs the intermediate values computed during the forward pass in order to apply the chain rule. Every layer output, every attention score matrix, every intermediate inside the feed-forward block must be retained from the moment it is produced until its gradient is computed.

The scaling is what makes this dangerous. Activation memory grows with batch size times sequence length times hidden dimension times layer count. Doubling the batch doubles it. Doubling the context doubles it again, and for the attention score matrices specifically, the cost has historically grown with the square of sequence length, which is what memory-efficient attention implementations exist to avoid.

So a configuration that fits comfortably at 2k context can fail at 8k with everything else unchanged. Model state stayed identical; activations quadrupled.

This is the term that most often causes the first out-of-memory error, and it has two standard remedies.

Both remedies trade compute for memory, and both are almost free in accuracy terms, which makes them the first things to reach for.

Gradient accumulation splits a large batch into micro-batches. Run a forward and backward pass on each, add the gradients into a buffer without stepping the optimizer, and step once at the end. Activation memory is now set by the micro-batch, while the optimizer sees the full batch. The mathematics of the update are identical to running the large batch at once, aside from floating point summation order.

Activation checkpointing, also called recomputation, discards most activations during the forward pass and keeps only the boundaries between segments. When the backward pass reaches a segment, it re-runs that segment's forward pass to regenerate what it needs.

The canonical version keeps only layer boundaries, which reduces activation memory roughly to the square root of its original size at the cost of one extra forward pass, commonly around thirty percent more compute. Selective variants recompute only the cheap-to-recompute, expensive-to-store operations, which gets most of the saving for a fraction of the cost.

The oldest form of distributed training is data parallelism, and its limitation is precise and worth stating exactly.

Every GPU holds a complete copy of the model, the gradients and the optimizer state. Each processes a different slice of the batch. After the backward pass, the gradients are averaged across all GPUs with an all-reduce, so every replica applies the same update and stays in lockstep.

What this buys is throughput. Sixty-four GPUs process sixty-four times as much data per step.

What it does not buy is a single byte of memory relief. Every one of those sixty-four GPUs still holds the full 1120 GB of state for a 70B model, which means the model has to fit on one device before data parallelism is available at all.

Stated plainly: data parallelism scales the batch, not the model. A model too large for one accelerator cannot be trained this way regardless of how many accelerators you own, and that gap is what the rest of this path fills.

Since every distributed strategy is ultimately a choice about communication, it is worth knowing what the basic operation costs.

An all-reduce combines a value held by every participant, typically by summing, and leaves every participant with the result. The naive implementation sends everything to one node, which makes that node's link the bottleneck and scales badly.

The standard implementation is ring all-reduce. Participants form a logical ring, the data is split into as many chunks as there are participants, and the operation runs in two phases: a reduce-scatter that walks chunks around the ring accumulating sums, then an all-gather that walks the completed chunks back around.

The result is that each participant sends and receives roughly twice the size of the gradient buffer, independent of how many participants there are. Communication volume per GPU does not grow with cluster size, which is what makes data parallelism scale at all.

What does grow is latency, since the ring has more hops, and sensitivity to the slowest link. One degraded network interface slows every participant.

Laying the terms beside the strategies makes the rest of this path a lookup rather than a list.

There are four consumers of memory. Parameters, gradients, optimizer states, and activations. Every technique targets specific ones, and knowing which tells you what to reach for when a run fails.

Sharding, meaning ZeRO and FSDP, attacks the first three by splitting them across data-parallel ranks so no GPU holds a full copy. It leaves activations untouched.

Tensor parallelism splits individual matrices across GPUs, so it reduces parameters, gradients and optimizer states within a layer, and it also reduces activations because each GPU only computes part of them.

Pipeline parallelism assigns whole layers to different GPUs, so each holds a fraction of every term, but introduces idle time.

And recomputation attacks activations alone, paying in compute.

The diagnostic question when a run runs out of memory is which of the four terms grew. Increasing sequence length points at activations; switching optimizer points at states.

It is tempting to treat batch size purely as a memory and throughput knob. It is also a learning-dynamics decision, and forgetting that produces runs that fit perfectly and train badly.

Larger batches give lower-variance gradient estimates, which permits larger learning rates and better hardware utilisation. But the benefit saturates. Past a certain point, doubling the batch stops halving the variance in any useful way, and you are spending twice the compute per step for a marginally better direction.

The common heuristic when changing batch size is to scale the learning rate with it, either linearly or with the square root, combined with a warmup period long enough that the early steps do not diverge.

This matters for distributed training specifically because the global batch is the product of micro-batch size, gradient accumulation steps, and data-parallel degree. Changing your parallelism layout to fix a memory problem silently changes the global batch, and therefore the optimisation problem.

Hold the global batch fixed when you reconfigure, or you are no longer running the same experiment.

Every strategy in this path moves data between accelerators, and the hardware makes a sharp distinction between two kinds of link.

Inside a node, GPUs are connected by a dedicated high-bandwidth fabric, NVLink on NVIDIA systems, offering bandwidth on the order of hundreds of gigabytes per second between peers.

Between nodes, traffic crosses a network, typically InfiniBand or high-speed Ethernet, at bandwidth substantially lower than the intra-node fabric and with meaningfully higher latency.

That gap of roughly an order of magnitude is the single most important physical fact in choosing a parallelism layout, and it explains a rule you will meet everywhere: put the chattiest parallelism inside the node.

Tensor parallelism communicates twice per transformer layer, so it belongs on NVLink and is conventionally capped at the number of GPUs in one node, typically eight. Pipeline parallelism communicates once per stage boundary, a much smaller volume, so it tolerates crossing nodes. Data parallelism communicates once per step and tolerates it best.

That ordering is not a convention. It follows directly from the bandwidth hierarchy.

Given a model that does not fit, the techniques have a natural ordering by cost, and following it avoids a great deal of unnecessary complexity.

First, reduce the micro-batch and add gradient accumulation. Costs nothing but a little efficiency, changes no mathematics, and often solves the problem outright.

Second, enable activation checkpointing. Costs roughly thirty percent more compute for a large reduction in the activation term.

Third, shard the model state with ZeRO or FSDP. This is where genuine distribution begins, and stage 3 alone takes a 70B model from 1120 GB per GPU to 17.5 GB across 64 GPUs.

Fourth, add tensor parallelism inside each node, when a single layer's activations are themselves too large.

Fifth, add pipeline parallelism across nodes, when the model is too large even for sharding plus tensor parallelism, and accept the bubble.

The ordering matters because each step adds operational complexity, and complexity is what makes long runs fail. Most models people actually train stop at step three.