Compute Budgets & Scaling Laws

The problem

Before a training run starts, a few lines of arithmetic tell you how much compute it needs, how long it will take, and whether the model size fits the budget. You will write four of them, with plain Python floats. All of them are built on one rule from the theory page: training a model with $N$ parameters on $D$ tokens costs

FLOPs (floating-point operations), where 2 per parameter per token are for the forward pass and 4 for the backward pass.

1. Compute of a run

$C = 6ND$, with $N$ = n_params and $D$ = n_tokens.

Example. $N = 10^9$ and $D = 2 \cdot 10^{10}$: $C = 6 \cdot 10^9 \cdot 2 \cdot 10^{10} = 1.2 \times 10^{20}$ FLOPs.

2. The best model size for a budget

Hoffmann et al. (Chinchilla) found that, for a fixed budget, about 20 training tokens per parameter gives the lowest loss: $D = 20N$. Put it into $C = 6ND$:

so

Example. $C = 10^{21}$: $N = \sqrt{10^{21} / 120} \approx 2.89 \times 10^9$ parameters, and $D = 20N \approx 5.77 \times 10^{10}$ tokens.

3. Days of training

Each of the n_gpus GPUs (graphics processing units) can do at most peak_flops FLOPs per second, and a real run uses a fraction mfu of that peak, the MFU (model FLOPs utilization), for example 0.4. So the run does n_gpus * peak_flops * mfu useful FLOPs per second, and lasts $C$ divided by that, in seconds. A day has 86,400 seconds.

Example. The run of part 1 ($1.2 \times 10^{20}$ FLOPs) on 8 GPUs with a peak of $10^{15}$ FLOPs per second and an MFU of 0.4: $8 \cdot 10^{15} \cdot 0.4 = 3.2 \times 10^{15}$ FLOPs per second, so $37{,}500$ seconds, about $0.434$ days.

4. The MFU of a running job

The other way round: from the measured speed of a run, in tokens per second for the whole cluster, compute its MFU. The model needs $6N$ FLOPs per token, so it uses $6N \times$ tokens per second; divide by the peak of the whole cluster, n_gpus * peak_flops.

Example. A model of $N = 10^9$ processing $2 \times 10^5$ tokens per second on the same 8 GPUs: $6 \cdot 10^9 \cdot 2 \cdot 10^5 = 1.2 \times 10^{15}$ FLOPs per second, out of a peak of $8 \times 10^{15}$: an MFU of $0.15$.

What to implement

• train_flops(n_params, n_tokens): part 1.
• chinchilla_optimal(compute): part 2, the tuple (n_params, n_tokens).
• training_days(n_params, n_tokens, n_gpus, peak_flops, mfu): part 3.
• model_flops_utilization(n_params, tokens_per_second, n_gpus, peak_flops): part 4.

Requirements

• Return floats (for chinchilla_optimal, a tuple of two floats).