Distributed Training Strategies #

Distributed training combines devices to increase throughput, fit larger models, or shorten time to solution. The central design question is what to partition: examples, model parameters, layers, tensors, or the training schedule.

Course coverage:

Data parallelism
Model parallelism
Pipeline parallelism and micro-batches
Gradient checkpointing
Mixed-precision training and loss scaling

Learning Objectives #

By the end of this page, you should be able to:

compare data, model, and pipeline parallelism
calculate global batch size, speedup, efficiency, and effective throughput
explain gradient aggregation through all-reduce
calculate simple pipeline time and utilisation
quantify the memory–computation trade-off of checkpointing
explain mixed precision and loss scaling

Big Picture #

flowchart TD
    A["Training Constraint"] --> B["More Data Throughput"]
    A --> C["Model Too Large"]
    A --> D["Activation Memory Too High"]
    B --> E["Data Parallelism"]
    C --> F["Model or Pipeline Parallelism"]
    D --> G["Checkpointing or Mixed Precision"]

    style A fill:#E1F5FE
    style B fill:#C8E6C9
    style C fill:#FFF9C4
    style D fill:#EDE7F6
    style E fill:#E1F5FE
    style F fill:#C8E6C9
    style G fill:#FFF9C4

1. Data Parallelism ☆ #

Each worker holds a complete model replica and processes a different shard of the mini-batch.

For worker i:

\[ g_i = \frac{1}{b_i}\sum_{x\in B_i}\nabla\ell(w;x) \]

For equal local batch sizes, gradients are averaged:

\[ g = \frac{1}{p}\sum_{i=1}^{p}g_i \]

Every worker applies the same update and therefore retains the same parameter values.

Global Batch Size #

\[ B_{\text{global}}=pB_{\text{local}} \]

With 8 workers and local batch size 32, global batch size is 256.

All-Reduce #

An all-reduce combines gradients across workers and returns the aggregate to every worker. It avoids a single central server but introduces a synchronisation point. Training step time is determined by the slowest worker plus the collective.

\[ T_{\text{step}} \approx \max_i(T_{\text{compute},i}) + T_{\text{all-reduce}} \]

2. Model Parallelism ☆ #

Model parallelism divides model parameters or operations across devices. It is necessary when the model or its intermediate state does not fit on one device.

Two common forms are:

layer or stage partitioning: consecutive groups of layers reside on different devices
tensor parallelism: one large matrix or tensor operation is sliced across devices

Model parallelism communicates activations in the forward pass and activation gradients in the backward pass. It can reduce per-device parameter memory but introduces dependencies between partitions.

Tensor-Slicing Example #

For a matrix multiplication Y = XW, split the output columns of W:

\[ W=[W_1\;W_2], \qquad Y=[XW_1\;XW_2] \]

Two devices can compute XW₁ and XW₂ concurrently, after which the output slices are concatenated.

3. Pipeline Parallelism ☆ #

Pipeline parallelism assigns consecutive model stages to different devices and splits a training batch into m micro-batches. Once the pipeline is full, different stages work on different micro-batches at the same time.

For s balanced stages, m micro-batches, and per-stage time t, ideal forward pipeline time is:

\[ T_{\text{pipeline}}=(m+s-1)t \]

Sequential stage execution would take mst, so ideal speedup is:

\[ S=\frac{mst}{(m+s-1)t}=\frac{ms}{m+s-1} \]

Ideal stage utilisation is:

\[ U=\frac{m}{m+s-1} \]

The unused slots during pipeline fill and drain form the pipeline bubble.

Worked Numerical: Pipeline Bubble #

Let s = 4, m = 8, and t = 10 ms.

\[ T_{\text{pipeline}}=(8+4-1)10=110\text{ ms} \] \[ S=\frac{8\times4\times10}{110}\approx2.91, \qquad U=\frac{8}{11}\approx72.7\% \]

Increasing micro-batches reduces the bubble fraction, but very small micro-batches may underutilise the accelerator and increase scheduling overhead.

4. Comparing Parallel Strategies #

Strategy	Partitioned Quantity	Main Communication	Best When	Main Limitation
Data parallel	Training examples	Gradients	Model fits on each device	Gradient synchronisation and global batch growth
Tensor parallel	Large tensor operations	Partial activations/results	Individual layers are too large	Frequent fine-grained communication
Pipeline parallel	Consecutive layer groups	Activations between stages	Model has balanced sequential stages	Pipeline bubbles and stage imbalance
Hybrid	Data, tensor, and stages	Several collectives	Very large models and clusters	Complex placement and tuning

5. Gradient Checkpointing ☆ #

Backpropagation normally stores intermediate activations from the forward pass. For a deep model, activation memory may exceed parameter memory.

Gradient checkpointing stores only selected activations. During the backward pass, discarded activations are recomputed from the nearest saved checkpoint.

If a model has L equally sized layer activations of size M_layer, storing all activations costs:

