The memory tax
A single graphics card has a fixed amount of memory. The largest cards available to us hold 80 GB. The model has 30 billion numbers in it, and every one of those numbers has to be stored somewhere while the model is being trained.
Start with the arithmetic, because it settles the question quickly. Each weight in the model carries more than just itself. It carries a record of how it should change, a high-accuracy copy of itself, and two running averages that the training procedure keeps between steps.
| Item | Bytes |
|---|---|
| the weight, in the 16-bit format used for arithmetic | 2 |
| its gradient, the number saying how it should change | 2 |
| a 32-bit copy of the weight, kept for accuracy | 4 |
| two running averages the optimizer keeps | 8 |
| total | 16 |
Sixteen bytes for every weight, not two. The weight you think of as the model is one eighth of what training actually holds.
Multiply that by the parameter count and the situation is clear before any clever engineering enters the picture.
Two unit systems, one number
That footprint gets quoted two ways and it is worth knowing they are the same thing. Multiply it out once and then divide by whichever base you are using.
30 × 109 weights × 16 bytes = 480,000,000,000 bytes
÷ 109 → 480 GB ÷ 230 → 447 GiB
Not two quantities. The same bytes counted in thousands or in 1024s. Storage vendors and
marketing use the decimal form; operating systems and nvidia-smi use the binary
one, which is why a card sold as 80 GB reports about 81,500 MiB.
The card is quoted the same way, so 80 GB is 74.5 GiB. What matters is not which system you pick but that you do not mix them inside one division.
| GiB | GB | Cards | |
|---|---|---|---|
| training footprint | 447 | 480 | |
| one card | 74.5 | 80 | |
| cards needed | 6.0 | 6.0 | 6 |
Six in both columns, because each column divided within its own system. Divide 480 GB by 74.5 GiB and you get 6.4, which is not a quantity of anything.
The model is not the problem. Training it is.
Calling 447 GiB the model size is the one mistake worth avoiding, because it makes the situation sound worse than it is. That figure is the training footprint. The model on its own is a fraction of it.
| Situation | Bytes per weight | Total |
|---|---|---|
| inference — weights only, 16-bit | 2 | 60 GB |
| training — the full sixteen | 16 | 480 GB |
Sixty gigabytes fits on one card with room to spare. So the model is not too big for the hardware — it is too big only while you are training it, and only because of the twelve bytes of optimizer state that exist for no other purpose. That is a much sharper statement of the problem, and it points directly at which band to attack first.
And the gap does not close by buying a better card, because 80 GB is already the top of the market. It closes by putting the training state on several cards at once, which is the entire subject of this session.
The three bands
Group those four rows into three categories, because the rest of the session takes them apart one at a time and in this order. Naming them now saves a lot of words later.
| Parameters the working copy the forward pass reads |
2 |
| Gradients produced by the backward pass, consumed by the update |
2 |
| Optimizer states the 32-bit master copy plus the two running averages |
12 |
Notice the imbalance. The optimizer states are twelve of the sixteen bytes — three quarters of the footprint sits in the band that has nothing to do with the forward pass and is touched exactly once per training step. That lopsidedness is not incidental. It is the reason the first stage of ZeRO, which only shards that band, already wins most of the memory back.
Where the 12 comes from
Training runs the arithmetic in 16-bit for speed, but 16-bit numbers are too coarse to accumulate tiny updates without losing them. So the optimizer keeps a 32-bit master copy of every weight, updates that, and rounds it back down to 16-bit for the next forward pass. Adam then keeps two more 32-bit numbers per weight: a running average of the gradient and a running average of its square. Four bytes plus eight bytes is the twelve.
The four rows are one loop
Those four items are not four independent costs sitting side by side. They are four stages of a cycle that runs once per training step, and seeing the cycle explains why each one has the width it has.
The reason for the round trip is precision. An update is w − lr × g,
and lr × g is routinely millions of times smaller than w. A
16-bit number carries about three decimal digits. Add something that much smaller to it and
the result rounds straight back to where it started — the update disappears. Run a few
thousand steps that way and the weight has learned nothing.
So the subtraction is done in 32 bits, where the small number survives, and that value is what persists between steps. The 16-bit weight is refreshed from it after every update and used for nothing but speed.
The two averages
They are the optimizer's memory. One tracks the average of recent gradients for that weight, the other tracks the average of their squares. Both are held in 32-bit floating point, so they cost four bytes each.
| Held in 32-bit | Purpose | Bytes |
|---|---|---|
| master copy of the weight | the authoritative value updates are applied to | 4 |
| average of recent gradients | which direction this weight has been moving | 4 |
| average of their squares | how erratically it has been moving | 4 |
| optimizer states | 12 |
The model this session counts
The model has 30 billion parameters. That is the number every calculation in this session uses, on every page, without restating it each time.
One more thing this table leaves out, deliberately. It counts only what is stored per weight. Activations — the intermediate values a forward pass produces on its way to a loss — are also held in memory, and they scale with batch size rather than with parameter count. They are a real cost and a different problem, handled by different tools. Everything in this session is about the sixteen bytes.