Gradients and Optimizers: From SGD to Adam, and Why m and v Cost 8 Bytes

Read a gradient as an instruction to one weight

Here is the version that sticks. Training is rolling a ball down into a valley, where height is the loss (how wrong the model is) and position is the weights. The gradient is the slope under the ball. It tells you which way is uphill and how steep. To reduce the loss you step the opposite way, downhill.

Formally the gradient is the partial derivative of the loss with respect to a weight. That phrase is the only piece of calculus vocabulary in this part, so take it apart once.

A derivative is a rate: how fast one number changes when another number changes. The word partial means you nudge exactly one input and hold every other input perfectly still. That restriction is what makes the idea usable here. A model has 1.54 billion weights, and asking what happens when all of them move at once tells you nothing you can act on. Asking about one weight at a time gives you 1.54 billion separate, simple answers.

So a weight’s gradient reads as “if I nudge this one weight up a hair and change nothing else, does the loss rise or fall, and how fast”. Nothing more mysterious than that. Put numbers on it. The loss is 2.4000. You increase one weight by 0.001, and the loss becomes 2.4030. The loss rose by 0.0030 in response to a nudge of 0.001, so the rate is 0.0030 divided by 0.001, which is 3. That weight’s gradient is 3. It is positive, which says up on this weight is uphill on the loss, so the weight should go down.

The backward pass is simply the machinery that computes this instruction for every weight at once, running back from the loss through the layers, and it never has to nudge anything and remeasure to do it. The full derivation of every derivative involved is available if you want the calculus rather than the ball.

One consequence matters for memory. There is exactly one gradient number per weight, which is why the gradients tenant is the same size as the weights tenant, 2 bytes per parameter in a bf16 run.

Separate what the gradient decides from what the optimizer decides

The gradient gives a direction but not a policy. Do you step exactly along it? A fraction of it? Do you remember which way you were heading and carry some momentum? Do different weights get different step sizes? Those decisions belong to the optimizer.

What it remembers between steps is the optimizer state, and the optimizer state is the memory tenant. Remembering is neither free nor temporary. Whatever an optimizer keeps is allocated once per trainable weight and then sits in VRAM from the first step of the run to the last. SGD remembers almost nothing. AdamW remembers two numbers per weight. That difference is the whole story of why the optimizer is 12 of the 16 bytes per parameter that the part on the four memory tenants derived.

Price SGD and momentum, and see where one learning rate runs out

Stochastic gradient descent is the plainest policy: take the gradient, multiply by a learning rate, subtract it from the weight. The learning rate is a single small number, often around 0.0001, that scales every step everywhere in the model. SGD remembers nothing between steps, so it costs zero extra memory per parameter.

Carry the earlier weight through it. Its gradient was 3 and the learning rate is 0.0001, so the step is 3 times 0.0001, which is 0.0003. The gradient was positive, meaning uphill, so the step is subtracted. A weight of 0.5000 becomes 0.4997. That is a complete SGD update, and there is nothing else in it.

Cheap and simple. The catch is that it uses one global learning rate for every weight in the network.

Momentum costs one number per weight, 4 bytes in fp32, and genuinely helps. But it does not fix the core limitation. One learning rate still has to suit every weight in the model, and in a transformer that is a much harder ask than it sounds.

Explain m and v, and why they cost 8 bytes per parameter

Adam’s idea is to stop using one global learning rate and instead give every weight its own adaptive step, derived from that weight’s own gradient history. To do that it remembers two running statistics per weight, conventionally called m and v.

In words: m is a running average of the gradient, so it points the way with the noise smoothed out. v is a running average of the squared gradient, so it measures how large and erratic this particular weight’s gradients have been. Dividing by the square root of v means a weight with large, noisy gradients takes cautious small steps, while a weight with small steady gradients is allowed to move faster.

