Precision formats, autocast and gradient scaling

Mixed precision is the cheapest memory and speed win in training, but it moves range failures into the gradient stream where a wrong threshold or a stale flag silently changes the update.

explain
implement

You can now

Explain what TF32, fp16 and bf16 change about range, precision and accumulation, and state the current defaults.
Implement an autocast forward with a gradient scaler, including save and resume of scaler state.
Apply unscale before clipping with an unchanged threshold and diagnose skipped steps from the scale.

Covers

FP32, TF32-style settings, FP16 and BF16
Autocast, gradient scaling and unscale-before-clip order

Study, in this order

4 steps

Automatic Mixed Precision recipe
PyTorch - reliability and performance recipes

Adding torch.autocast with the fp16 linear output and fp32 loss asserts; Adding GradScaler with scale, backward, step and update; Inspecting/modifying gradients with unscale_ before clip_grad_norm_ at max_norm 0.1; Saving/Resuming with scaler state_dict

Done
TF32 precision controls
PyTorch - checkpointing, numerical accuracy and CUDA semantics

TensorFloat-32 on Ampere and later devices: the new fp32_precision settings per backend and operator, matmul versus cuDNN defaults, and the rule against mixing old and new APIs

Done
Reduced-precision accumulation
PyTorch - checkpointing, numerical accuracy and CUDA semantics

Reduced Precision Reduction for FP16 and BF16 GEMMs with the off switches; Reduced Precision Reduction for FP16 and BF16 in SDPA with the fp32 upcast default

Done
HOML: why float16 underflows and how mixed precision fixes it
Hands-On Machine Learning with Scikit-Learn and PyTorch

Appendix B, Reduced Precision Models and Mixed-Precision Training (printed pp. 798-802): float16 and bfloat16 ranges, the 123 + 0.045 rounding example, fp32 primary weights, and the train_mpt loop with autocast and GradScaler

Depth · One correction: p. 801 says GradScaler also scales down the learning rate. It does not; it unscales the gradients before the step and skips the step when they hold inf or NaN.

Done

Focus on

Autocast chooses per operation: matmul-like ops go narrow, reductions stay wide. A fp16 linear output with a fp32 loss is the contract working, not a bug.
The scaler exists because small fp16 gradients flush to zero. It scales the loss, unscales before the optimizer sees grads, and skips the step on inf or NaN.
Current defaults are full precision for CUDA matmul and TF32 allowed for cuDNN convolutions. Set them with fp32_precision, never by mixing allow_tf32 with the new API.
The recipe assumes a CUDA GPU. The CPU fallback is autocast with bf16 and no scaler, with its own dtype asserts.

Reproduce closed-book

On a CUDA GPU, run a fp16 autocast forward over one linear layer and assert the output is fp16 while an mse loss over it is fp32. Construct torch.amp.GradScaler('cuda'), confirm the initial scale reads 65536.0, then run scale, backward, unscale_, clip_grad_norm_ and step plus update without error.

Traps

Mixing allow_tf32 and fp32_precision. In PyTorch 2.14, reading allow_tf32 after setting fp32_precision raises a RuntimeError, while the reverse just reports the mapped value. The docs do not support mixing, so pick one API and stay on it.
Clipping .grad between backward and scaler.step without unscale_ first. The threshold then applies to scaled values and silently clips the wrong updates.
Copying CUDA fp16 settings to CPU. The recipe requires a CUDA GPU; on CPU use bf16 autocast and do not expect the fp16 asserts to transfer.

Practice

PyTorch AMP recipe: vary participating sizes

The Exercise above the batch_size cell: rerun Default Precision and All together with batch sizes such as 128, 256 and 513, using the recipe's synchronized timers

Report the speedup per size on your GPU. Compare the 513 run with the Troubleshooting advice that matmul sizes be multiples of 8, and say whether your GPU and library versions still show that effect. Then add the unscale_ plus clip and scaler save and resume sections to the same loop.

TrenTorch: Loss Scaling

