Lensa ML
Lensa ML

Distributed Training

Step 1 of 6: Why Distribute?

Too much memory and too much compute for one GPU

Training tokens2.00T
Training state, 16 bytes per parameter × 6.74B parameters2 + 2 + 12 B = 107.8 GB80 GB GPUfp16 weightsfp16 gradientsfp32 Adam statesTime to train at 150 TFLOP/s per GPU, perfect scaling (illustrative)1 GPU17.1 years64 GPUs98 days1,024 GPUs6 days
Training state
107.8 GB
Over one 80 GB GPU
1.35×
Compute, 6ND
8.1e22 FLOP
On one GPU
17.1 years

Training on 2.00T tokens takes 8.1e22 FLOP, about 17.1 years on one GPU at this rate. The 107.8 GB of state also overflows the GPU before the first step runs.

Drag Training tokens from 50B up to 2T and watch the time bars. The 16 bytes per parameter follow the mixed-precision Adam accounting in the ZeRO paper: 2 for fp16 weights, 2 for fp16 gradients, and 12 for the fp32 master weights, momentum, and variance. Activations come on top and grow with batch size and sequence length. The 6ND rule counts about 2 FLOP per parameter per token for the forward pass and 4 for the backward pass. Llama 2 7B was trained on 2T tokens, so its compute is about 8.1e22 FLOP.

Step 2 of 6: Data Parallelism

Full copies, split batch, all-reduce the gradients

GPUs (data-parallel replicas)8
GPU 0128 seqGPU 1128 seqGPU 2128 seqGPU 3128 seqGPU 4128 seqGPU 5128 seqGPU 6128 seqGPU 7128 seqring all-reduce13.5 GB of gradientsEach holds the full 107.8 GB copy (red bar: over 80 GB)Global batch 1,024 sequences of 4,096 tokens, split evenly
State per GPU
107.8 GB
Sequences per GPU
128
Sent per GPU, per step
23.6 GB
Compute per step
2.4 min

At 8 GPUs each one computes only 128 sequences (2.4 min), while the all-reduce still sends 23.6 GB per GPU, close to the 2× limit. Across 50 GB/s links that is about 0.47 s per step (illustrative), and every copy still needs 107.8 GB.

Drag GPUs from 1 to 64 and compare the sent bytes with the compute time. In a ring all-reduce each GPU sends its gradient in N chunks, first summing around the ring (reduce-scatter), then passing the summed chunks on (all-gather), which totals 2(N − 1)/N times the gradient. That figure barely grows with N, so the method scales well. Frameworks such as PyTorch DDP overlap the all-reduce with the backward pass, starting on the last layers' gradients while earlier layers are still being computed.

Step 3 of 6: ZeRO Sharding

Split optimizer states, gradients and weights

Data-parallel GPUs N8
Sharding
What GPU 0 keeps of each tensor group (highlighted part)fp16 weights13.5 GB, full copyfp16 gradients13.5 GB, full copyAdam states80.9 GB / 8 = 10.1 GBGPU 0 total37.1 GB80 GBActivations and buffers not counted
State per GPU
37.1 GB
Smaller by
2.9×
Sent per GPU, per step
23.6 GB
Fits 80 GB
Yes

ZeRO-1 splits the optimizer states into 8 shards, bringing GPU 0 to 37.1 GB. Communication matches plain data parallelism, 23.6 GB per GPU per step, and the 27.0 GB left unsharded sets the floor.

Drag Data-parallel GPUs N and switch between None, ZeRO-1, ZeRO-2, and ZeRO-3. The per-GPU formulas from Rajbhandari et al. are (4 + K/N)Ψ bytes for stage 1, (2 + (2 + K)/N)Ψ for stage 2, and (4 + K)Ψ/N for stage 3, with K = 12 for Adam. The paper's own example, a 7.5B model on 64 GPUs, drops from 120 GB to 1.9 GB. Stages 1 and 2 add no traffic because an all-reduce is already a reduce-scatter followed by an all-gather, and ZeRO updates each shard between the two.

Step 4 of 6: Tensor Parallelism

Split each matrix multiply across GPUs

Tensor-parallel GPUs t4
Link between them
input Xup: column blocklocal act.down: row blockpartialall-reduceXGPU 0GPU 0GPU 1GPU 1GPU 2GPU 2GPU 3GPU 3ΣZOne MLP block split 4 ways; attention splits the same way, by headsIllustrative: 150 TFLOP/s per GPU, NVLink 450 GB/s, network 50 GB/s
Weights per GPU, one layer
50.6M
State per GPU
27.0 GB
Sent per GPU, per layer
201 MB
Comm ÷ compute
0.05×

Each GPU multiplies by 1/4 of every matrix, and the partial results are summed by 4 all-reduces per layer, 201 MB per GPU for one sequence. Over NVLink that takes 0.45 ms against 8.3 ms of compute.

Move Tensor-parallel GPUs t through 1, 2, 4, and 8 and switch the link. Megatron-LM splits the first MLP matrix by columns, so each GPU applies the nonlinearity to its own slice with no exchange, then splits the second matrix by rows so each GPU produces a partial sum. Attention splits by heads in the same pattern. That leaves one all-reduce after attention and one after the MLP in the forward pass, two more in the backward pass. Because these exchanges sit on the critical path, t rarely exceeds the 8 GPUs in one server.

Step 5 of 6: Pipeline Parallelism

Split the layers, stream micro-batches

