How much will a training run cost? C = 6ND

Build a Small GPT From Scratch · Part 6 of 8

Before you spend a GPU-hour, you can estimate what a training run will cost with one line of arithmetic. It is rough, but it is accurate enough to stop you from planning something impossible.

The rule of thumb

Training a dense transformer costs about 6 floating-point operations per parameter per training token: roughly 2 for the forward pass and 4 for the backward pass. So for N parameters and D tokens, the total compute is about C = 6 × N × D.

Turning FLOPs into days

FLOPs mean little until you convert them to time on a real GPU. An A100 peaks at about 312 trillion operations per second in bfloat16, and a good run reaches perhaps 40% of that. Dividing the total by 40% of peak, and by the seconds in a day, gives A100-days. Using the 20-tokens-per-parameter ratio that the next post explains:

def train_flops(N, D):
    return 6 * N * D

def a100_days(flops, mfu=0.4, peak=312e12):          # bf16 peak; 40% utilisation is a good run
    return flops / (mfu * peak) / 86400

for name, N in (("tiny GPT", 10e6), ("1B", 1e9), ("7B", 7e9), ("70B", 70e9)):
    D = 20 * N
    print(name, "A100-days:", round(a100_days(train_flops(N, D)), 3))
# tiny GPT A100-days: 0.001
# 1B A100-days: 11.129
# 7B A100-days: 545.317
# 70B A100-days: 54531.695

What the table teaches

At a fixed number of tokens per parameter, cost grows with the square of model size, because parameters and tokens both grow. A 70B run costs a hundred times a 7B run, not ten times. And the whole of Lab 3 is about one five-hundred-thousandth of a 7B run, which is what “Colab-scale” honestly means. These are estimates: real runs differ in architecture, utilisation and wasted work.

Open in Colab Change the model size and the tokens-per-parameter ratio and see the bill.

Takeaways

• Training compute is about 6 × parameters × tokens.
• Dividing by 40% of an A100’s peak speed turns FLOPs into GPU-days.
• At a fixed tokens-per-parameter ratio, 10 times the parameters costs 100 times the compute.
• The tiny GPT of Lab 3 is about 1/500,000 of a 7B run.

Next in the series (Part 7): Chinchilla, and why recent models are over-trained on purpose