Parameter & FLOP Counting

The problem, from first principles

Before training a model you need two numbers: how many parameters it has and how much compute a run will cost. A decoder-only transformer's parameters are dominated by two kinds of weights. Each layer has four d x d attention matrices (query, key, value, output) and a two-matrix MLP of width d_ff. The embedding table has vocab x d weights and is counted once if the output head shares it. Training compute then follows a famous rule of thumb: each parameter takes part in about 6 floating-point operations per training token (2 for the forward pass, 4 for the backward pass), so total FLOPs are about 6 * N * D for N parameters and D tokens.

From theory to code

Implement transformer_params, training_flops and chinchilla_tokens.

Constraints

• transformer_params(n_layers, d_model, vocab, d_ff=None, tied=True) counts: per layer 4 * d_model**2 + 2 * d_model * d_ff, plus the embedding vocab * d_model once if tied else twice. d_ff defaults to 4 * d_model. Ignore biases and normalization weights.
• training_flops(n_params, n_tokens) is 6 * n_params * n_tokens.
• chinchilla_tokens(n_params) is the compute-optimal token count 20 * n_params.
• Return Python ints for the first function and floats or ints from the arithmetic for the others.

Hints