Compute Needs of ML Training
Training large models is an HPC problem: it demands data, model, and pipeline parallelism, fast collectives, and careful precision and memory management.
Training as an HPC workload
Training a large neural network is dominated by dense matrix multiplication repeated over billions of examples. This makes it a data-parallel, accelerator-bound workload with the same concerns as scientific HPC: parallel decomposition, communication cost, memory limits, and numerical precision.
Three axes of parallelism
- Data parallel: replicate the model, split the batch, average gradients with all-reduce
- Tensor (model) parallel: split individual layers across devices
- Pipeline parallel: place successive layers on successive devices
Communication is the bottleneck
Data-parallel training synchronizes gradients every step with an all-reduce, whose cost grows with model size and device count. The interconnect often limits scaling, which is why large training clusters use high-bandwidth fabrics and overlap gradient communication with backpropagation.
Memory and precision
Model weights, activations, gradients, and optimizer states must fit in accelerator memory, which is the binding constraint for the largest models. Mixed precision on tensor cores, activation checkpointing (recompute instead of store), and sharding optimizer state across devices are the standard techniques for fitting and speeding up training.
The same lessons
Large-scale training rediscovered the lessons of scientific HPC: expose massive parallelism, minimize and overlap communication, respect the memory hierarchy, and choose precision deliberately. The two fields now share hardware, tools, and techniques, and increasingly the same machines. Frameworks hide much of this behind high-level interfaces, but understanding what happens underneath is what separates training that scales from training that stalls.