Implement scale_loss, unscale_gradients and has_inf_or_nan; the last test walks one dynamic-scaling step

Map each function onto GradScaler: scale(loss), unscale_(optimizer), the inf or NaN check that skips step(), and update(), which halves the scale after an overflow and doubles it after 2000 clean steps.

Self-test

What are the current defaults for TF32 on CUDA matmul and cuDNN convolutions?
Code sets torch.backends.cuda.matmul.fp32_precision to tf32, and a library later reads allow_tf32. What happens?
Scaled gradients sit in .grad. Where does clipping go, and does max_norm change?
Why does fp16 training need a gradient scaler while bf16 training usually does not?

Completion criterion

Stop when

Your autocast plus scaler loop asserts the right dtypes, clips at an unchanged threshold, and you can explain any skipped step from the scale history.

Mark passed, then open 8.5 Only mark this when you can do it without notes.

Leads to

8.5 Overflow, underflow and finding the first invalid value Trace the overflow and underflow that scaling is designed to prevent.
8.9 Profiling training phases and verifying optimizations Measure whether narrower math actually moved the bottleneck.

Comes after

1.8 Batched shapes, operation costs and conditioningUse conditioning and arithmetic precision to reason about mixed-precision risk.
5.6 Floating point, rounding, overflow and cancellationApply range-versus-precision reasoning when selecting dtypes.
5.7 Stable log-sum-exp, online reductions and mixed precisionApply the distinction between storage precision and accumulation precision.
3.6 Schedules, gradient accumulation and clippingSee the full mixed-precision scaler contract behind the clipping order.
8.3 GPU memory: allocated, reserved, peak and memory modelsHalve the parameter, gradient and activation terms with precision choice.

Resources for this unit

PyTorch - reliability and performance recipes
PyTorchprimarycourse
Use the Profiler, Automatic Mixed Precision, model loading and Benchmark recipes, then the Performance Tuning Guide. Apply each to one compact model from an earlier section and keep a CPU fallback when CUDA is unavailable.
PyTorch - checkpointing, numerical accuracy and CUDA semantics
PyTorchrequiredreference
Use the official user-guide entries for torch.utils.checkpoint, numerical accuracy, CUDA semantics, reproducibility and serialization. Predict memory categories before measuring and define resume equivalence before running.
Hands-On Machine Learning with Scikit-Learn and PyTorch
Aurelien Geron (O'Reilly, 2025)supporttheory
Read Appendix B, Reduced Precision Models and Mixed-Precision Training, for why float16 training underflows and what fp32 primary weights and loss scaling fix.

Also in this section

PyTorch - compiler troubleshooting
PyTorchrequiredreference
Read graph breaks, guards, recompilations and logging. Compile one stable fragment, intentionally trigger a graph break or recompilation, inspect it, and compare eager and compiled outputs and gradients.
Deep Learning Training Debugger
Roadmap-Finalreferenceassessment
Use the failure-injection, measurement and evidence rules as the deeper standard. The full Tier 2 project is a later route and is not added to this stage milestone.
TensorTonic - mixed-precision memory and transformer FLOPs
TensorTonicsupportpractice
Solve only Mixed-Precision Training Memory Accountant and Transformer Training FLOP Estimator. Reconcile both with the compact model's profiler and document every included and excluded term.
Learn PyTorch - profiler walkthrough
Daniel Bourkeoptionalimplementation
Use notebook 10 only if the official profiler route remains unclear. Transfer the relevant annotations into the compact model instead of keeping the tutorial code.
TrenTorch - precision, checkpointing and resume
TrenTorchsupportpractice
Solve Loss Scaling, FP16/BF16 Representable Range, Gradient Checkpointing and Resume-from-checkpoint in NumPy, then reproduce each behavior with the PyTorch API in the compact model.
Python os - `os.replace`
Python Software Foundationreferencereference
Read os.replace for the atomic-rename step of checkpoint writes: save to a temp file in the same directory, flush and fsync, then replace the final name.