Lesson 66: Mixed Precision, Loss Scaling & Gradient Accumulation

Discussion prompt

Before we open Lesson 66: Mixed Precision, Loss Scaling & Gradient Accumulation: without looking back, what was the main idea of Naive Bayes Classifiers, and what could you do by the end of it that you could not do before?

Hint: One sentence for the idea, one for the skill. If the second one is blank, that is the part to revisit.

Comparison matrix

From Two floating-point formats: refill the bits column from what you know. The rest of the table is as it appeared.

Sort into buckets

Sort these by dtype, from memory, without looking back at Why mixed precision helps. Telling them apart on the spot is the skill; the table is only where the answer happens to be written down.

FP16

activations (fwd); gradients (bwd)

FP32

master weights; optimizer state (m, v)

g1

dtype is "FP16" for activations (fwd), gradients (bwd) — that is what the table on "Why mixed precision helps" records, and it is the single property separating this group from the rest.

g2

dtype is "FP32" for master weights, optimizer state (m, v) — that is what the table on "Why mixed precision helps" records, and it is the single property separating this group from the rest.

Comparison matrix

From autocast: letting PyTorch pick the dtype: every row here is a choice with a cost. Fill the why column, then say which row you would actually pick and what you give up for it.

Step through it

Step through The underflow problem one row at a time. What is driving the change, and what would the row after the last one be?

1. Step 1: gradient value is 1.0e-04
2. Step 2: gradient value is 6.1e-05
3. Step 3: gradient value is 1.0e-08

Predict first

