Vol. XXXI · History and Projects Plate 8 of 14project, training plan, compute budget, Chinchilla, tensor parallelism, pipeline parallelism, ZeRO, memoryCA

Project: Plan a Training Run.

A cluster, a month and a loss target: pick the model, the tokens and the layout before spending a GPU-hour.

The brief: 2 048 H100s (989 TFLOP/s dense BF16, 80 GB each) for at most 30 days at an assumed 40% model FLOPs utilisation. Choose one of the Llama 3 shapes (8B, 70B or 405B), a number of training tokens, and a layout — tensor and pipeline parallel degrees and a ZeRO stage, with the rest data parallel — so that each GPU needs at most 72 GB, the run takes at most 30 days by 6ND (plate V.6), and the loss predicted by the Besiroglu et al. refit of the Chinchilla law (plate X.3) is at most 1.955.

The starting plan, 70B on 15 trillion tokens with no parallelism, needs 1,140 GB per GPU and 91 days. A project plate grades the design against a brief: every check below the workshop must pass, and a passing design is saved in this browser.

Left: per-GPU memory by component against the 72 GB line. Right: the predicted loss against the target and the days needed against the month.

What you will learn
Workshop
Memory per GPU from plate XIV.1’s model: BF16 weights and gradients, FP32 Adam state and master weights sharded by ZeRO stage, activations with full recomputation and FlashAttention for 8 192-token sequences, one sequence per GPU. Loss from Besiroglu et al. (2024): E = 1.82, A = 482, B = 2,085, α = 0.348, β = 0.366. MFU 40% is an assumption.

Three constraints, solved in order

The budget comes first: the cluster delivers about 2.10 × 10²⁴ FLOPs in 30 days at 40% utilisation, and 6ND fixes the products of size and tokens it allows. The loss target then rules out the 8B model — however long it trains within the month, its predicted loss stays above the target — and the 405B model, which cannot see enough tokens in a month. Under this fit the compute-optimal model for the budget would have about 140 billion parameters; among the published shapes, 70B with a few trillion tokens is the one that fits.

Memory comes last and is solved by layout, not by model choice: without sharding, a 70B model’s weights, gradients and Adam state need 1,140 GB on every GPU. Tensor parallelism splits each layer’s matrices across a node’s eight GPUs, ZeRO shards the optimiser state across the data-parallel replicas, and recomputation keeps activations small — the levers of Vol. XIV. A real plan adds what this one leaves out: communication time, failures and restarts (plate XV.1), and an MFU that has to be measured, not assumed.

Llama 3 405B was trained on about 15.6 trillion tokens with 16 384 H100s; its 8B sibling saw the same tokens, far past compute-optimal, because it was meant to be cheap to serve.

Terms

FLOP budget
GPUs × peak × utilisation × time.
Compute-optimal
The size and token count minimising predicted loss for a budget.
Tensor parallelism
Splitting each layer’s matrices across GPUs.
ZeRO
Sharding optimiser state, gradients or weights across data-parallel replicas.
A training plan solves budget, loss and memory together: the scaling law picks size and tokens, and the parallel layout makes them fit.

Sources

  1. Hoffmann et al., “Training Compute-Optimal Large Language Models”, NeurIPS 2022.
  2. Besiroglu, Erdil, Barnett and You, “Chinchilla Scaling: A replication attempt”, arXiv:2404.10102, 2024.
  3. Rajbhandari et al., “ZeRO: Memory Optimizations Toward Training Trillion Parameter Models”, SC 2020.
  4. Narayanan et al., “Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM”, SC 2021.
  5. Llama Team, “The Llama 3 Herd of Models”, arXiv:2407.21783, 2024.