What is data parallelism in model training?

Training on one GPU is limited by how many examples it can get through. The simplest way to use 8 GPUs is to give each a full copy of the model and a different eighth of the data. That is data parallelism, and it is the first kind of parallel training almost every team uses.

One step, 8 GPUs

Each training step goes like this:

1. A batch of 256 examples is split into 8 slices of 32.
2. Each GPU runs its slice through its own copy of the model and computes the gradientsGradientThe vector of partial derivatives of the loss function with respect to model parameters, used to update weights during training..
3. The GPUs average their gradients with each other, an operation called all-reduce.
4. Each GPU applies the same averaged gradients, so all 8 copies stay identical.

The result is the same as one huge GPU training on all 256 examples, at close to 8 times the speed. PyTorch's DistributedDataParallel (DDP) does this, and starts averaging some gradients while others are still being computed, so the waiting overlaps the work.

The memory problem, and sharding

Plain data parallelism keeps a full copy of everything on every GPU: the weightsWeightA trainable parameter in a neural network that scales input values and is updated during training., the gradients and the optimizerAdam OptimizerThe most used way to train neural networks. A version of gradient descent that gives every weight its own step size, from a running memory of its past gradients. state. With Adam that is roughly 16 bytes per parameter in mixed-precision training, so a 7B model needs over 100 GB on each GPU before counting activations. The copies are identical, which is wasteful.

Microsoft's ZeRO, from 2019, removed the waste. Each GPU keeps only its share of the optimizer state, the gradients, and in its fullest form the weights too, and fetches the other shares just in time when it needs them. PyTorch's Fully Sharded Data Parallel (FSDP) and DeepSpeed both work this way. It is still data parallelism, since each GPU still trains on its own slice of the data, but the model no longer has to fit on one GPU.

How it fits with the other kinds

Data parallelism scales the data. TensorTensor ParallelismSplitting each layer's weight matrices across several GPUs, so every GPU does part of every layer and they combine their results after each one. and pipeline parallelismPipeline ParallelismPutting different layers of a model on different GPUs and passing work from one to the next like an assembly line, with small batches keeping every stage busy. split the model. Large runs use all 3 together. For serving a model, the same idea is simply running several copies behind a load balancer, each answering different requests.

When it is enough

If the model, its gradients and optimizer state fit on each GPU, plain data parallelism is the simplest and fastest choice. If they do not, sharded data parallelism (FSDP or ZeRO) is usually the next step before splitting the model any further.