Compute-Optimal Training in Practice — isoFLOP Frontiers and Budget Allocation

From a frontier to a spending plan

The derivation gives you a shape: loss falls as a power law in both parameters N and training tokens D, and for a fixed compute budget C there is one split of that budget between the two that minimizes loss. What the derivation does not give you is the numbers for your setup — the exponents and coefficients depend on your architecture, tokenizer, data mix, and optimizer, so they have to be measured, not assumed.

That reframes the exercise as a two-step loop: first, run a small sweep to locate the frontier for your stack; second, extrapolate along it to the budget you can afford and read off the target N and D. Everything downstream — learning rate, batch size, schedule — exists to make the real run land on the point the frontier promised.

Reading an isoFLOP frontier

The empirical tool is the isoFLOP sweep. Fix a compute budget — say a few times 10^19 FLOPs — and train several models of different sizes that each consume roughly that same budget: a smaller model sees more tokens, a larger model sees fewer, because D ≈ C / (6N) holds them on the same FLOP contour. Plot final loss against N and you get a shallow U: too-small models underfit, too-large models run out of tokens, and the bottom is the compute-optimal size for that budget.

Two things matter when you read this curve. First, the bottom is flat — loss barely moves across a wide band of sizes, so the optimum is a comfortable region, not a knife-edge; being 30% off in N costs almost nothing. Second, you need runs that actually bracket the minimum on both sides; a monotone set of points has not found a bottom, it has only found a slope. Repeat the sweep at two or three budgets and the minima trace the frontier you extrapolate along.

Allocating a fixed compute budget

With a frontier in hand, allocation is mechanical. Compute-optimal scaling says the optimal N and D each grow like a power of C — in the balanced Chinchilla regime both grow like C^0.5, so they scale in lockstep and the ratio D/N stays roughly constant across budgets. That constant is the familiar tokens-per-parameter number (around 20 in the original fit, but measure your own).

The recipe: (1) take your budget C in FLOPs; (2) read the optimal N* off the extrapolated frontier; (3) get tokens from the identity, D* = C / (6 N*); (4) check that D*/N* matches your measured tokens-per-parameter ratio. If not, an input is off. Then convert to wall-clock: FLOPs divided by your hardware’s sustained throughput (peak FLOP/s times realized utilization, often 35–50%) gives the GPU-hours finance actually cares about.

A worked budget allocation

Suppose the budget is C = 1.2 × 10^21 FLOPs and your sweep’s frontier says the optimum for that budget sits near N* = 2.2 billion parameters. Tokens follow from the identity, no derivation needed:

D* = C / (6 N*)
   = 1.2e21 / (6 × 2.2e9)
   = 1.2e21 / 1.32e10
   ≈ 9.1e10 tokens  (~91 billion)

check:  D*/N* = 9.1e10 / 2.2e9 ≈ 41 tokens/param

Forty-one tokens per parameter is on the token-heavy side of the classic ratio — fine if your frontier put the minimum there, a red flag if you expected ~20 and mis-read N*. Now cost it: at 4 × 10^14 sustained FLOP/s per accelerator, the run needs 1.2e21 / 4e14 ≈ 3.0 × 10^6 accelerator-seconds — about 830 accelerator-hours, under a day on 64 devices. The allocation is arithmetic once the frontier hands you N*.

Learning rate at the optimum

Hitting the frontier requires a learning rate that actually converges the model, and the peak LR is not constant across sizes: as width grows, the LR that trains well shrinks. Under standard parameterization the safe peak LR falls off roughly with model width, which is why a recipe tuned on a 100M model can silently diverge or underperform at 3B.

The clean fix is μP (maximal-update parameterization), which rescales initializations and per-layer LRs so that the optimal learning rate becomes nearly invariant to width. That turns LR selection into a transfer problem: tune the LR once on a small proxy model, then apply it at target scale — a huge saving, because you cannot afford a sweep at the full budget. Without μP, budget a modest LR search at each new scale and expect the best value to decrease as you grow. Either way, LR decides whether the run reaches the frontier or stalls above it.

