Gradient descent, Adam, and the memory bill nobody mentions

You already have the gradient: a vector telling you, for every weight, which direction increases the loss. Optimization is the answer to “now what do I do with that vector?” The simplest answer, SGD (stochastic gradient descent), is to subtract a small multiple of the gradient from the weights and repeat:

$W \leftarrow W - \text{lr} \cdot \nabla L(W)$

This is gradient descent exactly as it was proven to work: moving opposite the gradient’s direction is provably the direction of steepest decrease, nothing added. Adam is a smarter version of that same subtraction. Instead of using the raw gradient directly, it keeps a running memory of past gradients to smooth out noise, and it adapts how big a step to take for each individual parameter, since some parameters genuinely need bigger steps than others and Adam learns which is which as training proceeds.

It’s worth saying up front why this matters beyond convergence speed, because that’s not actually the headline reason it matters at frontier scale: Adam’s memory cost is the direct cause of a large fraction of the distributed-training engineering that exists in this industry. That claim gets justified with real numbers below, not asserted.

Adam’s two moving averages#

Adam tracks two exponential moving averages per parameter, updated at every step $t$ from the gradient $g_t$:

$m_t = \beta_1 m_{t-1} + (1-\beta_1)g_t, \qquad v_t = \beta_2 v_{t-1} + (1-\beta_2)g_t^2$

$m_t$, the first moment, is a smoothed running average of the gradient itself: which direction has been good, on average, recently. $v_t$, the second moment, is a smoothed running average of the squared gradient: how large or variable this specific parameter’s gradient has been, regardless of sign. Typical values are $\beta_1=0.9$, $\beta_2=0.95$.

The update divides the (bias-corrected) first moment by the square root of the (bias-corrected) second moment:

$\hat m_t = \frac{m_t}{1-\beta_1^t}, \qquad \hat v_t = \frac{v_t}{1-\beta_2^t}, \qquad W_t = W_{t-1} - \text{lr}\cdot\frac{\hat m_t}{\sqrt{\hat v_t}+\epsilon}$

Dividing by $\sqrt{\hat v_t}$ is Adam’s adaptive part: a parameter whose gradient has consistently been large gets its effective step size shrunk, since you’re dividing by a big number; a parameter with a small, quiet gradient history gets a relatively larger step. Every parameter effectively earns its own learning rate, derived from its own recent gradient history, rather than sharing one global rate the way plain SGD does. $\epsilon$ is a tiny constant, the same defensive pattern as every other numerical-stability fix in this series, preventing division by zero when $\hat v_t$ is near zero.

Bias correction, derived rather than pasted in#

The $1-\beta_1^t$ and $1-\beta_2^t$ denominators look like a detail to memorize, but they fix a specific, derivable problem. Both moving averages start at $m_0 = v_0 = 0$. At step 1:

$m_1 = \beta_1(0) + (1-\beta_1)g_1 = (1-\beta_1)g_1$

That’s only $(1-\beta_1)$ times the true gradient, a severe underestimate, because the moving average hasn’t had time to “fill up” yet. Dividing by $(1-\beta_1^t)$ fixes this exactly: at $t=1$, $(1-\beta_1^t) = (1-\beta_1)$, so $\hat m_1 = m_1/(1-\beta_1) = g_1$ exactly, the correction exactly cancels the early-step underestimate. As $t$ grows, $\beta_1^t \to 0$, so the correction factor fades to a no-op precisely when the moving average no longer needs it. It’s a self-retiring fix, not a permanent adjustment.

AdamW: decoupling weight decay from the adaptive scaling#

AdamW adds weight decay as a separate term, subtracted after the adaptive update rather than folded into the gradient itself:

$W_t = W_{t-1} - \text{lr}\cdot\frac{\hat m_t}{\sqrt{\hat v_t}+\epsilon} - \text{lr}\cdot\lambda \cdot W_{t-1}$

