Parallel training
Describe a pretraining loop and reason about data, batch size, precision and a simplified compute budget.
THE IDEA
Using many devices
Data parallelism gives each device a copy of the model and a share of the batch; their gradients are averaged ("all-reduce") before the update. With equal shares the result equals the single-device gradient. Tensor parallelism splits individual matrices across devices and pipeline parallelism splits the layers into stages that work on different microbatches, leaving some idle time (the bubble) when the pipeline fills and drains.
Data parallelism: one pass takes 80 s on one device. Each device handles a share, then gradients are averaged across devices (communication).
Pipeline parallelism cuts the model into stages. At the start and end some stages sit idle (the bubble):
Predict first. Is the ideal arithmetic division by GPU count a complete wall-clock model?
Chapter context
Pretraining optimizes next-token predictions over a dataset. Its costs depend on the architecture, sequence lengths, implementation and hardware. A useful budget model states which work it approximates and which work it ignores.
TRY IT YOURSELF
Part 1. Do it hands-on
Compute a real speedup and a pipeline bubble by hand, then show that data-parallel gradients match the single-device gradient.
Data parallelism: one pass takes 80 s on one device. Each device handles a share, then gradients are averaged across devices (communication).
Pipeline parallelism cuts the model into stages. At the start and end some stages sit idle (the bubble):
- To do: Data parallel on 4 devices; each processes batch 4 × length 8 and accumulates 2 steps. Token positions per optimizer update across all devices?
- To do: One device needs 80 s per pass. On 8 devices the compute takes 10 s but communication adds 2 s. Real speedup = 80 / 12 (3 decimals)
- To do: Pipeline parallelism with p = 4 stages and m = 8 microbatches wastes a fraction (p−1)/(m+p−1) of the time idle (the bubble). Type it (3 decimals)
Part 2. Do it in code
- global_tokens(per_device_batch, seq, accum, devices): the token positions per optimizer update across all devices.
- allreduce_mean(grads): grads is a list of equal-shaped arrays, one per device. Return their elementwise mean.
- data_parallel_grad(w, X, y, devices): split the rows evenly across devices (ValueError if not divisible), compute each device's mean-squared-error gradient and return allreduce_mean of them.
- speedup(t_single, t_compute, t_comm): t_single / (t_compute + t_comm).
- pipeline_bubble(stages, microbatches): (p − 1) / (m + p − 1).
Stuck, or something not working? You can come back to it any time.
CHECK YOUR UNDERSTANDING
Is the ideal arithmetic division by GPU count a complete wall-clock model?
Chapter reference: math, worked example & code
This section belongs to the whole chapter, so it is the same on every lesson in it and may cover a different lesson's topic. The code for the experiment you just ran is inside the experiment itself.
Math and notation
A commonly used rough dense-transformer training estimate is
with N model parameters, D training tokens and C floating-point operations. It omits some architecture-dependent work and is not a hardware throughput guarantee. Effective token batch is approximately for microbatch B, length T and A accumulation steps.
Worked example
For N=1 million parameters and D=100 million training tokens, the rough compute estimate is 6×10^6×10^8=6×10^14 FLOPs. A microbatch of 4 sequences of length 8 accumulated twice contributes 64 token positions per optimizer update. Neither calculation includes padding exclusions or communication time.
Reference code
The following example uses the Python standard library and does not call a remote service.
N,D=1_000_000,100_000_000
assert 6*N*D == 600_000_000_000_000
batch,length,accumulation=4,8,2
assert batch*length*accumulation == 64
Python examples are provided for study and local execution; this page does not execute Python.
Common mistakes
Compute-optimal ratios are model- and study-dependent. Train/test leakage can arise from shared documents and overlap. Tokens per update and optimizer steps are different units.