\[ M_{\text{full}}=L M_{\text{layer}} \]

If only C checkpoints are stored:

\[ M_{\text{checkpoint}}\approx C M_{\text{layer}} \]

Worked Numerical: Activation Memory #

A 48-layer model produces 20 MB of saved activation per layer.

\[ M_{\text{full}}=48\times20=960\text{ MB} \]

Storing 8 checkpoints uses approximately:

\[ M_{\text{checkpoint}}=8\times20=160\text{ MB} \]

The approximate activation-memory reduction is 800 MB, but the forward operations between checkpoints must be recomputed during backpropagation.

More checkpoints use more memory and less recomputation. Fewer checkpoints save more memory and require more recomputation.

6. Mixed-Precision Training ☆ #

Mixed precision uses lower precision for suitable tensor operations while retaining higher precision where numerical range or accumulation accuracy is important.

A common arrangement is:

FP16 or BF16 for forward and backward tensor operations
FP32 for master weights or sensitive accumulations
conversion between representations around the optimiser update

Potential benefits include:

lower parameter, gradient, and activation memory
higher accelerator throughput
lower memory-bandwidth demand

Memory Numerical #

One billion FP32 values require:

\[ 10^9\times4\text{ bytes}=4\text{ GB} \]

The same number of FP16 values require 2 GB. Actual training memory includes weights, gradients, optimiser states, and activations, so the complete reduction may not be exactly one half.

7. Loss Scaling #

FP16 has limited dynamic range. Very small gradients can underflow to zero. Loss scaling multiplies the loss by scale S before backpropagation:

\[ L'=SL, \qquad \nabla L'=S\nabla L \]

Before the optimiser update, gradients are divided by S:

\[ g=\frac{g'}{S} \]

A scale that is too small may not prevent underflow; a scale that is too large may cause overflow. Dynamic loss scaling adjusts the scale after checking for invalid values.

Worked Numerical: Loss Scaling #

If a gradient is 3 × 10⁻⁸ and S = 1024, the scaled gradient is:

\[ 3\times10^{-8}\times1024=3.072\times10^{-5} \]

After safe computation, dividing by 1024 restores the original mathematical gradient.

8. Speedup, Efficiency, and Overhead #

\[ S_p=\frac{T_1}{T_p}, \qquad E_p=\frac{S_p}{p} \]

Parallel overhead expressed in processor-time units is:

\[ T_o=pT_p-T_1 \]

Worked Numerical: Training Run #

A training job takes 480 minutes on one GPU and 75 minutes on eight GPUs.

\[ S_8=\frac{480}{75}=6.4, \qquad E_8=\frac{6.4}{8}=0.8 \] \[ T_o=8(75)-480=120\text{ GPU-minutes} \]

The run achieves 80% efficiency and incurs 120 GPU-minutes of parallel overhead.

9. Strategy Selection #

Use data parallelism when the model fits on one device and throughput is the goal. Use tensor or layer partitioning when individual operations or the model exceed one device. Add pipelining when sequential layer groups can be balanced. Use checkpointing when activation memory is the immediate constraint, and mixed precision when the hardware supports efficient lower-precision arithmetic.

These techniques can be combined, but every added dimension of parallelism creates another communication and scheduling surface.

Common Mistakes #

Calling data parallelism model parallelism because several model replicas exist.
Forgetting that global batch size grows with worker count.
Assuming pipeline speedup equals the number of stages without accounting for bubbles.
Treating checkpointing as free memory reduction; it adds recomputation.
Assuming every training value can safely use FP16.
Applying loss scaling without unscaling gradients before the update.

Practice Questions #

Sixteen workers use local batch size 24. Find global batch size.
Explain how an all-reduce keeps data-parallel replicas consistent.
Compare tensor parallelism with pipeline parallelism.
For s = 6, m = 18, and t = 5 ms, find ideal forward pipeline time and utilisation.
A model has 60 layers with 12 MB of activation per layer. Compare full storage with 10 checkpoints.
Why are master weights often retained in FP32?
A job takes 900 seconds on one GPU and 140 seconds on eight GPUs. Calculate speedup, efficiency, and parallel overhead.
What trade-off is controlled by the number of pipeline micro-batches?

Key Takeaways #

Data parallelism partitions examples and synchronises gradients.
Model parallelism partitions parameters or tensor operations.
Pipeline parallelism overlaps micro-batches across layer stages but creates bubbles.
Gradient checkpointing exchanges activation memory for recomputation.
Mixed precision reduces memory and may increase throughput, but numerical range must be managed.
Speedup alone is incomplete; efficiency and parallel overhead show resource cost.

Checklist #

I can calculate global batch size.
I can compare data, tensor, and pipeline parallelism.
I can calculate ideal pipeline time and utilisation.
I can quantify checkpoint memory savings.
I can explain mixed precision and loss scaling.
I can calculate speedup, efficiency, and parallel overhead.

Home | ML System Optimisation