The Arithmetic of OLMo 3's Training Run

The problem

The OLMo 3 report gives enough numbers to check each other: if they are consistent, you understand how a training run is planned. For the pre-training of the 32B model:

• 5.5T tokens, in batches of 8,388,608 tokens (1,024 sequences of 8,192);
• a measured throughput of 1,960 tokens per second per GPU (graphics processing unit), on up to 1,024 H100 GPUs;
• 32,233,522,176 parameters;
• a real duration of about 9.5 days on 512 GPUs followed by 35 days on 1,024 GPUs, restarts included;
• about 56 days on 1,024 GPUs from the start of training to the evaluated 32B Think model, priced at $2 per H100-hour.

You will write the five small functions of the arithmetic.

The five calculations

1. Steps. The number of optimizer steps is the tokens divided by the tokens per batch, rounded up (a partial last batch is still a step).
2. Compute. The training FLOPs (floating-point operations) are $6ND$ (section 8).
3. Ideal days. If the cluster ran at the measured throughput the whole time: $\text{tokens} / (\text{tokens per second per GPU} \times \text{GPUs})$ seconds, divided by 86,400.
4. GPU days. A run can change size: a list of (days, n_gpus) segments adds up to $\sum \text{days} \times n_{\text{gpus}}$ GPU days.
5. Cost. GPU days times 24 hours times the price of one GPU hour.

Example. With the numbers above:

1. $5.5 \times 10^{12} / 8{,}388{,}608 = 655{,}651.09$, so $655{,}652$ steps.
2. $6 \cdot 32{,}233{,}522{,}176 \cdot 5.5 \times 10^{12} \approx 1.06 \times 10^{24}$ FLOPs.
3. $1960 \cdot 1024 = 2{,}007{,}040$ tokens per second for the cluster, so $5.5 \times 10^{12} / 2{,}007{,}040 \approx 2.74 \times 10^6$ seconds: $31.7$ days. The real run took $9.5 + 35 = 44.5$ days: the difference is crashes, restarts and the first 9.5 days on half the GPUs.
4. [(9.5, 512), (35, 1024)] is $4864 + 35840 = 40{,}704$ GPU days.
5. 56 * 1024 GPU days at $2 per hour: 57,344 times 24 times 2 is $2,752,512, the $2.75M of the report.

What to implement

• n_steps(tokens, batch_tokens): optimizer steps, rounded up, an int.
• train_flops(n_params, tokens): $6ND$.
• ideal_days(tokens, tokens_per_sec_per_gpu, n_gpus): days at the measured throughput.
• gpu_days(segments): total GPU days of a list of (days, n_gpus).
• cost_usd(gpu_days, price_per_gpu_hour): dollars.

Requirements

• Return Python numbers (int for n_steps).
• Use math.ceil for the steps.