Home / Fundamentals / Distributed Training

Memory and compute arithmetic for training

Where training memory goes (16 bytes per parameter plus activations), activation recomputation, FLOPs per token, MFU, and how to estimate the wall-clock time of a run.

13 min read

Before choosing a parallelism strategy, you need two numbers: how much memory the job needs and how many FLOPs it must perform. Everything else follows from those.

Model and optimizer state

With mixed-precision AdamW, each parameter costs:

So a 7B model needs ~112 GB of state and a 70B model ~1.1 TB, before a single activation. An 80 GB GPU can't even hold a 7B model's training state, which is why state is sharded.

Activations

The backward pass needs intermediate values from the forward pass. For one transformer layer with sequence length s, micro-batch b and hidden size h, a standard estimate (bf16, fused/flash attention so the s × s score matrix isn't stored) is:

activation bytes per layer ≈ 34 · s · b · h

Example: a 7B model (h = 4096, 32 layers) at s = 4096, b = 1:

34 × 4096 × 4096 ≈ 570 MB per layer  →  ≈ 18 GB for 32 layers

Activations grow linearly with batch and sequence length, so they often dominate memory for long contexts.

Activation recomputation (checkpointing) stores only each layer's input (2·s·b·h bytes, ~34 MB here) and recomputes the rest during backward. Cost: one extra forward pass, roughly +33% compute (the forward is about 1/3 of the forward + backward FLOPs). Selective recomputation recomputes only cheap-to-recompute but memory-heavy parts (such as attention internals) and recovers most of the memory for a few percent of compute.

Compute

Training costs about 6N FLOPs per token for a dense model with N parameters: 2N forward and 4N backward. Attention adds ≈ 12 · L · h · s FLOPs per token (forward plus backward), which matters for long sequences.

Model FLOPs utilization (MFU) = useful model FLOPs per second / hardware peak. Recomputation doesn't count as useful. Well-tuned large dense runs reach roughly 35–50% MFU on H100-class GPUs.

Estimating wall-clock time

time = 6 · N · D / (GPUs · peak FLOP/s · MFU)

Example: 70B parameters on 1.4T tokens, 2,048 H100s (989 TFLOP/s dense bf16), 40% MFU:

6 × 70e9 × 1.4e12 = 5.9e23 FLOPs
2,048 × 989e12 × 0.4 = 8.1e17 FLOP/s
5.9e23 / 8.1e17 ≈ 7.3e5 s ≈ 8.4 days

Real runs add time for failures, restarts, evaluations and data stalls, so budget extra (often 10–30%).

What to remember

• 16 bytes per parameter of state; activations ≈ 34·s·b·h per layer.
• Recomputation trades ~33% compute for most of the activation memory.
• 6ND FLOPs; MFU of 40% is a good large-scale result.