The reason this needs to be a separate term rather than just adding $\lambda W$ to the gradient before it enters $m$ and $v$: doing it the naive way means the decay itself gets distorted by Adam’s adaptive per-parameter scaling, so a weight with a large historical gradient would get less decay applied to it, for no principled reason. Decoupling makes decay behave predictably regardless of a parameter’s gradient history, which is why AdamW rather than plain Adam is standard in essentially all LLM training, typically with $\beta_1=0.9$, $\beta_2=0.95$.

Learning rate schedule: warmup, then cosine decay#

Training almost never uses a single fixed learning rate. It warms up, linearly increasing from 0 to a peak over roughly the first 2,000 steps, to avoid the large, noisy gradients that come from a fresh random initialization destabilizing training before it’s found its footing. Then it decays, most commonly on a cosine schedule:

$\text{lr}(t) = \frac{1}{2}\cdot\text{peak}\cdot\left(1+\cos\left(\frac{\pi t}{T}\right)\right)$

At $t=0$: $\cos(0)=1$, giving $\text{lr}=\text{peak}$. At $t=T$: $\cos(\pi)=-1$, giving $\text{lr}=0$. The schedule starts at the peak and decays smoothly to zero exactly at the end, trading large, fast-but-risky steps early for small, precise steps as training converges. Batch size typically ramps the opposite direction, small at first for stability, larger later for throughput.

def cosine_lr_schedule(t, T, peak_lr):
    return 0.5 * peak_lr * (1 + np.cos(np.pi * t / T))

def linear_warmup_then_cosine(t, warmup_steps, total_steps, peak_lr):
    if t < warmup_steps:
        return peak_lr * (t / warmup_steps)          # linear ramp, 0 to peak
    decay_t, decay_T = t - warmup_steps, total_steps - warmup_steps
    return cosine_lr_schedule(decay_t, decay_T, peak_lr)

print(linear_warmup_then_cosine(0,     2000, 10000, 3e-4))   # 0.0,     start of warmup
print(linear_warmup_then_cosine(2000,  2000, 10000, 3e-4))   # 3e-4,    peak, end of warmup
print(linear_warmup_then_cosine(10000, 2000, 10000, 3e-4))   # ~0.0,    end of decay

What a real run actually uses. GPT-3’s published training recipe is a good calibration point for what these knobs look like at full scale, rather than in a toy example: $\beta_1=0.9$, $\beta_2=0.95$, gradients clipped to a global norm of $1.0$, weight decay of $0.1$, a linear warmup over the first 375 million tokens, and cosine decay down to 10% of the peak learning rate spread across the next 260 billion tokens. Batch size itself was ramped too, starting small and linearly increasing to its full value over the first few billion tokens of training, the same batch-size-ramp pattern mentioned above, not a simplification of it.

Gradient clipping: capping the vector’s own length#

$\texttt{clip\_grad\_norm\_(params, max\_norm=1.0)}$

caps the gradient vector’s L2 norm before the optimizer step, using exactly the norm concept from earlier in this series: if $\|\nabla L\|_2$ exceeds max_norm, every entry of the gradient gets scaled down proportionally so the whole vector’s length is capped. This is a safety valve against one anomalous, oversized gradient producing a destructive weight update, and skipping it is a common, specific cause of a training run dying to NaN partway through.

Worked example: one Adam step by hand#

$\beta_1=0.9$, $\beta_2=0.95$, $\text{lr}=0.001$, $\epsilon=10^{-8}$, starting from $m_0=v_0=0$, $W_0=1.0$, with gradient $g_1=0.5$ at step $t=1$.

m1 = 0.9(0) + 0.1(0.5)     = 0.05
v1 = 0.95(0) + 0.05(0.5²)  = 0.0125

bias-corrected:
m̂1 = 0.05 / (1 - 0.9¹)  = 0.05 / 0.1  = 0.5
v̂1 = 0.0125 / (1 - 0.95¹) = 0.0125 / 0.05 = 0.25

note: m̂1 = g1 = 0.5 exactly -- confirming the bias-correction derivation above

