Four Ways to Split a Model

The Training at Scale board draws why one GPU cannot hold a model in training, and how the work gets cut up.

Before you start. From The GPU (thousands of cores, and the memory that feeds them): a GPU is thousands of simple calculators with memory that is very fast but small, 80 GB on an A100. From What a Neural Network Is: a weight is one adjustable number inside the network, a dial, and a model’s size is how many weights it has. From How It Learns (loss, gradient descent, backprop): training shows the model examples, measures how wrong it was, and nudges every weight a little; the list of which-way-to-nudge numbers, one per weight, is the gradient. From Why the Loss Won't Go Down: Adam, the standard way to take those nudges, remembers two extra numbers for every weight. A number can be stored in 4 bytes (careful) or 2 bytes (half the room, a little rougher).

The idea in 30 seconds

Training a 7B model needs about 16 bytes per weight: the weight, its gradient, a full-precision 4-byte master copy of each weight, and the two extra numbers the Adam optimizer keeps. That is 112 GB, more than one 80 GB A100.
Data parallelism copies the whole model to every GPU and splits the batch of examples; an all-reduce, a step where the GPUs swap and average their gradients, makes the result exact. It saves time, not memory.
Tensor parallelism cuts every matrix into columns across GPUs; it talks inside every layer, so it stays inside one machine.
Pipeline parallelism gives each GPU a run of layers; GPUs sit idle while the pipeline fills and drains unless the batch is cut into many small slices.
ZeRO (a 2019 Microsoft method) keeps data parallelism but gives each GPU one slice of the optimizer numbers, gradients and weights, trading memory for network traffic.
Mixed precision does the heavy arithmetic in bf16, a 2-byte number format with float32’s range but little detail, and keeps a 4-byte float32 copy of the weights so tiny updates are not rounded away.

Read the story · Skip to the numbers