Design Distributed Training for a Large LLM
Distributed LLM training means training a large language model across many GPUs, often thousands, that coordinate their work and update the same model. It combines several forms of parallelism to fit the model into memory and complete training in weeks rather than years.
Popular Examples: Megatron-LM, DeepSpeed, PyTorch FSDP
"Design distributed training for a large LLM" can come up in ML engineer and research engineer interviews. Although it involves training a model, much of the discussion focuses on systems: memory capacity, network bandwidth, hardware failures, storage throughput, and cost.
The basic problem is straightforward. A 70-billion-parameter model and its training state can require more than a terabyte of memory, while a single H100 GPU has 80 GB. Even if everything fit, one GPU would need hundreds of years to process the training data. We therefore need thousands of GPUs working together for weeks, on a cluster where hardware failures are a normal part of operation.
In this chapter, we'll design the system that runs this training job. We'll start with the requirements and high-level architecture, then examine how to divide the work across GPUs, recover from failures, manage checkpoints, and estimate training time and cost.