Mixed precision without NaNs
Mixed precision is trustworthy only when dtype choices, loss scaling, skipped steps, and quality parity are all observable.
Training repeatedly OOMs and misses its nightly completion window.
This chapter moves from lock a full-precision reference to unscale before inspection. We will make the idea concrete in a small code example before placing it back inside a production system with 24 GB device and nightly completion.
For FP16, unscale gradients before clipping and log the scale plus skipped-step count. Compare convergence and final metrics to the reference.
Verify the behavior illustrated in amp_step.py, then measure decision utility, slice reliability, and system latency.
Unsynchronized timing reports impossible speedups while retained tensors keep computation graphs alive and leak memory. An unsafe release must degrade to a named baseline or the last known-good version.
Build the idea one layer at a time.
Mixed precision is a controlled numerical trade: use smaller numbers where they are safe, and preserve range or accuracy where training needs it.
Lower-precision tensors reduce memory traffic and can unlock specialized matrix hardware. FP16 has limited exponent range, so small gradients may underflow to zero and large values may overflow. BF16 keeps an FP32-sized exponent range but has fewer mantissa bits, which often makes it easier to train while still requiring numerical checks.
Automatic mixed precision chooses lower precision for suitable operations and keeps sensitive operations in safer types. FP16 loss scaling multiplies the loss before backpropagation, then unscales gradients before clipping or stepping. The scaler must skip an optimizer update when gradients contain infinities; hiding those events removes the earliest signal of instability.
- pinned memory
- FP16/BF16
- loss scaling
- activation memory
Follow the reasoning before touching the code.
- 01
Lock a full-precision reference
Record loss curves, held-out quality, and a few golden logits in FP32. Speed without a numerical comparison is not an optimization result.
- 02
Autocast the forward path
Let the framework select operation dtypes instead of casting the entire model blindly. Keep targets and loss semantics unchanged.
- 03
Unscale before inspection
For FP16, unscale gradients before clipping and log the scale plus skipped-step count. Compare convergence and final metrics to the reference.
import torch
model = torch.nn.Linear(16, 1).cuda()
optimizer = torch.optim.AdamW(model.parameters())
scaler = torch.amp.GradScaler('cuda')
x = torch.randn(64, 16, device='cuda')
y = torch.randn(64, 1, device='cuda')
optimizer.zero_grad()
with torch.autocast(device_type='cuda', dtype=torch.float16):
loss = torch.nn.functional.mse_loss(model(x), y)
scaler.scale(loss).backward()
scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
scaler.step(optimizer)
scaler.update()One FP16 training step with safe scaling and clipping order.
What to notice
- Clipping happens after unscaling, so the threshold applies to the true gradient magnitude rather than the amplified values.
- The optimizer step is mediated by the scaler, which can skip it when overflow is detected and adjust the scale for later steps.
YOUR TURN Pause and predict before revealing the answer.
Why is calling `optimizer.step()` directly after `scaler.scale(loss).backward()` incorrect?
Need a hint?
The parameter gradients are still multiplied by the current scale.
Check your reasoning
The optimizer would apply scaled gradients and bypass overflow detection. `scaler.step(optimizer)` first ensures gradients are usable and performs the real step only when safe.
Lessons 1, 2 in this course.
Connect the chapter to the course’s quantitative spine.
Mixed precision is trustworthy only when dtype choices, loss scaling, skipped steps, and quality parity are all observable. The roofline bound below is the shared course-level lens that helps quantify the surrounding system. Name each symbol, check its units or shape, and use only values that were available at the production cutoff.
Roofline bound
| Symbol | Meaning / shape / unit |
|---|---|
P_peak | peak compute |
I | arithmetic intensity |
B_mem | memory bandwidth |
Open derivation and numerical substitution
Start from the production quantity being optimized, substitute the observed values with their declared units, then isolate the model-controlled term. Preserve shape annotations at each step so broadcasting or aggregation cannot silently change the result.
- Write the named inputs: P_peak, I, B_mem.
- Substitute one small, hand-checkable batch before vectorizing.
- Calculate an independent reference value and compare within a declared tolerance.
# equation → code contract
inputs = validate_shapes_and_units(batch)
value = compute_c12(inputs)
assert is_finite(value)Change one assumption and make the tradeoff visible.
Use the worked chapter as your prediction, then change one input in the lab and explain why the output moves. The lab runs predefined TypeScript and never executes learner code, so the focus stays on the concept rather than setup.
GPU memory ledger
Change microbatch size and calculate a simplified training-memory budget.
Assumption: 3.2 GB weights+grads+optimizer, 0.34 GB activations per sample, 24 GB device, 10% reserve.
Open nonvisual data table
| Item | Computed state | Interpretation |
|---|---|---|
| block 1 | 1.09 GB activations | resident |
| block 2 | 2.18 GB activations | resident |
| block 3 | 3.26 GB activations | resident |
| block 4 | 4.35 GB activations | resident |
| block 5 | 5.44 GB activations | resident |
Trace the complete operating path.
Carry the chapter’s central idea into the system: Mixed precision is trustworthy only when dtype choices, loss scaling, skipped steps, and quality parity are all observable. The path below shows where that reasoning needs an artifact, a test, an owner, or a fallback.
- 01
Turn “Lock a full-precision reference” into a written invariant and a failing test.
- 02
Promote amp_step.py from illustration to a versioned implementation with representative fixtures.
- 03
Use “Unscale before inspection” as the release checkpoint: For FP16, unscale gradients before clipping and log the scale plus skipped-step count. Compare convergence and final metrics to the reference.
- 04
Serve or materialize behind explicit 24 GB device and nightly completion.
- 05
Join telemetry to mature outcomes and retain a rollback path.
Observability
Join service health, input quality, prediction distributions, slice behavior, and mature outcomes by exact version.
Cost
Measure storage, preprocessing, compute, queueing, and human review under a representative arrival pattern.
Failure modes
Unsynchronized timing reports impossible speedups while retained tensors keep computation graphs alive and leak memory. Add a detector, owner, mitigation, and stop condition for this class of failure.
Alternatives
Compare a rule, a simpler statistical baseline, and a different system boundary before adding model complexity.
Explain the contract, not just the vocabulary.
Single-GPU optimization plan
Profile and optimize a correctness-locked trainer without sacrificing model quality.
- Profiler diagnosis
- Mixed precision parity
- Memory plan
- Throughput/cost report
Read primary material with a purpose.
Return to the opening failure.
Record hardware/software, shapes, precision, seeds, throughput method, peak memory, and cost per sample.
For this lesson, the release evidence is the worked code example, a hand-checked formal result, deterministic simulation output, a ≥80% checkpoint, the production rubric, and a named fallback. The course resolves when the system can produce a profiled, mixed-precision, memory-budgeted training loop.