Scaling Laws & Modern Training

TL;DR

Intuition

Training a frontier LLM is a three-way trade between compute, data, and parameters. Spend all your compute on a tiny model and you under-fit. Spend all of it on parameters and you over-fit, or run out of data. Spend it on a moderate-size model trained for a long time and you sit on the compute-optimal frontier [14]. The scaling-law literature is the empirical map of that frontier: how loss falls as you spend more, and which axis pays the most.

For most of 2018 to 2022 the consensus from Kaplan 2020 was “make the model bigger; data scales sub-linearly”. Hoffmann et al. ran a broader sweep with the right hyperparameter retuning per model size and found the opposite: data scales as fast as parameters. The whole field shifted in a year. Modern open-weights LLMs (LLaMA-2 [8], Mistral [9], Qwen [11], DeepSeek-V2 [12]) train past the Chinchilla optimum on purpose, because at inference time the cheaper, smaller, longer-trained model wins on tokens-per-dollar.

1. 2017Hestness et al.: deep-learning scaling is predictablepaper

   A Baidu paper shows that translation, language-modelling, and image-classification loss all fall as smooth power laws in dataset size, holding architecture and compute fixed. Plants the empirical seed for scaling laws.

2. 2020Kaplan 2020 scaling lawspaper

   Clean joint fit of loss versus parameters, data, and compute on dense decoder-only Transformers. Predicts compute-optimal N grows roughly as C^0.73 and D as C^0.27. Becomes the planning law for GPT-3.

3. 2022Hoffmann et al. (Chinchilla)paper

   Per-size hyperparameter retuning and a wider sweep find compute-optimal N and D scale at near-equal exponents. Chinchilla 70B beats Gopher 280B at the same compute, validating the new prescription D approx 20 N.

4. 2023LLaMA-2: deliberately past Chinchillarelease

   LLaMA-2 7B and 70B both train on 2T tokens. The 7B is more than 10x past Chinchilla; the bigger training compute is paid back in cheaper inference at deployment.

5. 2024MoE re-emerges at scalearchitecture

   Mixtral 8x7B and DeepSeek-V2 ship sparse mixture-of-experts open weights, decoupling active parameters per token from total parameter count. The training stack that supports them is the modern 3D parallelism + ZeRO + mixed-precision recipe end to end.

Compute, parameters, and data: the FLOP identity

Before any scaling law you need a flat-FLOP accounting. Per training token, a dense Transformer does roughly $2 N$ forward FLOPs and $4 N$ backward FLOPs (one re-evaluate plus one gradient matmul of the same size as forward), giving the textbook approximation

where $C$ is total training FLOPs, $N$ is parameter count, and $D$ is the number of tokens trained on. The factor $6$ is a rule of thumb: it ignores the attention $n^2$ term (relevant only at very long context) and small heads like the embedding matrix. Plug in GPT-3: $N = 1.75 \times 10^{11}$, $D = 3 \times 10^{11}$. Then $C \approx 6 \cdot 1.75 \times 10^{11} \cdot 3 \times 10^{11} = 3.15 \times 10^{23}$ FLOPs, which matches the $\approx 3 \times 10^{23}$ FLOPs reported in the GPT-3 paper. Hold this identity in your head; every scaling-law derivation rides on it.

Kaplan 2020 scaling laws

The Kaplan exponents motivate two practical tricks. First, you can predict loss before training by fitting the law on small models and extrapolating. The GPT-3 budget was set this way: extrapolate the law, see that $\sim 10^{23}$ FLOPs land you near $L \approx 2.0$, allocate the cluster. Second, the law gives a sharp diagnostic for “is this run on track”: if at compute $C$ your loss is $0.1$ above the predicted $L(C)$, something is wrong (a hyperparameter, a data issue, a numerical bug) and another epoch will not save you.

Chinchilla: compute-optimal training

The Chinchilla optimum is a guideline for the compute-frontier run. In practice, modern open-weights LLMs train past the optimum on purpose. LLaMA-2 7B sees 2T tokens at $D / N = 286$, more than $10\times$ past Chinchilla. The reason is inference cost: a 7B model is much cheaper to serve than a 70B, and since serving accounts for the vast majority of total project cost, paying extra training compute to ship a smaller deployable model is a clear win [8]. The Chinchilla rule names the training-loss-optimal point; the inference-economics-optimal point is to the right of it.

Training infrastructure: the three parallelism axes

A 70B-parameter model in BF16 needs $140$ GB just for weights, and Adam adds another $\sim 12 \times$ on top of that during training. No single GPU has that memory. The three orthogonal parallelism axes are how a modern cluster fits a model that does not fit on one chip [14].

