ML0039 Distributed Training

Written by

What are the two main distributed training approaches for machine learning?

Answer

The two main distributed training approaches are data parallelism and model parallelism, and they answer different bottlenecks. In data parallelism, the dataset is split across many devices, each holding a complete copy of the model; every device computes gradients on its own mini-batch, and after each step the gradients are aggregated and synchronized so all copies stay identical. It scales training when the data is large but the model still fits in one device’s memory. In model parallelism, the model itself is too large for one device, so different parts of the model (e.g., different layers) are placed on different devices, and data flows through them in sequence, with activations passed between devices. It is the enabling technique for today’s very large models, at the cost of more complex communication. Modern large-scale training often combines both (plus pipeline and tensor parallelism) in hybrid strategies.

(1) Data Parallelism: Full model copy per device, different data shards, gradients synchronized after each step: scales with data.
(2) Model Parallelism: The model is split across devices, data flows sequentially through the parts: scales with model size.
(3) Trade-off: Data parallel communicates gradients per step; model parallel communicates activations between stages.

Data parallelism with model replicas on three GPUs training on data shards versus model parallelism with layers split across three GPUs

Figure 1: The two approaches: data parallelism (left) replicates the whole model on every device and splits the data, synchronizing gradients after each step; model parallelism (right) splits the model itself across devices and passes activations down the chain.

Mathematical Formulation:
$g = \frac{1}{D} \sum_{d=1}^{D} g_d$
$g_d = \nabla_\theta \mathcal{L}_d(\theta) \quad \text{(data-parallel gradient sync)}$
$h^{(d+1)} = f^{(d+1)}\big(h^{(d)};\, \theta^{(d+1)}\big)$
$\theta^{(d+1)} \text{ lives on device } d+1 \text{ (model parallel)}$

Where:

$D$ is the number of devices; $g_d$ the gradient computed by device $d$ on its data shard; $g$ the synchronized average applied to every replica.
$h^{(d)}$ is the activation handed from device $d$ to device $d+1$ ; $\theta^{(d)}$ the slice of parameters owned by device $d$ .