The braking is easier to believe with numbers in it. Take two weights in the same model. The first sees gradients of roughly 1.0 every step, so its running average of squared gradients settles near 1.0, and the square root of that is 1.0. The second sees gradients of roughly 0.01, so its v settles near 0.0001, and the square root is 0.01. Adam divides each weight’s smoothed gradient by its own square root. For the first weight that is 1.0 divided by 1.0, which is 1. For the second it is 0.01 divided by 0.01, which is also 1. Both weights end up taking a step of about the learning rate, even though their raw gradients differ by a factor of a hundred. Adam does not care how big a weight’s gradients are. It cares how they compare with that weight’s own recent history.

Every weight is tuned individually, and there is no single global learning rate to agonise over. The cost of that per-weight intelligence is storing m and v: 4 bytes plus 4 bytes, so 8 bytes per parameter.

Hold the running example against that figure. Qwen2.5-1.5B, the 1.54 billion parameter model this series sizes everything against, needs 1.54 billion times 8 bytes for m and v alone. That is 12.3 GB. Two numbers per weight, whose only job is to remember, take up more of the card than the weights and the gradients put together, which come to 6.2 GB between them at 2 bytes each. Those are static-state figures. The activations from the forward pass are counted separately and Part 3 covers them.

Why per-weight step sizes matter so much for transformers

This is the part that justifies the memory. A transformer’s gradients vary enormously across its parts: attention against feed-forward, early layers against late, layer norms against projections, embeddings against everything. With one global learning rate, a value that is right for one part is wrong for another. Too large and some weights diverge. Too small and others barely learn. Pascanu et al. described the first failure in its sharpest form, the exploding gradient, and proposed the guardrail still used against it: scale the whole set of gradients down whenever their combined size crosses a threshold. Part 6 covers that guardrail and why it is a separate setting from anything in this part.

Adam’s v term auto-scales each weight’s step to its own gradient magnitude, so one setting works across the whole varied network. You et al. took the same principle one level up, adding a per-layer rescaling that kept training stable at batch sizes in the tens of thousands, where a single global rate had stopped working. That robustness is why practically every large language model is trained with an Adam-family optimizer despite the memory cost. It removes a brutal tuning problem, and tuning problems cost far more than VRAM.

Say what the W in AdamW changes, and what it costs

AdamW is Adam with one fix, and the fix is about weight decay. Weight decay is a small pull applied to every weight on every step, toward zero, so that no weight grows larger than it needs to be. It is there to stop the model fitting its training data too tightly.

AdamW applies that decay directly to the weight, outside the adaptive step, instead of folding it into the gradient where Adam’s own per-weight scaling used to distort it. The distortion is worth one sentence. Fold the decay into the gradient and it gets divided by the square root of v along with everything else, so a weight with large noisy gradients quietly receives less decay than a quiet one, which is not a policy anybody chose. The fix costs nothing in memory, since it changes the update rule rather than the stored state.

The next part dissects the update equations line by line, including why bias correction exists.

Answer the argument for dropping back to SGD, and cut the right thing instead

AdamW’s power is also its cost. The m and v it stores are the biggest memory tenant in a training run. There are two ways to react to that, and they are not equally good.

The tempting one is to drop back to SGD and save the 8 bytes. The better one is to keep AdamW and store m and v in fewer bits.

Take the SGD argument apart line by line, because it comes up in every memory-constrained conversation.

• Yes, you would save memory. SGD stores nothing extra, or 4 bytes with momentum, against AdamW’s 12.
• But transformer training with SGD is notoriously finicky. With one global learning rate you must hand-tune schedules, and it often still underperforms.
• The v term is precisely what makes Adam robust to the wildly different gradient scales across a transformer’s layers, so one config works across the whole network.
• So the field pays 8 bytes per parameter for not having to babysit training.
• The right move under pressure is to shrink the state (8-bit Adam) or remove it for most weights (LoRA), rather than to abandon per-weight adaptivity.

What block-wise quantisation actually does

Shrinking the state deserves its own paragraphs, because “store the same two numbers in fewer bits” hides a real problem. Quantisation means storing a number by mapping it onto a small set of allowed values. Eight bits gives you 256 of them, so every value in a list has to be rounded to one of 256 rungs. Where you put those rungs decides how much you lose. Spread them evenly between the smallest and largest value in the whole list, and one outlier stretches the scale so far that every ordinary value near zero collapses onto the same few rungs.

