Distributed Model Training: Data, Model and Pipeline Parallelism

Distributed training explained: data, model and pipeline parallelism, all-reduce, sharded optimizers, checkpointing and failure recovery on GPU clusters.

Advanced 3 min read · Updated

Distributed training spreads model training across many GPUs or machines, either because the data is too large to process in reasonable time on one device or because the model itself doesn't fit in one device's memory. The main strategies are data parallelism (every worker has a full copy of the model and trains on different data), model or tensor parallelism (each worker holds part of each layer) and pipeline parallelism (each worker holds different layers). Large models combine all three.

Data parallelism

Each worker holds a full copy of the model, processes a different slice of each batch, computes gradients, and then all workers average their gradients before updating the weights, so every copy stays identical. The averaging step is an all-reduce, usually implemented as a ring or tree over fast interconnects so each worker sends and receives about the same amount of data regardless of how many workers there are.

periodicTraining datashardedGPU 1full modelGPU 2full modelGPU 3full modelAll-reduceaverage gradientsCheckpointsobject storage
Synchronous data parallelism: workers compute gradients on different shards of the batch, all-reduce them, then apply the same update.
  • Synchronous training (all workers wait for the all-reduce) is the norm: it's deterministic and converges like single-device training with a larger batch.
  • Scaling the global batch usually requires adjusting the learning rate and warm-up.
  • The slowest worker sets the pace (stragglers), so homogeneous hardware and balanced data loading matter.

When the model doesn't fit: sharding and model parallelism

TechniqueWhat's splitCost
Sharded data parallelism (ZeRO, FSDP)Optimizer state, gradients and optionally weights are sharded across data-parallel workers and gathered when neededExtra communication; often the first step for models that almost fit
Tensor (model) parallelismEach layer's matrices are split across devices, which compute parts of every layerHeavy communication every layer; needs very fast links, usually within one machine
Pipeline parallelismGroups of layers live on different devices; micro-batches flow through like an assembly linePipeline bubbles (idle time) at the start and end of each batch
Mixed precision and activation checkpointingLower-precision numbers; recompute activations instead of storing themSome compute overhead for large memory savings

Very large models use 3D parallelism: tensor parallelism within a machine, pipeline parallelism across a few machines, and data parallelism across groups of those.

Input pipelines

Accelerators are expensive, and an idle GPU waiting for data is wasted money. Prefetch and decode data in parallel with training, store training data in large sequential files or shards rather than millions of small objects, shuffle across shards, and make sure each worker reads a disjoint shard of each epoch.

Failures and checkpointing

With thousands of GPUs running for weeks, hardware failures are routine. A synchronous job stops when any worker fails, so recovery is built around checkpoints:

  • Save model weights, optimizer state, data-loader position and random seeds periodically to durable storage.
  • Choose the interval from the failure rate: checkpoint too rarely and a failure loses hours; too often and saving dominates. Asynchronous and sharded checkpointing reduce the pause.
  • Restart automatically on healthy nodes, ideally with spare capacity on standby, and resume from the last checkpoint.
  • Detect slow or faulty nodes (not just dead ones) and replace them, since one bad GPU slows every step.

The training platform around it

  • Scheduling: gang scheduling (all workers of a job start together), quotas per team, and topology-aware placement on well-connected nodes.
  • Reproducibility: record code version, data snapshot, configuration and seeds for each run in an experiment tracker and model registry.
  • Cost control: preemptible capacity for fault-tolerant jobs, utilisation dashboards, and stopping runs whose loss curves have diverged.

Where it’s used

Large language models

Combine tensor, pipeline and sharded data parallelism across thousands of accelerators with frequent checkpoints.

Recommendation models

Huge embedding tables are sharded across devices while the dense layers use data parallelism.

Vision and speech models

Data parallelism across a few machines speeds up training on large datasets.

In an interview

  • Start with data parallelism and explain all-reduce.
  • Move to sharding, tensor or pipeline parallelism only when the model doesn't fit.
  • Design checkpointing and automatic restart from the failure rate.
  • Keep accelerators busy with a fast input pipeline and gang scheduling.

Further reading