W1 = 1.0 - 0.001 × (0.5 / (√0.25 + 1e-8))
   = 1.0 - 0.001 × (0.5 / 0.5)
   ≈ 1.0 - 0.001 = 0.999

import numpy as np

def adam_step(W, g, m, v, t, lr=0.001, beta1=0.9, beta2=0.95, eps=1e-8, weight_decay=0.0):
    m = beta1 * m + (1 - beta1) * g
    v = beta2 * v + (1 - beta2) * (g ** 2)
    m_hat = m / (1 - beta1 ** t)
    v_hat = v / (1 - beta2 ** t)
    # AdamW: decay subtracted separately, never folded into the adaptive term above
    W = W - lr * (m_hat / (np.sqrt(v_hat) + eps)) - lr * weight_decay * W
    return W, m, v

W_new, m_new, v_new = adam_step(W=1.0, g=0.5, m=0.0, v=0.0, t=1)
print(W_new)   # 0.999, matching the hand-worked example exactly

The memory bill: why this is a systems problem, not just an algorithms one#

A single training step needs four tensors, each the same size as the model itself: the FP32 weights, the FP32 gradients, and Adam’s own FP32 first and second moments, $m$ and $v$. Four tensors at 4 bytes each is 16 bytes per parameter:

$\text{Total memory} = N_{\text{params}} \times 16 \text{ bytes}$

def adam_memory_bytes(n_params, bytes_per_value=4):
    return 4 * n_params * bytes_per_value   # weights, grads, m, v

print(adam_memory_bytes(7e9) / 1e9)   # 112.0 GB

For a 7-billion-parameter model, that’s exactly $7\times10^9 \times 16 = 112$GB, on a single GPU, before activations or anything else is even counted, and an H100 has 80GB. This isn’t an approximation, it’s exact arithmetic, and it’s why techniques like ZeRO and FSDP exist at all: splitting the optimizer state itself across multiple GPUs, rather than replicating it whole on every one, is the direct engineering answer to this specific 16-bytes-per-parameter number. Optimization, at this scale, isn’t a self-contained algorithms topic, it’s a large fraction of the reason distributed-training infrastructure needs to exist in the first place.

How production actually pays this bill#

Real mixed-precision training doesn’t keep all 16 bytes in FP32 either. The standard recipe (the one behind ZeRO’s own accounting) splits it as: 2 bytes for BF16 weights, 2 bytes for BF16 gradients used in the forward/backward pass, then 4+4+4 bytes for the FP32 master weights, momentum, and variance that the optimizer actually updates. Same 16 bytes total, just split by which pieces need full precision (the optimizer’s running statistics) and which can tolerate a lower-precision copy (the values used for the matmuls themselves).

ZeRO shards that 16-byte total across $N$ data-parallel GPUs, in three stages of increasing aggressiveness:

Stage	What gets sharded across GPUs	Per-GPU memory
ZeRO-1	FP32 master weights, momentum, variance (12 bytes)	$4 + 12/N$ bytes/param
ZeRO-2	+ BF16 gradients (2 bytes)	$2 + 14/N$ bytes/param
ZeRO-3 (equivalent to FSDP)	+ BF16 weights themselves	$16/N$ bytes/param

At $N=8$ GPUs, ZeRO-3 brings that 7B model’s 112GB down to 14GB per GPU, comfortably inside an 80GB H100 with room left for activations. This is the exact mechanism, not a hand-wave: nobody holds a full spare copy of the optimizer state on every GPU, each GPU holds a $1/N$ slice and the full state is reassembled only transiently, communicated just in time, when a given shard’s update is actually needed.

The complementary, non-distributed answer to the same bill is 8-bit optimizers: quantizing the momentum and variance tensors themselves down to INT8 with block-wise dynamic scaling (the bitsandbytes library is the widely used implementation), cutting the optimizer-state portion of the memory roughly in line with the quantization mechanics from earlier in this series applied to $m$ and $v$ instead of to the model’s weights. It’s the same 16-bytes-per-parameter problem, attacked from the precision side instead of the sharding side, and the two techniques compose: fine-tuning stacks commonly run both at once.