Batch size at the optimum

Batch size trades statistical efficiency against parallel throughput, and the governing idea is the critical batch size set by the gradient noise scale (McCandlish et al.). Below the critical size, doubling the batch nearly halves the steps to a target loss — almost perfect scaling. Above it, you burn compute for diminishing returns.

The subtlety that trips people up is that the critical batch size grows during training: early gradients are noisy so a large batch is wasteful, while late gradients are cleaner so larger batches pay off. The compute-optimal move is to sit near the critical batch size and optionally ramp it upward over the run. A giant fixed batch ‘for throughput’ quietly pushes you off the frontier: you spend the budget in fewer, less efficient steps and land at a higher loss than promised.

Scheduling to the token budget

Compute-optimal training assumes the model actually consumes D* tokens under a schedule that lets it converge, and the learning-rate schedule has to be matched to that exact token count. The standard shape — a short linear warmup followed by a cosine decay to a small final value — is tuned to the total number of steps.

The classic mistake is decaying on the wrong horizon. If the cosine is sized to bottom out at D* tokens but you stop early, the LR never fully anneals and the final loss is worse than predicted; the decay endpoint should coincide with the token budget you actually spend. Warmup, by contrast, is short and mostly about stability — long enough to clear the high-variance early steps at your batch size, then out of the way. Get the schedule wrong and a perfectly good allocation still misses the optimum.

When you cannot reach D-star tokens

Compute-optimal allocation quietly assumes you have D* unique tokens. Often you do not: a specialized domain, a single language, or a small institution can hit a hard ceiling of high-quality data well below what the budget wants to spend. This is the data-constrained regime, and the recipe changes.

The two levers are repetition and reallocation. Repeating data (training for multiple epochs) works better than the ‘never repeat data’ folklore suggests — up to roughly a handful of epochs, repeated tokens are nearly as useful as fresh ones, after which returns fall off sharply and eventually the model just memorizes. When unique data is the true bottleneck, the compute-optimal response is to spend the surplus budget on a smaller model trained for more epochs, or on better data curation, rather than a larger model that will starve. For CPU-served small language models this regime is the common case: you are usually data-bound long before you are FLOP-bound.

Overtraining on purpose

Compute-optimal minimizes training loss for a fixed training budget — and that is often the wrong objective. If a model will serve billions of inference requests, the training FLOPs are a one-time cost amortized over an enormous inference bill, so it pays to train a smaller model on more tokens than Chinchilla-optimal. The smaller model is cheaper on every single forward pass forever, and you deliberately overspend on training data to buy that permanent discount.

This is exactly the Llama-family logic: their smaller models were trained far past the compute-optimal token count because the deployment math dominates. The frontier still tells you the training-optimal point; you then choose to move off it, toward smaller-and-longer, by how much your inference volume justifies. The correct objective is total lifetime cost — training plus inference — and the compute-optimal point is only the right answer when you will not serve the model at scale.

How the recipe gets misapplied

Most failures are not math errors; they are category errors in applying the recipe. The frequent ones:

The through-line: compute-optimal is a measured, local, training-only statement — true for the stack you measured, near the budget you measured, for the objective of training loss. Carry it outside those bounds without re-checking and it stops being optimal.

A practical checklist

Compressed to a workflow you can run:

1. Survey. Do an isoFLOP sweep at two or three small budgets; confirm each U-curve brackets its minimum. 2. Fit. Extract your exponents and tokens-per-parameter ratio — do not borrow them. 3. Allocate. Extrapolate to your real budget, read off N*, get D* = C/(6N*), and cross-check the ratio. 4. Adjust for deployment. If the model will be served heavily, shift toward smaller-and-longer. 5. Check data. If D* exceeds your unique tokens, plan epochs or shrink the model. 6. Tune the run. Set LR by μP transfer or a small sweep (peak LR falls with size), sit near the critical batch size, and match the cosine decay to the exact token budget.

Do these six and the expensive run lands where the cheap survey said it would — which is the whole point of compute-optimal training: spend a little compute to learn the frontier so you do not waste a lot missing it.