Training at scale: GPUs, parallelism and memory

In short: Large models do not fit on one GPU. They are spread across thousands of GPUs with several kinds of parallelism, fast links and mixed precision.

In more detail: In training each weight needs, besides its value, a gradient and optimiser states. With Adam in mixed precision that is about 16 bytes per parameter. A model with 70 billion parameters then needs over 1 terabyte for this data alone.

In detail

• Data parallelism (DP): each GPU has a copy, gets different data, gradients are averaged.
• ZeRO/FSDP: distributes weights, gradients and optimiser states across the GPUs, each fetches on demand.
• Tensor parallelism: individual matrices are split across GPUs, see Multi-GPU and tensor parallelism.
• Pipeline parallelism: different layers on different GPUs, micro-batches flow through.
• Expert parallelism: with mixture of experts experts sit on different GPUs.
• Mixed precision: compute in BF16/FP8, master weights in FP32, see Number formats FP32, FP16, BF16 and FP8.
• Links: NVLink inside a server, InfiniBand or Ethernet between them. Communication is often the bottleneck.
• Failures: with thousands of GPUs something fails daily, hence checkpoints.
• Hardware: NVIDIA H100/B200, Google TPU, AMD MI300, see NPU and AI accelerators.
• Utilisation (MFU): the share of theoretical compute that is used, typically 30 to 50%.

See also: How a language model is trained, Energy consumption of AI