Efficiency Frontiers › Training Efficiency

The cost of training large models has been a fundamental constraint on AI development. Training GPT-3 once cost around $4 million; training GPT-4 is estimated to have cost over $100 million. Efficiency improvements here compound directly into more training runs, more experimentation, and faster iteration.

The shift from FP32 to mixed-precision training (BF16 or FP16 for forward and backward passes, FP32 for optimizer states) was one of the first major training efficiency wins. It halves memory usage for most tensors and doubles throughput on hardware with tensor cores.

Loss scaling is required for FP16 training because the smaller exponent range causes underflow in gradients. The loss is multiplied by a large scale factor $S$ before the backward pass, the gradients are unscaled before the optimizer step, and the scale factor is adjusted dynamically (increased when no overflow occurs, decreased when it does).

BF16 avoids this problem. With the same exponent range as FP32 but only 8 bits of mantissa (versus 24 for FP32), BF16 covers the same dynamic range but at lower precision. The reduced precision in the mantissa is generally acceptable because gradient noise in stochastic optimization already introduces similar-scale errors.

Standard backpropagation requires storing all layer activations computed during the forward pass, because they are needed to compute gradients during the backward pass. For a transformer with $L$ layers and batch size $B$, this requires storing $O(L \cdot B \cdot n \cdot d)$ activations in memory, which grows linearly with model depth.

Gradient checkpointing (also called activation recomputation) trades compute for memory. During the forward pass, only a subset of layer outputs (checkpoints) are stored. During the backward pass, the activations at non-checkpointed layers are recomputed from the nearest preceding checkpoint. By checkpointing every $\sqrt{L}$ layers, the memory for activations is reduced from $O(L)$ to $O(\sqrt{L})$ at the cost of approximately 33% additional compute per training step.

At large scale, this tradeoff is almost always worthwhile. Training a 70B-parameter model with full activation storage would require enormous memory, while checkpointing reduces this to a manageable level at a modest compute overhead.

The datasets used for training have grown alongside models, but data quality matters more than quantity. Recent work on data filtering and dataset curation shows that training on a smaller, higher-quality dataset often outperforms training on a much larger noisy one.

Chinchilla scaling laws (Hoffmann et al., 2022) established that previous large models were undertrained. The optimal training configuration allocates roughly equal compute to model size and training data: for a model with $N$ parameters, the compute-optimal number of training tokens is approximately $20N$. A 70B-parameter model should train on about 1.4 trillion tokens. This insight shifted the field from "largest possible model" to "right-sized model with enough data."

Curriculum learning pre-trains models in stages, presenting easier or more structured data first and more complex data later. Code training, for example, benefits from showing well-formatted code early before introducing partially broken or obscure examples. The rationale is that the model builds a strong foundation on clear examples before having to handle noisy or ambiguous ones.

Repetition and replay are important considerations for long training runs. Seeing the same data too many times causes memorization and reduces generalization. Many recent training pipelines do multiple passes over high-quality data (e.g., textbooks, scientific papers) but only one pass over noisier sources, explicitly managing the per-source repetition rate.

Standard Adam requires storing first and second moment estimates for every parameter, tripling the optimizer memory footprint beyond the parameters themselves. For a 70B model, Adam adds roughly 420 GB of optimizer state (at FP32) on top of the 140 GB for the parameters.

Adafactor reduces this by factoring the second moment estimate. For a weight matrix $W \in \mathbb{R}^{m \times n}$, instead of storing a full $m \times n$ second moment matrix, Adafactor stores rank-1 factors $\hat{v}_r \in \mathbb{R}^m$ and $\hat{v}_c \in \mathbb{R}^n$ such that $v_{ij} \approx \hat{v}_{r,i} \hat{v}_{c,j}$. This reduces the second moment storage from $O(mn)$ to $O(m + n)$, a substantial savings for large weight matrices.

CAME (Confidence-guided Adaptive Memory Efficient optimization) extends this with a confidence matrix that distinguishes reliable gradient directions from noisy ones, improving training stability on large models.

Sophia uses an approximate Hessian diagonal rather than gradient squared as the second moment estimate, which produces better step size calibration and has shown faster convergence on language modeling benchmarks.

Deep learning frameworks like PyTorch execute operations one at a time by default. Each operation involves a kernel launch, memory allocation, and HBM read/write. For many small operations (layer normalizations, activation functions, elementwise additions), these overhead costs dominate the actual computation.

Kernel fusion combines multiple operations into a single GPU kernel. Instead of separate kernels for the bias addition, GELU activation, and dropout in a feed-forward layer, a single fused kernel reads the input once, applies all three operations, and writes the output once. This can reduce memory bandwidth by 3-5x for these operations.

torch.compile (introduced in PyTorch 2.0) automates kernel fusion using the Triton JIT compiler. By tracing the computation graph and identifying fusible operations, it can speed up training loops by 20-30% with a single line of code: model = torch.compile(model). This makes kernel fusion accessible without requiring custom CUDA code.

XLA (Accelerated Linear Algebra), used by JAX and TensorFlow, takes a more aggressive approach by compiling entire training steps into optimized machine code. It performs whole-program optimization including operation reordering, buffer aliasing, and cross-operation fusion, often achieving higher utilization than PyTorch's per-operation approach.