Pipeline stages p4
Micro-batches m4
GPipe schedule: one row per stage, time runs rightGPU 0L0–7GPU 1L8–15GPU 2L16–23GPU 3L24–3111223344112233441122334411223344time, 21 units (backward counted as 2× forward); gray = idle bubbleforward, numbered by micro-batchbackward
Bubble (idle time)
42.9%
Layers per GPU
8
State per GPU
27.0 GB
Busy units per GPU
12 of 21

4 micro-batches across 4 stages leave a 42.9% bubble, (p − 1)/(m + p − 1). Between stages only one micro-batch's activations move, 34 MB for a 4,096-token sequence.

Drag Pipeline stages p and Micro-batches m and watch the gray triangles at the start and end. GPipe (Huang et al., 2018) stores the activations of all m micro-batches until their backward passes run, so memory grows with m. The 1F1B schedule from PipeDream and Megatron-LM alternates one forward and one backward per stage, which keeps at most p micro-batches in flight with the same bubble. Interleaved schedules give each GPU several smaller chunks of layers and cut the bubble further.

Step 6 of 6: 3D Parallelism

Tensor inside a node, pipeline across, data on top

Tensor parallel
Pipeline stages
ZeRO on DP
Data-parallel replicas8
Data-parallel replicas →Each box is one tensor-parallel group of 8 GPUs, kept inside one nodestage 0, L0–7stage 1, L8–15stage 2, L16–23stage 3, L24–318 TP × 4 PP × 8 DP = 256 GPUs on 32 nodesIllustrative configs; days assume 150 TFLOP/s per GPU and ignore communication
GPUs (nodes)
256 (32)
State per GPU
1.16 GB
Pipeline bubble
2.3%
2T tokens take
25 days

256 GPUs each hold 1.16 GB of state, leaving room for activations. Each pipeline gets 128 micro-batches, so the bubble is 2.3%. At this scale, 2T tokens take about 25 days with ideal scaling.

Click a preset, then change the toggles and Data-parallel replicas. Narayanan et al. (2021) give the usual recipe: keep tensor parallelism within a node, where NVLink is fast, use pipeline stages across nodes, since they only pass activations between neighbors, and scale out with data parallelism. They trained a 1-trillion-parameter model this way on 3,072 GPUs. Larger data-parallel degree means fewer micro-batches per pipeline at a fixed global batch, so the bubble rises as the cluster grows.

What is distributed training in deep learning?

Distributed training is training one neural network on many GPUs at once by splitting the work between them. Data parallelism splits the batch, tensor parallelism splits each matrix multiply, and pipeline parallelism splits the stack of layers. Sharding methods such as ZeRO split the training state itself, so no single GPU has to hold all of it.

Every large language model is trained this way, because both the memory and the compute exceed one accelerator. Llama 2, GPT-3, and similar models ran on thousands of GPUs at once, usually with tensor parallelism inside each 8-GPU server, pipeline parallelism across servers, and data parallelism on top. The same tools, such as PyTorch FSDP, DeepSpeed, and Megatron-LM, also make fine-tuning a 7B model practical on a handful of GPUs.

Key ideas

  • Training a 7B model with Adam in mixed precision needs about 108 GB of weights, gradients, and optimizer state, more than one 80 GB GPU.
  • Data parallelism copies the model to every GPU and all-reduces the gradients, which speeds up compute but saves no memory.
  • ZeRO shards optimizer states, then gradients, then weights across data-parallel GPUs, cutting per-GPU memory by up to N times.
  • Tensor parallelism splits each layer's matrices and needs all-reduces inside every layer, so it stays within one fast-linked node.
  • Pipeline parallelism assigns layers to stages and leaves an idle bubble of (p − 1)/(m + p − 1) of each step.
  • Large runs combine all three, with the total GPU count equal to the tensor, pipeline, and data-parallel degrees multiplied together.

Common questions

What is the difference between data parallelism and model parallelism?

Data parallelism gives every GPU the whole model and a different slice of the batch. Model parallelism splits the model itself, either within layers (tensor parallelism) or between layers (pipeline parallelism), so each GPU holds only part of the weights. Data parallelism is simpler and is tried first, and model parallelism is added when the model or its state no longer fits.

Is FSDP the same as ZeRO?

PyTorch Fully Sharded Data Parallel implements the same idea as ZeRO stage 3. Each GPU keeps a shard of the weights, gradients, and optimizer state and gathers full layers just before using them. Its sharding strategies also include a mode that matches ZeRO-2, keeping full weights but sharding gradients and optimizer state.

Does a larger global batch change training results?

Yes, it can. More data-parallel GPUs at the same per-GPU batch raise the global batch, which lowers gradient noise and usually needs a retuned learning rate and warmup. Past a model-dependent critical batch size, extra data per step gives little extra progress, so adding GPUs stops shortening training.

How many GPUs does it take to fine-tune a 7B model?

Full fine-tuning with Adam needs about 108 GB of state plus activations, so it typically runs on 2 to 8 80 GB GPUs with ZeRO or FSDP sharding. Parameter-efficient methods such as LoRA freeze the base weights and train small adapters, which brings the memory close to the 13.5 GB of fp16 weights. With 4-bit base weights (QLoRA), a single 24 GB GPU is often enough.

Related explainers

References

  1. Rajbhandari et al. (2019), "ZeRO: Memory Optimizations Toward Training Trillion Parameter Models"
  2. Shoeybi et al. (2019), "Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism"
  3. Huang et al. (2018), "GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism"
  4. Narayanan et al. (2021), "Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM"

Last updated