Block-wise quantisation refuses to use one scale for the whole list. It chops the list into small blocks, a couple of thousand values each in Dettmers et al.’s original 8-bit optimizer work, and gives every block its own scale factor stored alongside it. Inside one small block the values tend to be of similar size, so that block’s 256 rungs land where its numbers actually are. An outlier now spoils only its own block instead of the entire tensor.

The bookkeeping is cheap. One extra 4-byte scale per couple of thousand values is a rounding error next to the values themselves, which is why m and v at 1 byte each come out at about 2 bytes per parameter rather than exactly 2. Run that through the example model. The optimizer tenant goes from 12 bytes per parameter to about 6, because the fp32 master copy of the weights is untouched and still costs its 4. On Qwen2.5-1.5B that is 18.5 GB of optimizer state down to about 9.2 GB, and static state as a whole from about 24.6 GB to about 15.4 GB. Roughly 8 GB of activations is still owed on top of both figures, so a full fine-tune goes from about 33 GB, which does not fit a 32 GB card, to about 23 GB, which does with room to spare.

Two related variants are worth naming now because they appear in later parts. Paged Adam offloads optimizer state to host RAM during a memory spike and brings it back, converting a hard crash into a brief slowdown. And LoRA removes the tenant entirely for every weight it freezes, which is the subject of the part on low-rank adaptation.

Measure the 8 bytes yourself rather than trusting them

You can see the 8 bytes directly. This walks a single linear layer through one step and sums the optimizer’s own tensors. Pin your PyTorch version before you run it: optimizer internals are not a stable public API and have changed between releases.

# torch 2.x
import torch

layer = torch.nn.Linear(4096, 4096, bias=False)      # about 16.8M parameters, fp32
opt = torch.optim.AdamW(layer.parameters(), lr=1e-4)

layer(torch.randn(2, 4096)).sum().backward()          # produce gradients
opt.step()                                            # AdamW allocates m and v here

def state_bytes(optimizer):
    return sum(t.numel() * t.element_size()
               for slot in optimizer.state.values()
               for t in slot.values()
               if torch.is_tensor(t))

n = sum(p.numel() for p in layer.parameters())
print(f"{n:,} params, {state_bytes(opt) / n:.1f} bytes of optimizer state per param")

If you have never used PyTorch, that block is doing ten things. Here they are in order.

1. torch.nn.Linear(4096, 4096, bias=False) builds one layer. Its weights are a 4,096 by 4,096 grid, which is 16,777,216 numbers, created in fp32 unless you say otherwise. Switching the bias off keeps the parameter count to exactly that grid, so the division at the end comes out clean.
2. layer.parameters() hands over the list of tensors the optimizer is allowed to change, which here is that single grid. A tensor is a grid of numbers with any number of dimensions. torch.optim.AdamW(...) wraps that list and sets the learning rate. Nothing has been allocated for m and v yet.
3. layer(torch.randn(2, 4096)) pushes two rows of random numbers through the layer, which is a forward pass. .sum() adds the whole output up into a single number, standing in for a loss. What that number means does not matter here. What matters is that there is exactly one of it, because a backward pass has to start from one number.
4. .backward() is the backward pass. It computes a gradient for every parameter and attaches it to that parameter, so gradients now sit alongside the weights in memory.
5. opt.step() is the update. The first time it runs, AdamW finds it has no history for this parameter and creates some: two new tensors the same shape as the weights, one holding m and one holding v.
6. optimizer.state is where that history lives. It is a dictionary whose keys are the parameter tensors themselves and whose values are one small dictionary per parameter, holding whatever this optimizer needs to remember about it. .values() walks those small dictionaries and ignores the keys, which is all the function needs.
7. The inner loop, for t in slot.values(), walks the remembered items inside one of those small dictionaries. AdamW puts three things in each: exp_avg, which is PyTorch’s name for m, exp_avg_sq, its name for v, and a step counter.
8. torch.is_tensor(t) guards the arithmetic, because not every entry in that dictionary is guaranteed to be a grid of numbers. Current PyTorch stores the step counter as a one-element tensor, so it passes the guard and adds 4 bytes to a total of 134 million, which does not move the printed answer. Older versions and some third-party optimizers store the counter as a plain integer, which has no size to measure and would raise an error. The guard makes one function work across both, and it is the reason the version caution above is not boilerplate.
9. t.numel() is how many numbers a tensor holds, and t.element_size() is how many bytes each of those numbers takes, which is 4 for fp32. Multiply the two for the bytes that tensor occupies, then sum over every remembered tensor for the whole optimizer state.
10. The last two lines count the parameters and divide. The output is 16,777,216 params, 8.0 bytes of optimizer state per param. Two fp32 numbers per weight, exactly as advertised.

