Scaling Pre-Training in Practice: A Hierarchical Approach

gist

Jordan Sassoon of Aleph Alpha describes scaling pre-training of a 30B-A3B mixture-of-experts model to 512 Nvidia B200 GPUs at 35.3% MFU. MFU, or model FLOPS utilisation, expresses throughput as a fraction of the hardware's peak. The team swept parallelism, activation checkpointing, and local batch size on 16 GPUs, then used PyTorch profiles to confirm compute kernels ran back to back. Scaling to 128 GPUs with full sharding and then four data-parallel replicas cut per-GPU throughput only 6%, from 28.4k to 26.7k tokens per second. System- and kernel-level optimisations are out of scope.

ideas

• Tokens per second per GPU is the objective, and MFU is the indicator. Sassoon says the per-token FLOP count that MFU needs is hard to compute, and every MFU figure uses BF16 peak FLOPS as the denominator.
• A 16-GPU sweep ruled out most configurations. No checkpointing ran out of memory too early, selective checkpointing peaked at a local batch size of about 18, and full checkpointing at about 38.
• Profiles show whether training stays compute-bound. Traces showed weight all_gather prefetches overlapping compute, and skipping the data-parallel all_reduce on non-final gradient accumulation steps removed that communication.
• Fully sharded data parallelism grows until communication nearly catches compute, then data parallelism takes over. At 128 GPUs the forward pass sat at the communication-bound boundary, so the team kept FSDP at 128 and added four data-parallel replicas to reach 512 GPUs.
• The final setup could train the model on 20T tokens in 17 days. Sassoon says two more replicas would cut that to 11 days, and the results cover one architecture, one sequence length, and FSDP and DP only.

quotes