Memory-Efficient Optimisers

Training memory is four things: parameters, gradients, optimiser state, and activations. Checkpointing addresses the last. This entry is about the third, which is usually the largest.

For a 7B model that is 112 GB before a single activation. Two thirds of it is optimiser state.

Adafactor: factorise the second moment

Adam’s $v$ is a full-shape tensor of squared-gradient averages. For a matrix parameter, Shazeer and Stern approximate it by a rank-one outer product of a row vector and a column vector.

This is not an arbitrary approximation. Among all rank-one non-negative approximations it is the one minimising the generalised KL divergence to $V$, and the row and column sums are exactly the sufficient statistics — the derivation is the paper’s contribution, not the idea of factorising.

Memory for a $4096 \times 4096$ matrix falls from 16.7 M numbers to 8192. Momentum is usually dropped as well, at which point the state is effectively free.

8-bit Adam: quantise the state

Keep both moments in full shape and store them in one byte each.

The difficulty is dynamic range. Optimiser state spans many orders of magnitude, and a naïve 8-bit linear quantisation of a tensor with one large outlier destroys everything else in it. Dettmers et al. use two ideas together: block-wise quantisation, with an independent scale per 2048 elements so one outlier contaminates only its block, and a dynamic non-linear codebook that spends its 256 values logarithmically rather than uniformly.

The result matches 32-bit Adam’s final loss across models up to 1.5B — not approximately, but within run-to-run variance — at a quarter of the memory. It is what bitsandbytes provides and what most consumer-GPU fine-tuning uses.

Low-rank projection

GaLore observes that the gradient of a large weight matrix is empirically close to low-rank, and keeps the optimiser state in a projected subspace.

$P$ comes from an SVD of the gradient, recomputed every few hundred steps. State memory falls by the rank ratio, and unlike LoRA the weights are still updated at full rank — the constraint is on the optimiser, not on the parameterisation, so nothing is permanently confined to a subspace.

The claim is 7B pre-training on a 24 GB card. The caveats are that the periodic SVD is not free and that the quality comparisons are at smaller scale than the headline.

import torch
from torch import Tensor


class FactoredSecondMoment:
    """Row and column sums in place of the full V. Matrices only —
    vectors and 1-D parameters keep an ordinary full-shape state."""

    def __init__(self, shape: tuple[int, int], device):
        n, m = shape
        self.row = torch.zeros(n, device=device)
        self.col = torch.zeros(m, device=device)

    def update(self, g: Tensor, beta2: float) -> Tensor:
        sq = g.float().square() + 1e-30
        self.row.mul_(beta2).add_(sq.mean(dim=1), alpha=1 - beta2)
        self.col.mul_(beta2).add_(sq.mean(dim=0), alpha=1 - beta2)

        # Reconstruct: the rank-one approximation, normalised by the row mean.
        return torch.outer(self.row / self.row.mean(), self.col)

The 1e-30 is load-bearing. A zero row sum makes the reconstruction divide by zero, and a parameter whose gradient is exactly zero for a whole batch is common enough — an unused expert, a padded position — that it happens on real runs rather than in theory.

The last row of that table is the one to be careful with. The others are established; the projection methods are recent, and the gap between “matches AdamW at 1B” and “matches AdamW at frontier scale” has swallowed several optimisers already.