Where this actually breaks#

Failure	Cause	Fix
Loss explodes at step 0	Learning rate too high on a fresh random init produces large activations and large gradients immediately	Warmup
Training stalls or oscillates	Learning rate too low, or badly scheduled	Cosine decay with a correctly chosen peak
Embeddings grow unbounded	Weight decay missing or folded in undecoupled	AdamW’s decoupled decay
Entire run dies to NaN	No gradient clipping, one anomalous batch produces a destructive update	clip_grad_norm_
Repeated, unexplained loss spikes mid-run	A specific batch interacting badly with the current optimizer state, not a global LR problem	Restart from an earlier checkpoint, skip the offending batches (PaLM’s actual fix)
112GB doesn’t fit on one GPU	Adam’s four-tensor, 16-byte-per-parameter cost	ZeRO/FSDP sharding, 8-bit optimizer states, or both

These aren’t hypothetical. OPT-175B needed eight separate learning-rate adjustments over two months of training, each triggered by a real instability during the run, and Llama 3’s team ran weeks of smaller proxy experiments specifically to tune the batch-size ramp and cosine schedule before committing to the full run. PaLM’s training hit roughly 20 loss spikes despite gradient clipping already being in place, and the fix that actually worked wasn’t lowering the learning rate: it was restarting from a checkpoint saved about 100 steps before each spike and resuming training while skipping the 200 to 500 batches immediately around it. That the fix was “skip this specific data, not the whole schedule” is itself informative: it points at particular batches interacting badly with the current optimizer state, not a globally mistuned learning rate, which a simple loss-explodes-so-lower-the-LR response would have missed entirely. The mechanics above (momentum, bias correction, adaptive scaling) are the algorithm; the memory arithmetic and the instability firefighting are why it takes a team to actually run at this scale rather than a single correct implementation.

Common mistakes#

Treating the bias-correction denominators as boilerplate to copy rather than a fix for a specific, derivable problem: skipping them, or getting the exponent wrong, silently reintroduces the early-step gradient underestimate this section derived.

Folding weight decay into the gradient instead of applying it separately: this is exactly the mistake AdamW’s decoupling exists to prevent, and it’s easy to reintroduce by hand without noticing.

Assuming a bigger learning rate always trains faster: past a certain point it doesn’t train faster, it destabilizes the run outright, which is the entire reason warmup exists rather than just starting at the peak rate.

Skipping gradient clipping because a run “seems stable”: stability up to a point doesn’t rule out the one anomalous batch that produces a destructive update; clipping is a safety valve for a rare event, not a fix for a chronic problem, and its absence doesn’t show up until it does.

Try it yourself#

Beginner. Using adam_step above, run 5 consecutive steps with a constant gradient $g=0.3$ and watch $\hat m$ converge toward $0.3$ as $t$ grows, confirming the bias-correction fade-out derived above.

Intermediate. Using adam_memory_bytes, compute the memory cost for a 70-billion-parameter model and compare it against a single H100’s 80GB. By how many GPUs’ worth of memory does it exceed capacity?

Advanced. Sketch, as pseudocode, a monitoring rule that flags whether a training run’s logged loss curve looks like “learning rate too high” (a spike), “learning rate too low” (a plateau or oscillation), or a genuine bad batch (a spike that immediately recovers). Turn the failure-mode table above into an actual operational decision procedure rather than a static list.

---

The one-sentence version: SGD is subtracting a scaled gradient and nothing more, Adam adds a smoothed, bias-corrected, per-parameter-adaptive version of that same subtraction, and the reason any of this shows up in a systems conversation rather than staying a pure-math one is that Adam’s own bookkeeping costs 16 bytes per parameter, a number that by itself decides whether a model’s optimizer state fits on the GPU you have.