Data parallelism (gradient all-reduce)

Replicate the entire model on every GPU, split the batch across GPUs. Each rank computes its forward and backward on its slice of the batch and produces local gradients; the optimizer step needs the average across all ranks, so an all-reduce sums and broadcasts the gradients before the update.

Cost. Memory per GPU is unchanged (every rank has a full copy). Communication per step is $O(N)$ in the parameter count: each rank sends and receives roughly $2 N$ bytes (in FP16) per step under ring-all-reduce. The bandwidth is hidden behind compute when batch-per-GPU is large enough; data parallelism is the cheapest way to scale as long as the model fits on one GPU. Data-parallel training is the default when memory permits and is the only axis that does not change the per-step compute pattern. The all-reduce primitive predates the deep-learning era; the connection to gradient descent is just “add the per-shard gradients and divide by the number of shards” [14].

Tensor parallelism (Megatron, splitting one matmul)

When the model itself does not fit, split a single matrix multiply across GPUs. Megatron-LM [3] partitions the per-layer weights along a chosen dimension and synchronizes only at the boundaries.

For a Transformer block, the natural cuts are the $Q, K, V, O$ projections and the FFN. Take the FFN: $Y = \mathrm{GeLU}(X W_1) W_2$ with $W_1 \in \mathbb{R}^{d \times d_\mathrm{ff}}$, $W_2 \in \mathbb{R}^{d_\mathrm{ff} \times d}$. Megatron splits $W_1$ column-wise across $T$ GPUs (each rank holds $d_\mathrm{ff} / T$ columns), runs the GeLU locally without communication, then splits $W_2$ row-wise (each rank holds $d_\mathrm{ff} / T$ rows). The output is a partial sum on each rank; one all-reduce at the FFN exit gives the full output. Two all-reduces per Transformer block: one after FFN, one after attention. Communication scales with hidden size $d$ per token, not with parameter count.

Tensor parallelism works inside one node (NVLink or NVSwitch, 600+ GB/s). Across nodes the all-reduce becomes the bottleneck; modern recipes cap tensor-parallel size at $T = 8$ (one node) and stack other axes on top.

Pipeline parallelism (GPipe, splitting the layer stack)

Split the layer stack across GPUs: ranks 0 hold layers 0-15, rank 1 holds 16-31, and so on. A microbatch flows through the pipeline rank by rank during forward, then gradients flow back during backward. Naive layer-split has a “bubble” where downstream ranks idle until the first microbatch arrives.

GPipe [4] fixes this with microbatching: split the global batch into $M$ microbatches and feed them through the pipeline back to back. With $P$ stages and $M$ microbatches the pipeline bubble is $\frac{P - 1}{P + M - 1}$ of the step time. With $M = 4 P$ the bubble is below $20\%$, which is the practical sweet spot.

Cost. Memory per GPU is $1 / P$ of the full model (each stage holds one $P$-th of the layers). Communication is point-to-point activation passes between adjacent stages; under interleaved 1F1B schedules the volume is $O(\mathrm{batch} \cdot \mathrm{seq} \cdot d)$ per stage. Pipeline parallelism slots cleanly between nodes when tensor parallelism saturates a single node.

The standard 3D recipe is TP $\times$ PP $\times$ DP: tensor-parallel inside a node ($T = 8$), pipeline-parallel across nodes ($P =$ stages), data-parallel on top ($D =$ replicas). Megatron-Turing-NLG-530B trained at $T = 8$, $P = 35$, $D = 6$ on 2240 A100s.

ZeRO: sharding the optimizer state

The cost of ZeRO is communication. ZeRO-1 adds essentially no extra traffic versus standard data parallelism. ZeRO-2 swaps one all-reduce for a reduce-scatter plus an all-gather of weights at next forward, which is bandwidth-equivalent. ZeRO-3 adds an all-gather of weights at every layer’s forward and backward, doubling the per-step communication. Practitioners pick the smallest stage that fits the model: ZeRO-1 if you can, ZeRO-2 if you cannot, ZeRO-3 only if the parameters themselves do not fit. DeepSpeed and PyTorch FSDP both implement the full ladder.

Mixed precision: FP16, BF16, and FP8

Training in FP32 doubles memory and halves compute throughput on Tensor Cores. Modern training is mixed precision: forward and backward in a low-precision format (FP16 or BF16), optimizer state and master weights in FP32 [6].

