Lesson 6. Gradients and Optimizer State

Mission link: This is the argument for adapters. Once you can derive the sixteen-bytes-per-parameter figure, LoRA stops being a trick and becomes the obvious response.
Primary source: Paper: "ZeRO: Memory Optimizations Toward Training Trillion Parameter Models", Rajbhandari et al., arXiv:1910.02054
Prerequisites: Lesson 4, Lesson 5

Warm-up

1. ▢ Weights of a 7B model in bf16: how many gigabytes?

2. ▢ Which two tensors does AdamW keep per trainable parameter?

3. ▢ Why does bf16 beat fp16 for training?

Know this

Everything a trainable parameter costs, per parameter:

The right-hand column is the part that surprises people. Mixed precision does not reduce the total.

Both bars reach the same mark, which is the whole finding. The two accented pieces are the tensors that really were halved, four bytes saved between them, and the master copy sitting beside the first one is four bytes put straight back. Everything to the right of it never moved at all. It halves the weight and gradient tensors, then adds back an fp32 master copy of the weights, because accumulating tiny updates into a bf16 value silently loses them. Optimizer moments stay in fp32 for the same reason. The win from mixed precision is arithmetic throughput, not memory.

So: 16 bytes per trainable parameter, before a single activation.

What that means

1.5B trainable × 16 bytes = 24 GB
7B   trainable × 16 bytes = 112 GB
70B  trainable × 16 bytes = 1120 GB

Full fine-tuning a 7B model needs 112 GB of optimizer-and-weight state alone, then activations on top. That is multi-GPU territory for a model whose weights are 14 GB. The gap between "the model fits" and "training the model fits" is a factor of eight, and it is entirely gradients and optimizer state.

The escape

Look again at the table and notice what the 16 bytes are multiplied by: the number of trainable parameters, not total parameters. Frozen parameters cost 2 bytes each, for their weights, and nothing more. No gradient. No moments.

That is the whole argument for adapters:

frozen parameter   = 2 bytes  (bf16 weights only)
trainable parameter = 16 bytes (weights + master + grad + 2 moments)

Freeze the base model, add a few million trainable parameters, and the second term nearly vanishes. Worked for a 7B base with 4M trainable adapter parameters:

frozen base    7B  × 2  = 14.0 GB
trainable      4M  × 16 = 0.064 GB

The adapter's optimizer state is 64 MB. It has become a rounding error against the base weights, and the base weights are the one part you can attack separately, with quantisation, in stage 4.

Partial measures, and why they are not enough

Before adapters existed, people reduced this cost in other ways. They are all still available and all still bounded:

Every one of these scales the constant. None of them changes the fact that the cost is proportional to the number of trainable parameters. Adapters attack the multiplier instead, which is why they win by orders of magnitude rather than by factors.

Why the moments cannot simply be dropped

A reasonable question: if AdamW's two moments are two thirds of the cost, why not use plain SGD? Because transformer training depends on Adam's per-parameter step scaling. Gradient magnitudes vary enormously across a transformer's parameters, and a single global learning rate serves them badly. The second moment is what lets each parameter take an appropriately sized step. This is empirical, consistent, and not a tuning problem you can wish away.

Practice

1. ▢ Derive the memory needed to full fine-tune a 3B model with AdamW in mixed precision, excluding activations.

2. ▢ Why doesn't mixed precision reduce total training memory?

3. ▢ A 7B base is frozen in bf16 and you train a 20M-parameter adapter. Total, excluding activations?

4. ▢ Which change reduces optimizer memory most for a fixed model?

   • a) Switching the optimizer from AdamW to SGD
   • b) Switching the compute precision to bf16
   • c) Freezing the base and training an adapter
   • d) Switching to completion-only loss masking

5. ▢ Someone proposes an 8-bit optimizer so they can full fine-tune a 70B model on one 80 GB accelerator. Do the arithmetic and respond.

6. ▢ In one sentence: why do frozen parameters cost so much less than trainable ones?

Real-world reps

• [ ] Build the 16-bytes table from memory on paper. Then compute total training memory for three model sizes you actually care about.
• [ ] Load a small model, run one training step with AdamW on all parameters, and record peak memory. Compare to your prediction.
• [ ] Tomorrow: repeat with the base frozen and only a small adapter trainable. The ratio you measure is the argument for stage 3.

Going further

• Paper: "ZeRO: Memory Optimizations Toward Training Trillion Parameter Models", Rajbhandari et al., arXiv:1910.02054: the canonical accounting of these categories
• Paper: "8-bit Optimizers via Block-wise Quantization", Dettmers et al., arXiv:2110.02861
• Memory budget
• Lesson 8. The Low-Rank Idea: the response to this lesson

Not landing? Reread the primary source at the top, since this lesson compresses it and compression is where understanding leaks. Check the glossary for any term that felt slippery.

If the lesson itself is unclear rather than the material, that is a defect: open an issue.