Distributed Training

Learn how to scale model training across many GPUs — by the end you'll be able to explain data vs model parallelism, hand-compute the gradient all-reduce, simulate gradient accumulation, and pick the right framework (PyTorch DDP, FSDP, DeepSpeed) for a given model size.

Learn Distributed Training in our free AI & Machine Learning course — a beginner-friendly interactive lesson with worked examples, a practice exercise and a…

Part of the free AI & Machine Learning course at LearnCodingFast — hands-on lessons with examples you run in your browser, plus practice exercises and a quick quiz.

Imagine a giant banquet that's too much for one cook. You hire a team and split the work — but there's more than one way to split it.

Training GPT-4-scale models took thousands of GPUs running for months. The same ideas scale down: even fine-tuning a 7B model usually needs more than one GPU.

Data parallelism is the most common way to scale. You put a full copy of the model on every GPU (each copy is a worker ), hand each worker a different slice of the batch, and let them all compute gradients at the same time.

The catch: if every worker just applied its own gradient, the copies would drift apart. So after the backward pass, the workers run an all-reduce — they sum every gradient across all workers and divide by the number of workers. Each GPU then applies that same averaged gradient , so all copies stay identical. The example below does the all-reduce by hand so you can see the math.

There are two ways workers can coordinate. In synchronous training, every worker finishes its step and waits at a barrier for the all-reduce before anyone moves on — all copies stay identical, but the slowest worker (a "straggler") holds everyone up.

In asynchronous training, workers push gradients to a central parameter server and grab the latest weights without waiting. It avoids stragglers but workers can train on slightly stale weights, which can hurt accuracy. Modern deep learning almost always uses synchronous all-reduce — it's simpler and more accurate, and fast GPU interconnects (like NVLink) make the waiting cheap.

Big batches train more smoothly, but they don't always fit in GPU memory. Gradient accumulation is the workaround: run several small micro-batches , add up their gradients, and only call the optimizer once at the end. Four micro-batches of 8 give the same effective batch of 32 as one big batch — you trade extra time for less memory.

Notice you divide the accumulated gradient by the number of micro-batches so the update has the same scale as a real batch-of-32 step. The same idea stacks with data parallelism: effective batch = micro_batch × accum_steps × num_gpus .

Data parallelism assumes the whole model fits on one GPU. When it doesn't — say a 70B-parameter model that needs hundreds of gigabytes — you have to split the model itself .

Different layers live on different GPUs — GPU 0 holds layers 1–8, GPU 1 holds layers 9–16, and the activations flow forward through them like an assembly line. To keep GPUs busy, the batch is sliced into micro-batches that are fed in a staggered pipeline.

A single huge layer is split across GPUs — each GPU holds part of the weight matrix and they sync activations within the layer. Used for very wide layers (the attention and MLP blocks in large transformers).

FSDP / ZeRO (shard everything, still data parallel)

FSDP (Fully Sharded Data Parallel) and DeepSpeed ZeRO keep the data-parallel idea but shard the model weights, gradients, and optimizer states so each GPU stores only 1/N of everything. Params are gathered just in time for each layer's forward pass, then released — letting you train models far larger than one GPU's memory.

A rough memory rule for training in FP16: you need about 4× the model size (weights + gradients + optimizer states + activations). A 7B model is ~14 GB in FP16, so ~56 GB to train — close to filling one 80 GB A100 before you've added the batch.

Mixed precision stores most numbers in 16-bit instead of 32-bit. That halves memory and roughly doubles throughput on modern GPUs. Two 16-bit formats matter:

Scaling out (more GPUs, more accumulation) grows the effective batch . Larger batches produce smoother, lower-variance gradients, so you can — and must — take bigger steps. The linear scaling rule : when you multiply the batch by k , multiply the learning rate by k too, and add a few hundred steps of LR warmup so the bigger steps don't blow up early in training.

In practice you rarely write the all-reduce yourself — the framework does it. The most common entry point is PyTorch DistributedDataParallel (DDP) : wrap your model in DDP(...) and it averages gradients automatically during loss.backward() . You launch one process per GPU with torchrun . Read this through — it ties together every idea from the lesson.

Other frameworks fill different gaps: FSDP (built into PyTorch) and DeepSpeed ZeRO shard the model when it no longer fits; Horovod is an older Uber library that adds ring all-reduce to TensorFlow/PyTorch and is still seen in multi-node clusters.

Two workers each computed a gradient for the same 2 parameters. Fill in the blanks so you sum across workers and divide by the worker count — the all-reduce.

Accumulate four micro-batches into an effective batch of 16, then apply the linear scaling rule to the learning rate. Fill in the blanks.

Distributed bugs are sneaky — the code runs but the model trains wrong. Here are the four that bite everyone.

❌ Not scaling the learning rate with the batch

8 GPUs means an 8× larger effective batch, but the LR is still tuned for one GPU:

✅ Fix: apply the linear scaling rule (and warmup):

Stepping the optimizer on the raw model skips DDP's gradient all-reduce, so each GPU diverges:

✅ Fix: always run forward/backward through the DDP-wrapped model:

Giving every worker the full dataset means each sample is trained on N times — and the "extra" data is wasted compute:

✅ Fix: use a DistributedSampler so each GPU gets a disjoint shard:

❌ Communication bottleneck (GPUs idle, waiting)

If the all-reduce of gradients takes longer than the compute, adding GPUs barely helps — you scale poorly:

Put it together with no hints. Simulate one synchronous step across 3 workers: all-reduce the gradients, then update the weights. The expected output is in the comments so you can self-check.

Lesson Complete — you can scale training now!

You can explain data parallelism and its gradient all-reduce, contrast it with model/pipeline/tensor parallelism and FSDP, simulate gradient accumulation, reason about FP16/BF16 mixed precision, and apply the linear LR scaling rule. You also know the four classic bugs to avoid.

🚀 Up next: Model Serving — once a model is trained across many GPUs, learn how to serve it in production reliably.

Practice quiz

• Browse all free courses
• Cheat sheets & quizzes
• Coding dictionary