Trace the scale-before / unscale-after pattern. Scale S = 65536 (2^16 = GradScaler's default initial value).

Commit before you compute: what does Loss scaling by hand come out to? A rough magnitude and the right form is enough — the point is to have something concrete to be wrong about.

Correct: Unscaled gradient = 6.554e-04 / 65536 ≈ 1.00e-08 — correctly recovered

Why: A prediction you can defend turns the computation into a check rather than a leap of faith — and an answer that contradicts it is caught on the spot. Scaling shifts the gradient into FP16's representable range.

Discussion prompt

Work backwards. The example finished here:

Unscaled gradient = 6.554e-04 / 65536 ≈ 1.00e-08 — correctly recovered

What was it asked to do, and what must it have been given? Reconstruct the problem from its answer.

Hint: Every quantity in the result had to enter somewhere. Account for each one.

Discussion prompt

torch.amp.GradScaler automates the scale-and-unscale loop. It also detects overflowed gradients (inf/NaN) and skips the optimizer step that iteration.

That is stated as though it always holds. Do one of two things: produce a case where it fails, or say precisely what rules such a case out. "It just does" is not on the menu.

Hint: Hunt at the extremes first — zero, one, negative, empty, equal. If every extreme survives, the reason they survive is the proof.

Discussion prompt

The canonical mixed-precision loop adds three things to the standard loop (Lesson 40): autocast, scaler.scale(loss), and scaler.step / scaler.update.

What do you need to know — or decide — before the first line can be written? List everything the problem has to hand you.

Hint: Anything you would have to invent to get started is a thing the problem must supply.

Step through it

Step through The AMP training loop one row at a time. One of these columns never changes — find it, and say why it cannot.

1. Step 1: epoch is 1
2. Step 2: epoch is 2
3. Step 3: epoch is 3
4. Step 4: epoch is 4
5. Step 5: epoch is 5

Predict first

A student writes this, and it looks reasonable:

After scaler.step(optimizer), the model is updated — the loop is done for this step.

It is wrong. Say what breaks — and say it before you turn the page.

Correct: The GradScaler's internal scale factor never changes.

Always call scaler.update() after scaler.step(optimizer) each iteration.

Why: The GradScaler's internal scale factor never changes. After any overflow step, the scale was supposed to halve — but without update() it stays dangerously high, causing repeated inf gradients and skipped optimizer steps.

Discussion prompt

The rule this trap just fixed:

Always call scaler.update() after scaler.step(optimizer) each iteration.

Now break it on purpose. Build a case that violates it and follow the consequences until something visibly fails. Where does the failure first show up — and would you have noticed it if you had not been looking?

Hint: The dangerous rules are the ones whose violation still produces an answer. If yours fails loudly, try to find one that fails quietly.

Discussion prompt

Explain Why accumulate? by analogy to something with no Machine Learning in it at all — a queue, a recipe, a map, a bank balance, whatever fits. Then say where your analogy breaks.

Hint: An analogy that never breaks is not an analogy, it is the same idea wearing a hat. Find the seam — that is the part that is actually new.

Predict first

128 samples, 2 features, logistic regression. Accumulate 4 mini-batches of 32 vs single pass on 128. Weights must match exactly.

Commit before you compute: what does Accumulation = full batch: the proof come out to? A rough magnitude and the right form is enough — the point is to have something concrete to be wrong about.

Correct: Key: divide loss by N_ACCUM before backward — else each mini-batch contributes a full loss instead of 1/N of one

Why: A prediction you can defend turns the computation into a check rather than a leap of faith — and an answer that contradicts it is caught on the spot. BCEWithLogitsLoss uses reduction='mean' inside each mini-batch.

Comparison matrix

From Accumulation = full batch: the proof: refill the scaled_loss column from what you know. The rest of the table is as it appeared.

Predict first

Run the single full-batch reference and confirm the weights match to 8 decimal places.

Commit before you compute: what does Verifying the equivalence come out to? A rough magnitude and the right form is enough — the point is to have something concrete to be wrong about.

Correct: Both methods produce w = [[0.03085, 0.02717]] — difference = 0.0 (bitwise identical)

Why: A prediction you can defend turns the computation into a check rather than a leap of faith — and an answer that contradicts it is caught on the spot. PyTorch accumulates gradients with += into the .grad tensors.

Discussion prompt

Work backwards. The example finished here:

Both methods produce w = [[0.03085, 0.02717]] — difference = 0.0 (bitwise identical)

What was it asked to do, and what must it have been given? Reconstruct the problem from its answer.

Hint: Every quantity in the result had to enter somewhere. Account for each one.

Predict first

A student writes this, and it looks reasonable:

In the accumulation loop, backward on the raw loss each mini-batch — grads just add up.

It is wrong. Say what breaks — and say it before you turn the page.

Correct: Each mini-batch contributes its own full mean loss, so N iterations accumulate N× the intended gradient.

Divide the loss by N_ACCUM before calling backward each mini-batch.

Why: Each mini-batch contributes its own full mean loss, so N iterations accumulate N× the intended gradient. The effective learning rate is multiplied by N — causing overshooting or divergence.

Sort into buckets

Some of these hold up and some are the exact mistakes this lesson is built to prevent. Sort them.

Holds up

FP32 and FP16 differ in how many bits they use for the exponent and mantissa — that gap determines both range and precision.; The idea: run the forward pass and backward pass in FP16 (fast, half the bandwidth), but keep optimizer state and master weights in FP32 (numerically stable).; torch.autocast wraps the forward + backward and auto-casts ops to the most efficient dtype — no manual .half() calls needed.

Breaks

After scaler.step(optimizer), the model is updated — the loop is done for this step.; In the accumulation loop, backward on the raw loss each mini-batch — grads just add up.

sound

These are stated as this lesson states them — each one survives the edge cases Lesson 66: Mixed Precision, Loss Scaling & Gradient Accumulation puts it through.

flawed

Each of these is lifted from a trap in this deck: reasonable-sounding, and wrong in a way that only shows up once you rely on it.

Sort into buckets

Sort these by FP32 Adam (bytes/param), from memory, without looking back at Where AMP memory is (and isn't) saved. Telling them apart on the spot is the skill; the table is only where the answer happens to be written down.

4

weights; gradients

4 + 4 = 8

optimizer m, v

16

total params+opt

g1

FP32 Adam (bytes/param) is "4" for weights, gradients — that is what the table on "Where AMP memory is (and isn't) saved" records, and it is the single property separating this group from the rest.

g2

FP32 Adam (bytes/param) is "4 + 4 = 8" for optimizer m, v — that is what the table on "Where AMP memory is (and isn't) saved" records, and it is the single property separating this group from the rest.

g3

FP32 Adam (bytes/param) is "16" for total params+opt — that is what the table on "Where AMP memory is (and isn't) saved" records, and it is the single property separating this group from the rest.

Discussion prompt

Run Mixed precision + gradient accumulation recipe with this step confiscated:

Step + update with scaler.step(opt); scaler.update() — unscales grads, skips overflow steps, adapts scale

Is it still possible? If it is, say what takes its place and what it costs you. If it is not, say exactly what that step was providing that nothing else does.

Hint: A step you can drop for free was never load-bearing. If you cannot drop it, name the thing that goes wrong the moment it is gone.

Discussion prompt

Mixed precision + gradient accumulation recipe works on the cases you have just seen. Push it to the edge: what is the most degenerate input it still handles — empty, zero, one item, everything equal — and what is the first case where it stops being true? Name the case, not just "it breaks".

Hint: Try the smallest legal input, then the largest, then the one where two things collide. Methods are specified at their edges; the middle takes care of itself.

Eliminate the wrong options

A gradient of 1e-08 is cast to FP16. What is the stored value?

3 of these 4 are wrong. Strike them one at a time, and say what rules each one out before you strike the next. The survivor is the answer.

• A. 0.0 — underflow, the value is lost
• B. 1.0e-08 — FP16 can represent it exactly
• C. 1.192e-07 — rounded up to the nearest normal
• D. inf — the value overflows

Survives elimination: A

Why: FP16's smallest representable normal is 6.10e-05; 1e-08 is below even the subnormal range (~6e-08). torch.tensor(1e-8).half().item() == 0.0. Loss scaling multiplies it by ~65536 before the cast to keep it alive.

Check your understanding

A gradient of 1e-08 is cast to FP16. What is the stored value?

• A. 0.0 — underflow, the value is lost (correct)
• B. 1.0e-08 — FP16 can represent it exactly
• C. 1.192e-07 — rounded up to the nearest normal
• D. inf — the value overflows

Answer: A

Why: FP16's smallest representable normal is 6.10e-05; 1e-08 is below even the subnormal range (~6e-08). torch.tensor(1e-8).half().item() == 0.0. Loss scaling multiplies it by ~65536 before the cast to keep it alive.

Why B tempts people

FP16 has only 10 mantissa bits; its representable range starts at ~6e-08 (subnormal) and 1e-08 falls below that.

Why C tempts people

1.192e-07 is the FP16 representation of 1e-07, not 1e-08. The value rounds down to 0, not up.

Why D tempts people

Overflow produces inf for values > 65504. 1e-08 is tiny; it underflows to 0, not inf.

Predict first

Which is the correct order of one AMP training step?

Answer it in your own words, now, with nothing to choose from. The options are on the next slide — and picking the right one off a list is an easier skill than producing it.

Correct: zero_grad → autocast(forward+loss) → scaler.scale(loss).backward() → scaler.step(opt) → scaler.update()

Why: zero_grad first (clear old grads); forward + loss inside autocast; backward on the scaled loss; scaler.step unscales grads and calls opt.step if no overflow; scaler.update adjusts the scale factor.

Check your understanding

Which is the correct order of one AMP training step?

• A. zero_grad → autocast(forward+loss) → scaler.scale(loss).backward() → scaler.step(opt) → scaler.update() (correct)
• B. autocast(forward) → zero_grad → loss.backward() → opt.step() → scaler.update()
• C. zero_grad → forward → scaler.scale(loss).backward() → opt.step() → scaler.update()
• D. zero_grad → autocast(forward+loss) → loss.backward() → scaler.step(opt) → scaler.update()

Answer: A

Why: zero_grad first (clear old grads); forward + loss inside autocast; backward on the scaled loss; scaler.step unscales grads and calls opt.step if no overflow; scaler.update adjusts the scale factor.

Why B tempts people

zero_grad must come before autocast/forward, not after. Also, backward must be on the scaled loss — not the raw loss — for loss scaling to work.

Why C tempts people

The loss must be inside autocast for correct dtype casting. C wraps only the forward, leaving the loss computation in FP32 by default — but more critically, zero_grad must precede autocast.

Why D tempts people

D omits scaler.scale(loss) before backward, so the loss is not scaled and FP16 gradients will underflow. The unscale step inside scaler.step would also then divide by scale unnecessarily.

Eliminate the wrong options

You want to accumulate gradients over 4 mini-batches to simulate batch-128 training (32 per mini-batch). Which backward call is correct?

3 of these 4 are wrong. Strike them one at a time, and say what rules each one out before you strike the next. The survivor is the answer.

• A. (loss_fn(model(xb), yb) / 4).backward() — divide loss before backward
• B. loss_fn(model(xb), yb).backward() — just accumulate the raw loss
• C. loss_fn(model(xb), yb).backward(); model.weight.grad /= 4 — scale after
• D. loss = sum of all mini-batch losses; loss.backward() — accumulate loss, not grads

Survives elimination: A

Why: With reduction='mean', each mini-batch BCELoss already averages over its 32 samples. To get the full-128-sample mean gradient, each mini-batch's contribution must be 1/4 of its per-sample mean. Dividing before backward achieves this exactly — verified: weight diff vs single full batch = 0.0.

Check your understanding

You want to accumulate gradients over 4 mini-batches to simulate batch-128 training (32 per mini-batch). Which backward call is correct?

• A. (loss_fn(model(xb), yb) / 4).backward() — divide loss before backward (correct)
• B. loss_fn(model(xb), yb).backward() — just accumulate the raw loss
• C. loss_fn(model(xb), yb).backward(); model.weight.grad /= 4 — scale after
• D. loss = sum of all mini-batch losses; loss.backward() — accumulate loss, not grads

Answer: A

Why: With reduction='mean', each mini-batch BCELoss already averages over its 32 samples. To get the full-128-sample mean gradient, each mini-batch's contribution must be 1/4 of its per-sample mean. Dividing before backward achieves this exactly — verified: weight diff vs single full batch = 0.0.

Why B tempts people

Raw accumulation counts each mini-batch's mean at full weight, producing 4× the intended gradient norm. The effective step size is 4× too large, causing oscillation or divergence.

Why C tempts people

Scaling the gradient after all 4 backward passes is mathematically equivalent — but it requires the gradients not to overflow FP16 before the correction. In practice, dividing the loss before backward keeps everything scaled safely throughout.

Why D tempts people

Summing scalars and calling backward once would try to backpropagate through all 4 mini-batches' computation graphs — which no longer exist after the loop. The forward graph is consumed per-batch; you must call backward inside the loop.

Discussion prompt

Build a training loop that combines mixed precision (autocast + GradScaler) with gradient accumulation (N=4 mini-batches). Verify the final weights match a plain FP32 full-batch reference.

That is stated as though it always holds. Do one of two things: produce a case where it fails, or say precisely what rules such a case out. "It just does" is not on the menu.

Hint: Hunt at the extremes first — zero, one, negative, empty, equal. If every extreme survives, the reason they survive is the proof.

Comparison matrix

From Milestone 1 — AMP loop (no accumulation yet): every row here is a choice with a cost. Fill the expected scale column, then say which row you would actually pick and what you give up for it.

Comparison matrix

From Milestone 2 — add gradient accumulation: refill the expected column from what you know. The rest of the table is as it appeared.

Draw it

One page, no notation unless you need it: draw how these connect — FP16 & FP32: what they can hold · Loss scaling: preventing underflow · Gradient accumulation: simulating large batches · Memory budget & the full picture · Your turn: AMP + accumulation. Put an arrow wherever one of them is what makes another possible, and label the arrow with why.