Training Resource Budgeting: FLOPs, Memory, and Wall-Clock

The Problem

You are in a planning meeting. Someone asks: "Can we train a 7B model on 100 billion tokens in a month with 64 H100s?" That question is not a hunch, and it is not a search through someone else's benchmark spreadsheet. It decomposes into three quantities:

1. FLOPs. How many floating-point operations does the run need?
2. Memory. Will the model state and activations fit in each GPU, and which parallelism or checkpointing levers change that?
3. Wall-clock time. At a realistic Model FLOPs Utilization (MFU), how many hours or days will the cluster run?

If any of the three are wrong, the project slips. FLOPs too low and you under-provision GPUs; memory too low and you hit OOM at step zero; wall-clock too optimistic and you miss the ship date. This lesson builds the pocket calculator every engineer on the call should be able to reproduce.

The Concept

The 6ND FLOPs rule

For a dense transformer with N parameters trained on D tokens, the total training FLOPs are approximately:

The factor 6 comes from splitting one forward-backward pass per token.

Forward pass. One token touches every parameter once in a matrix multiply or embedding lookup. Each multiply-accumulate is 2 FLOPs, so forward costs 2N FLOPs per token.

Backward pass. Backpropagation computes two things: the gradient flowing backward through each operation, and the gradient with respect to each weight. For a typical dense layer this is roughly twice the work of the forward pass, or 4N FLOPs per token.

The backward pass is therefore twice as expensive as the forward pass. This rule of thumb is the 6ND estimate used in the PaLM training report and in many subsequent capacity plans.

Memory line-item budget

Mixed-precision training with Adam keeps the following per parameter:

Activations add another term that depends on batch size B, sequence length S, hidden dimension H, number of layers L, and precision bytes b:

For a first-order estimate, activation memory for a transformer block without gradient checkpointing is often in the range 12 \, BSH to 24 \, BSH bytes per layer. Gradient checkpointing trades extra forward recomputation for memory by keeping only the inputs to each transformer block and recomputing the rest during backward, reducing the stored activation footprint by roughly a factor of L.

ZeRO and gradient checkpointing: what each lever removes

ZeRO-1 shards the optimizer states (m and v) across data-parallel ranks. Each GPU stores 8N / \text{world\_size} bytes of optimizer state instead of 8N.

ZeRO-2 additionally shards gradients. Each GPU stores (8N + 2N) / \text{world\_size} bytes for optimizer plus gradients.

ZeRO-3 additionally shards the model weights. Each GPU stores (8N + 2N + 2N) / \text{world\_size} bytes, leaving only a slice of weights, gradients, and optimizer states on each rank.

Gradient checkpointing does not shard state; it reduces activation memory by recomputing activations during the backward pass. The cost is an extra forward pass, so total FLOPs increase by roughly 33%.

MFU and wall-clock time

Peak hardware throughput is not the throughput you get. Model FLOPs Utilization captures the gap:

An NVIDIA H100 SXM5 delivers roughly 989 TFLOP/s in dense FP16/BF16 tensor-core mode. Large transformer training runs typically observe MFU in the 0.3–0.5 range because of data-loading, communication, kernel launch, and pipeline bubbles.

Given G GPUs with peak P and MFU \eta, the wall-clock time in seconds is:

If you know the cluster rental price per GPU-hour, the total cost is:

Chinchilla optimum

Hoffmann et al. (2021), the Chinchilla paper, showed that for a given compute budget C = 6ND, the loss is minimized when the number of tokens is about twenty times the number of parameters:

Substituting into the compute budget:

Solving for N and D given C:

For a 7B parameter model, the Chinchilla-optimal dataset size is approximately 20 \times 7\,\text{B} = 140\,\text{B} tokens.

Later work and production models have deliberately over-trained: using 2\times to 5\times the Chinchilla-optimal token count improves inference quality and downstream task performance even though the extra training is sub-optimal for loss-versus-compute. The choice depends on whether the model will be served at high volume, where inference quality dominates lifetime cost.

Build It

The lesson ships a command-line training-budget calculator. It takes a model size, token count, GPU type, GPU count, MFU, precision, and parallelism level, then prints:

• total training FLOPs,
• a per-parameter memory line item table,
• estimated wall-clock time and GPU-hour cost,
• an activation-memory estimate.

Run it:

consoleCopypython3 main.py

The calculator also asserts three sanity checks used throughout the lesson: the 7B/100B-token FLOPs estimate, the 70B-parameter Adam model-state footprint, and the internal consistency of the Chinchilla formulas.

Verify It

code/tests/test_main.py locks down the calculator's accounting:

• 7B parameters trained on 100B tokens produce approximately 4.2 \times 10^{21} FLOPs.
• A 70B-parameter mixed-precision Adam run allocates roughly 1.1 TB of model-state memory before activations.
• The Chinchilla optimum satisfies C = 6ND when D = 20N.
• Wall-clock time scales inversely with GPU count at fixed MFU.
• ZeRO-1/2/3 each remove the advertised memory component.
• GPU peak-FLOP lookup returns the expected values for H100, A100, and RTX 4090.

Why This Matters for Engineers

• Procurement starts with arithmetic. "Can we train this in a month?" is answered by 6ND / (P \cdot \eta \cdot G), not by hope. Bring the formula to the meeting and the conversation becomes about MFU assumptions, not about gut feeling.
• OOM is a budget bug, not a surprise. If the memory line items do not fit in GPU RAM, you decide between ZeRO stages, gradient checkpointing, or more GPUs before a single job is submitted.
• MFU is the real capacity. A cluster on paper is peak TFLOP/s times GPU count; the cluster you actually get is that multiplied by MFU. 0.3 versus 0.5 is a 67% difference in wall-clock time.
• Chinchilla is a starting point, not a law. 7B/140B is compute-optimal, but high-volume inference models may be trained longer. The calculator lets you compare the compute-optimal plan against an over-trained plan side by side.
• Inference teams need the same skills. KV-cache memory and batch-size throughput are the inference analogs of optimizer-state memory and MFU. The habit of building a line-item budget transfers directly.

Hints

• The 16 bytes per parameter rule already assumes mixed precision plus Adam master weights. If you use 8-bit optimizers or ZeRO sharding, scale the relevant line item down.
• Activation memory is the term most likely to surprise you; sequence length enters linearly and batch size enters linearly, so a 2x longer context is a 2x memory hit unless gradient checkpointing absorbs it.
• When MFU is unknown, use 0.35 as a conservative first guess for dense transformer training on a well-tuned cluster.
• Compare Chinchilla-optimal against over-trained by fixing wall-clock time and asking which model size and token count fit the budget.

Key Terms