Note that the optimizer state does not exist until the first step(), which is a common source of confusion when a run survives the forward pass and then OOMs on what looks like nothing. The figure printed is about 8, not 12, and the reason is worth understanding: this toy layer is pure fp32, so its parameters already are the master copy. The extra 4 bytes appear only in a mixed-precision run, where a separate fp32 master is kept alongside the bf16 compute weights.

Tell the two appearances of fp32 apart, and know which one you pay for

The letters fp32 turn up twice in any description of a mixed-precision run, meaning two completely different things. One of them lands on your memory bill and the other never does. Separating them is the last piece of the optimizer’s 12 bytes, and it is the single most common confusion about mixed precision.

The first appearance is transient. A bf16 matrix multiply accumulates its partial sums in fp32 inside the tensor core and hands back a result. Make that concrete. Multiplying one row of 1,536 numbers by one column of 1,536 numbers means 1,536 separate multiplications, all added into one running total. If that total were rounded back to bf16 after each of the 1,536 additions, the small rounding errors would pile up inside it. So the hardware keeps the running total in fp32 while it is being built and rounds once, at the very end. That total lives inside the chip for the duration of one multiply and is never written out as a stored tensor, so it costs no persistent memory. It is also not something you configure. It happens whether you know about it or not.

The second appearance is persistent, and it is the one in the memory accounting. It is a high-precision copy of the weights kept so that thousands of tiny per-step updates add up instead of each rounding away.

That sentence is easy to nod at, so check it with real numbers. A 16-bit format does not store every decimal number. It stores a fixed set of representable values, and near any given magnitude those values sit a fixed distance apart. For bf16 near a weight of 0.5, the gap between one representable value and the next is about 0.0039. Nothing in between exists. So take a weight of 0.5 and apply an update of 0.0001. The result, 0.5001, is not a value bf16 can hold, and the nearest value it can hold is 0.5 itself. The weight is written back unchanged. The update was not made less accurate. It was deleted.

Now do the same in fp32, whose representable values near 0.5 are about 0.00000006 apart. An update of 0.0001 is more than a thousand steps along that finer grid, so it lands cleanly and the weight becomes 0.5001. Run two thousand steps of that size and fp32 has accumulated 0.2 of real movement, while bf16 is still sitting on exactly 0.5, having thrown away every one of the two thousand updates for the same reason each time. This is why the fix has to be storage rather than arithmetic. No amount of care during the multiply rescues an update that the destination cannot record.

So the correct one-liner is: mixed precision does the arithmetic in bf16, which is fast and small, and keeps the weights’ running total in fp32, which is precise, so that small updates survive. Not “the multiplication result is stored in fp32”, even though the multiply does accumulate in fp32 internally. The stored fp32 is the weight master copy, and it exists to defeat the vanishing-update problem step after step after step.

That stored copy is 4 bytes per parameter, and it is the third of the three fp32 numbers AdamW keeps for every weight: the master copy, m, and v. Three numbers at 4 bytes each is where 12 of the 16 bytes per parameter come from. It is also the 4 bytes the measurement script could not show you, because a pure fp32 layer has no separate master to count.