Compute-optimal training: Chinchilla-style scaling laws

The problem, from first principles

Training a large language model costs a roughly fixed, known amount of COMPUTE (measured in total floating-point operations, FLOPs), and that budget has to be SPLIT between two competing uses: making the model itself BIGGER (N, the parameter count) or training it on MORE data (D, the number of tokens seen). Kaplan et al. (2020, OpenAI's original scaling laws paper) initially suggested favoring bigger models; Hoffmann et al. (2022, DeepMind, "Chinchilla") revisited this far more carefully and found the earlier guidance had been systematically under-training models relative to their size. Chinchilla's headline finding: for a FIXED compute budget, loss is minimized when model size and training data are scaled roughly EQUALLY, both proportional to the SQUARE ROOT of the compute budget, not favoring one over the other. This directly motivated a real, visible shift in how the field trains models: Chinchilla itself, at 70B parameters trained on 1.4T tokens, OUTPERFORMED the much larger 280B-parameter Gopher, trained on far fewer tokens for the SAME compute budget, purely by reallocating that same budget toward more data instead of more parameters.

From theory to code

Implement flops_for_training(num_params, num_tokens), the standard 6*N*D compute approximation, and chinchilla_optimal_allocation(compute_budget_flops), returning the compute-optimal (N, D) split for a given budget: both proportional to sqrt(compute_budget_flops / 6).

Constraints

• flops_for_training(N, D) = 6 * N * D: this is a widely-used APPROXIMATION (one forward pass costs roughly 2N FLOPs per token, and a full forward-plus-backward pass costs roughly 3x that, 2N + 4N = 6N), not an exact count.
• chinchilla_optimal_allocation returns EQUAL values for N and D (Chinchilla's simplified headline finding: roughly N ≈ D at the optimum, in appropriately-scaled units), specifically sqrt(compute_budget_flops / 6) for both.
• The chosen (N, D) must satisfy flops_for_training(N, D) == compute_budget_flops EXACTLY (the whole budget gets allocated, none wasted, none exceeded).
• Doubling the compute budget scales BOTH N and D by sqrt(2) (not by 2, and not by different amounts for each).

Hints