FP16. Has $5$ exponent bits and $10$ mantissa bits, giving $\sim 6 \times 10^{-5}$ to $6.5 \times 10^{4}$ as the representable positive range. Activations and gradients near the small end of this range round to zero (an “underflow”); near the large end they overflow to infinity. The latter is silently catastrophic in deep stacks. The standard fix is dynamic loss scaling: multiply the loss by a large constant $S$ before the backward pass, so the gradients are $S$ times larger and stay in the FP16 range, then unscale before the optimizer step. If any unscaled gradient is non-finite, halve $S$ and retry; otherwise periodically double $S$ to track the largest representable gradient.

BF16. Has $8$ exponent bits (matching FP32) and $7$ mantissa bits. The wider exponent range means no loss scaling needed: gradients that would underflow in FP16 are still representable in BF16. The price is $3$ fewer mantissa bits, so individual operations are noisier. In practice BF16 trains modern LLMs with no quality loss versus FP32 and is the default on A100, H100, and TPUs.

FP8. Has either $4$ exponent + $3$ mantissa (E4M3, for forward) or $5$ exponent + $2$ mantissa (E5M2, for backward) bits. Halves memory and doubles throughput again versus BF16. Requires per-tensor scaling factors and is currently used selectively (linear layers) rather than across the whole model. The H100 Transformer Engine and recent H200 / B200 launches lean heavily on FP8 for the next throughput jump.

The interaction with the optimizer matters. Adam’s variance term $v$ accumulates the square of the gradient, which can underflow even in BF16 if you store it there. Standard practice is to keep $v$ (and $m$, and the master weights) in FP32 regardless of the forward-pass precision. This is the ”$12$ FP32 bytes per parameter” line in the ZeRO budget above.

Gradient checkpointing

Forward activations are the second-largest memory line item after Adam state. A Transformer block at $d = 4096$, $n = 4096$, batch $B = 1$ stores $\sim 4 \cdot d \cdot n \cdot B \cdot 2 = 134$ MB per layer in BF16; for $80$ layers that is $\sim 10$ GB per sample.

Chen et al. [7] showed how to trade compute for memory: keep only a $\sqrt{L}$ subset of the activations, and recompute the others on the backward pass. The optimal recomputation strategy stores activations every $\sqrt{L}$ layers, giving total memory $O(\sqrt{L})$ at the cost of one extra forward pass per backward (a $\sim 33\%$ compute overhead).

In practice the call is tactical: apply checkpointing to the layers you cannot fit otherwise, leave the rest un-checkpointed. Modern frameworks expose this as torch.utils.checkpoint.checkpoint_sequential or torch.utils.checkpoint.checkpoint, which wraps a forward function and re-runs it on demand. Combining ZeRO-3 sharding with selective checkpointing is the standard recipe for training $\geq 70$B-parameter models on commodity 80 GB GPUs.

Learning rate, batch size, and warmup at scale

A scaling-law plot tells you what target loss is reachable, not how to actually drive a training run there. The optimization-side recipe has three knobs that interact: peak learning rate, batch size, and the warmup-plus-decay schedule.

Peak learning rate. Scales inversely with model width. The Kaplan and Chinchilla papers both used a peak around $3 \times 10^{-4}$ for AdamW on 100M-class models and dropped it for the largest sizes. As a rule of thumb the peak learning rate at width $d$ scales as $\eta_{\max} \propto 1 / \sqrt{d}$ in the original Vaswani schedule and as $\eta_{\max} \propto 1 / d$ under maximal-update parameterization (muP). Most production runs sit between these and use empirical sweeps at small scale, then transfer to the full run.

Global batch size. Modern LLMs train at very large global batch sizes: GPT-3 used $3.2$M tokens per step, LLaMA-2 used about $4$M, and Chinchilla used $1.5$M scaling to $3$M. Big batches reduce the number of optimizer steps to reach a target loss but eventually hit a “critical batch size” past which doubling the batch buys less than half a step of progress. The critical batch size grows during training, so warmup batch-size schedules are common at the very top of the scale.

Warmup and decay. A linear warmup over the first $\sim 1\%$ of steps, followed by a cosine decay to $10\%$ of peak, is the modern default. Warmup is required because the optimizer’s variance estimate $v$ is noisy at step 0; running at peak learning rate before $v$ stabilizes blows up the parameter trajectory. The cosine decay tail is empirical: a flat-then-cosine schedule fits final loss curves better than a constant-then-step decay, with no clear theory [14].

