Planning a training run

Planning a run means picking a layout that fits in memory, then checking that it trains efficiently. This page works through a 70B model on a 64-GPU cluster. It covers the two tools that fit activations into memory, checkpointing and gradient accumulation. It ends with MFU, the number that tells you whether the layout is any good.

A worked layout

Take 8 nodes with 8 H100 GPUs each, which is 64 GPUs. The model is 70B parameters, trained in bf16 with Adam.

• Parameters alone at 2 bytes each are about 140 GB.
• With the full training state at about 16 bytes per parameter, the total is about 70e9 × 16 = 1.1 TB, before activations.

One GPU cannot hold that, and one node is tight even with sharding. A common layout is:

TP = 8   inside each node, for the frequent all-reduces within layers
PP = 4   across nodes, for depth
DP = 2   for the remaining replicas

8 × 4 × 2 = 64 GPUs

Sharding the optimizer state with FSDP or ZeRO can be added on top of this layout, depending on the stack.

Fitting activations

Two tools reduce activation memory and keep the batch size you want.

Activation checkpointing stores only some activations, usually at block boundaries, and recomputes the rest during the backward pass. The cost is roughly 30% more compute. In return, the microbatch can grow, which often makes the difference between a batch of 1 and a batch of 4.

Gradient accumulation runs several forward and backward passes, adds up the gradients, and steps the optimizer once. The effective batch size is:

effective batch = microbatch × accumulation steps × DP

MFU: the number that matters

Model FLOPs Utilization (MFU) is the fraction of the hardware's peak compute the model actually uses:

MFU = achieved model FLOPs per second ÷ peak GPU FLOPs per second

For a dense transformer, the useful work per training step is often estimated as about 6 × parameters × tokens. The estimate ignores kernels that do not contribute to the model's math, so treat it as a yardstick, not an exact count.

Well-tuned LLM pretraining often lands around 35 to 55% MFU on mature stacks. Below about 20%, something is wrong. The usual causes are:

• Tensor parallelism spanning slow links.
• A pipeline bubble from too few microbatches.
• A data loader that does not prefetch, so the GPUs wait for input.
• Stragglers, where one slow rank stalls every step.
• Gather traffic from sharding competing with tensor-parallel collectives on the same network cards.

Measure MFU after a few hundred steady steps. Warmup steps are slower and do not represent the run.

Common mistakes

• Declaring success because the job runs. A layout can run at 20% MFU. Measure it.
• Changing tensor, pipeline and data parallelism together. Change one axis at a time, so you know what helped.
• Measuring during warmup. The first steps are not representative.

Check yourself

1. Why measure MFU after a few hundred steady steps? Warmup steps are slower and not representative of the run.
2. Microbatch 2, accumulation 16, DP = 2. What is the effective batch? 2 × 16 × 2 = 64 sequences.

Key takeaways

• A 70B model needs about 1.1 TB of training state before activations. It needs a layout across all 64 GPUs.
• Tensor parallelism inside nodes, pipeline across nodes and data parallelism for replicas is a common starting point.
• Checkpointing trades about 30% compute for activation memory. Accumulation gives a large effective batch.
• MFU of 35 to 55% is typical for well-tuned pretraining. Much lower means a layout or input problem.

Further learning

The scaling laws page covers how much data and compute a run needs, which decides how long a run like this should take.