Parallelizing Transformer Training

The core goal of model scaling is achieving strong scaling, where throughput increases linearly as the number of chips grows. While single-chip performance depends on the balance between memory bandwidth and FLOPs, at the cluster level, it is crucial to overlap inter-chip communication with useful computation to hide latency.

Increasing the number of chips increases communication overhead while reducing the computation per device available to hide it, creating a non-trivial problem. In particular, sharded matrix multiplications incur high-cost communications such as AllGather and ReduceScatter, which can block effective work on TPUs.

This section covers five parallelization strategies: data parallelism, FSDP, tensor parallelism, expert parallelism, and pipeline parallelism, analyzing the communication cost limits at which each strategy turns computation costs into bottlenecks. Since memory capacity constraints are largely ignored when using rematerialization, the focus remains on inter-chip communication costs.