The 3D-parallel + ZeRO + mixed-precision stack changes which of these knobs is easy to vary. Global batch size is set by the data-parallel degree times microbatch times pipeline-microbatches; halving the DP degree halves the global batch but doubles the per-rank microbatch, which shifts the activation memory budget. Practitioners pick global batch first, then choose parallelism degrees that produce that batch at the per-rank memory ceiling.

Wiring it up: PyTorch FSDP and DeepSpeed ZeRO

The two production training frameworks for ZeRO-style sharding are PyTorch FSDP (built into torch.distributed) and DeepSpeed (the original ZeRO implementation from [5]). Both expose roughly the same configuration surface: pick a stage, pick the precision, pick a checkpoint policy.

import torch
import torch.distributed as dist
from torch.distributed.fsdp import (
  FullyShardedDataParallel as FSDP,
  ShardingStrategy,
  MixedPrecision,
  CPUOffload,
)
from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
from transformers.models.llama.modeling_llama import LlamaDecoderLayer
from functools import partial

dist.init_process_group(backend="nccl")
torch.cuda.set_device(dist.get_rank() % torch.cuda.device_count())

model = build_llama_7b()  # your nn.Module

# ZeRO-3 equivalent: SHARD_GRAD_OP shards optimizer + gradients;
# FULL_SHARD also shards parameters (true ZeRO-3).
mp = MixedPrecision(
  param_dtype=torch.bfloat16,
  reduce_dtype=torch.bfloat16,
  buffer_dtype=torch.bfloat16,
)

wrap_policy = partial(
  transformer_auto_wrap_policy,
  transformer_layer_cls={LlamaDecoderLayer},
)

model = FSDP(
  model,
  sharding_strategy=ShardingStrategy.FULL_SHARD,   # ZeRO-3
  mixed_precision=mp,
  auto_wrap_policy=wrap_policy,
  cpu_offload=CPUOffload(offload_params=False),
  device_id=torch.cuda.current_device(),
)

opt = torch.optim.AdamW(model.parameters(), lr=3e-4, betas=(0.9, 0.95))

for batch in loader:
  out = model(batch.input_ids, labels=batch.labels)
  out.loss.backward()
  opt.step()
  opt.zero_grad(set_to_none=True)

The choice between FSDP and DeepSpeed is mostly ergonomic. FSDP integrates more cleanly with vanilla PyTorch training loops and the wider PyTorch ecosystem. DeepSpeed has older, more battle-tested implementations of CPU and NVMe offload (ZeRO-Infinity), MoE layer support, and a broader set of debugging tools. Both produce identical training loss curves on the same configuration; both inherit the trade-offs [5].

The 2023-2024 open-weights landscape

The practical consequences of all this stack show up in the open-weights LLM lineage of 2023 and 2024. Each model below sits at a different point on the (Chinchilla, inference-cost, license) trade-space.

LLaMA-2 7B and 70B trained on $2$T tokens each, the same data budget regardless of model size [8]. The 70B is on the Chinchilla optimum at $D / N \approx 28$; the 7B is way past at $D / N \approx 286$, paying extra training compute for a smaller serving model. Mistral 7B keeps the same 7B parameter count but adds grouped-query attention and sliding-window attention of $4096$ tokens, dropping inference cost while matching or beating LLaMA-2 13B on most benchmarks [9]. Mixtral 8x7B is the same lineage gone MoE: $8$ experts of $7$B-FFN each, top-$2$ routing, $13$B active parameters per token but $47$B total parameters in memory [10]. Qwen 72B trained on $\geq 3$T tokens with a similar architecture to LLaMA-2 [11]. DeepSeek-V2 took the MoE idea further: $236$B total parameters, $21$B active per token, $8.1$T training tokens, and a novel multi-head latent attention to compress the KV cache [12]. Phi-3-mini is the contrarian point: $3.8$B parameters trained on $3.3$T heavily filtered “textbook-quality” tokens, deliberately past Chinchilla but tuned for on-device inference [13].

A few patterns are visible across this table. First, the bigger dense models (LLaMA-2 70B, Qwen 72B) sit close to the Chinchilla optimum at $D / N \approx 20$ to $40$. Second, the smaller dense models (LLaMA-2 7B, Mistral 7B, Phi-3-mini) train far past the optimum to lower inference cost, accepting more training compute for a cheaper serving footprint. Third, the MoE family (Mixtral, DeepSeek-V2) decouples active from total parameters, which scales total capacity without scaling per-token compute; the price is the routing complexity and the load-balance auxiliary loss. The license column matters as much as any architectural choice for downstream users.

Pitfalls

Pop quiz

Interview Q&A

Interactive Demo