Mixed precision vs FP32: how to train a neural network

The usual culprit is straightforward: every matrix multiplication, activation, and gradient is processed in FP32 even though modern GPUs are designed to execute many of those operations faster in 16-bit formats.

Mixed precision training addresses that mismatch. It runs numerically tolerant operations in FP16 or BF16, while keeping sensitive values—such as master weights, optimizer states, and selected reductions—in FP32. On architectures with Tensor Cores, the result can be up to a 3x speedup for compute-heavy models, alongside a substantial reduction in the memory used by weights and activations.

That does not mean we can replace every float32 tensor with float16 and call the job optimized. The gotcha is numerical stability. FP16 has a narrow dynamic range, gradients can underflow, and a training loop that looks correct may quietly stop learning. We need the right precision policy, the right hardware, and a few sanity checks in the loop.

FP32, FP16, and BF16: what actually changes

The difference between these formats is not just the number of bits. The distribution of those bits determines whether a format can represent very large values, very small values, or fine-grained numerical differences.

FP32 uses 32 bits:

• 1 sign bit
• 8 exponent bits
• 23 mantissa bits

It offers a broad dynamic range and relatively high precision. That is why FP32 remains the safe default for optimizer updates, master weights, many reductions, and operations where small numerical errors can accumulate.

FP16 uses 16 bits:

• 1 sign bit
• 5 exponent bits
• 10 mantissa bits

Its maximum representable value is approximately 65,504. Values below roughly 6.1e-5 can underflow, depending on the operation and representation. That range is often sufficient for activations and matrix multiplications, but it is not a reliable place to store every gradient produced during backpropagation.

BF16 also uses 16 bits:

• 1 sign bit
• 8 exponent bits
• 7 mantissa bits

The important property is its exponent. BF16 has approximately the same dynamic range as FP32—up to around 3.4e38—but less precision in the mantissa. In practice, BF16 is usually more forgiving than FP16 when a model produces very large or very small intermediate values.

The practical pattern is mixed precision rather than pure half-precision training. We use FP16 or BF16 for operations that benefit from Tensor Core acceleration, and retain FP32 where the computation is sensitive.

A typical training run may therefore contain:

• FP32 master weights used for stable parameter updates.
• FP16 or BF16 activations during the forward pass.
• FP16 or BF16 matrix multiplications and convolutions.
• FP32 optimizer states, particularly with optimizers such as Adam.
• FP32 handling for loss calculation, Softmax, Batch Normalization, and selected reductions.

Mixed precision is not “make the model half precision.” It is a controlled allocation of numerical risk.

This distinction matters when estimating memory. We can reduce the footprint of weights and activations by up to roughly 50%, but optimizer states do not disappear. An Adam-based model can still carry FP32 state tensors for each parameter, so the total memory reduction will not necessarily equal the reduction visible in the activation tensors.

Why mixed precision training accelerates GPU throughput

The speedup comes from hardware utilization, not from a software shortcut. Modern NVIDIA GPUs include Tensor Cores that are optimized for low-precision matrix operations. When the model presents compatible tensor shapes and datatypes, those units can process more arithmetic per clock than the traditional FP32 execution path.

This is most visible in workloads dominated by:

• Transformer attention and feed-forward layers.
• Convolutional backbones.
• Large matrix multiplications.
• Language-model pretraining and fine-tuning.
• Batch-heavy workloads with enough parallelism to keep the GPU occupied.

NVIDIA documentation reports speedups of up to 3x for arithmetically intense architectures when mixed precision is used with Tensor Cores. That figure is a ceiling, not a promise for every repository or custom model. A small network, irregular kernel, input pipeline bottleneck, or poorly shaped batch may show little improvement.

The first sanity check is therefore not the final epoch time. It is GPU utilization. If the GPU is waiting on tokenization, image decoding, dataloader workers, or host-to-device copies, changing FP32 to BF16 will not fix the main bottleneck.

A useful benchmark compares the same configuration across precision modes:

Shape and layout still matter. A model may contain a few large GEMMs that accelerate well, but spend substantial time in operations that remain memory-bound or execute in FP32. Framework overhead, synchronization points, custom CUDA kernels, and frequent tensor conversions can also dilute the gain.

For that reason, we should record more than samples per second. A useful pytorch amp benchmark tracks:

1. Step time — measured after warm-up rather than on the first few iterations.

2. Throughput — tokens per second, images per second, or examples per second.

3. Peak allocated and reserved memory — both can reveal different bottlenecks.

4. GPU utilization — to separate compute pressure from input-pipeline starvation.

5. Loss and validation metrics — speed is irrelevant if convergence changes materially.

6. Overflow or skipped steps — especially with FP16 and dynamic scaling.

The comparison should use the same batch size, sequence length, data order policy, optimizer, learning-rate schedule, gradient accumulation, and evaluation frequency. Otherwise, we are measuring a configuration change rather than a precision change.

The FP16 gotcha: gradients can disappear

FP16’s narrow exponent range is the main engineering problem. During backpropagation, a gradient can be mathematically valid but too small to represent in FP16. When that happens, it underflows to zero.

The model then receives no update for that parameter on that step. Nothing necessarily crashes. The loss may continue to print. The optimizer may continue to run. We can end up debugging a training curve that looks merely slow when the real issue is silent gradient loss.

Loss scaling is the standard workaround. Before backpropagation, we multiply the loss by a scale factor. This shifts small gradients into a representable range. After gradients are computed, the scaling factor is removed before the optimizer updates the weights.

PyTorch’s GradScaler manages this process dynamically. It can increase the scale when recent steps are stable and reduce it when an overflow is detected. This is preferable to choosing a fixed value and hoping it works across every model, batch, and learning-rate schedule.

The logic is:

1. Compute the forward pass under an autocast context.

2. Calculate the loss.

3. Scale the loss before calling backward().

4. Unscale gradients before clipping or inspecting them.

5. Check for non-finite gradients.

6. Apply the optimizer step only when the gradients are valid.

7. Update the scaler for the next iteration.

Gradient clipping has one implementation detail that is easy to miss. If we clip before unscaling, we are clipping the artificially enlarged gradients. The correct order is to call the scaler’s unscale operation first, then apply clipping to the original gradient magnitude.

A robust FP16 loop should therefore treat skipped optimizer steps as a signal, not an annoyance. Occasional skips during warm-up may be recoverable. Persistent skips usually point to a more serious issue:

• The learning rate is too high.
• A custom operation is numerically unstable.
• The loss is already non-finite before scaling.
• A reduction or normalization is running in an unsafe dtype.
• The input data contains extreme values.
• The model is producing exploding activations.

BF16 changes this trade-off. Its 8-bit exponent gives it the same broad dynamic range as FP32, so it generally does not require dynamic loss scaling. That makes the training loop simpler and often more stable on supported hardware. The reduced mantissa precision still matters, but BF16 is less exposed to the specific underflow problem that makes FP16 fragile.

Implementing PyTorch AMP without rewriting the model

The good news is that PyTorch Automatic Mixed Precision usually does not require a model rewrite. We wrap the forward pass in torch.autocast, select the target dtype, and use a scaler for FP16.

The exact scaler constructor varies across PyTorch releases, so it is worth checking the installed version before copying a repository’s boilerplate. Current code commonly uses torch.amp.GradScaler("cuda"), while older projects often use torch.cuda.amp.GradScaler().

The core structure looks like this in compact form:

1. Move the batch to the selected device with non-blocking transfers where the input pipeline supports them.

2. Clear gradients with optimizer.zero_grad(set_to_none=True).

3. Enter torch.autocast(device_type="cuda", dtype=torch.float16) or use torch.bfloat16.

4. Run outputs = model(inputs).

5. Compute the loss inside the autocast region unless a specific loss requires FP32.

6. For FP16, call scaler.scale(loss).backward().

7. Call scaler.unscale_(optimizer) before gradient clipping or diagnostics.

8. Apply scaler.step(optimizer) and then scaler.update().

9. For BF16, use the same autocast structure but normally call loss.backward() and optimizer.step() directly.

The important part is not the API spelling. It is the boundary around the forward and backward computations.

A typical FP16 iteration can be represented with inline operations such as with torch.autocast(device_type="cuda", dtype=torch.float16):, followed by scaled_loss = scaler.scale(loss), scaled_loss.backward(), scaler.unscale_(optimizer), and scaler.step(optimizer). Keeping these calls in the expected order prevents several common implementation errors.

For BF16, the corresponding policy is usually dtype=torch.bfloat16. We still keep the model’s master parameters and optimizer state in FP32. Autocast decides which eligible operations can use BF16 and which should remain in FP32.

What autocast handles—and what it does not

Autocast is an operation-level policy. It does not blindly cast every tensor in the model. PyTorch maintains rules for operations that are generally safe or beneficial in reduced precision, while keeping sensitive operations in FP32.

That means we should not manually call .half() on the entire model as a first step. Manual casting can force layers into a dtype they were not designed to handle and can make debugging much harder. Autocast provides a safer default because it preserves higher precision where the framework expects it.

Operations commonly kept in FP32 include:

• Loss calculation in numerically sensitive cases.
• Softmax and probability normalization.
• Batch Normalization.
• Some reductions and accumulation operations.
• Custom functions without reliable reduced-precision behavior.

This does not remove the need to understand the model. Custom CUDA extensions, third-party layers, and unusual loss functions may not follow the same stability assumptions as standard PyTorch operators.

When a specific operation produces NaNs or unstable metrics, we can force that local section to FP32 rather than abandoning mixed precision across the whole model. The workaround is to disable autocast for the problematic block, convert its inputs to FP32 where required, and then return to the surrounding autocast context.

For example, a numerically sensitive block can run under torch.autocast(device_type="cuda", enabled=False), with the input explicitly converted through .float(). This is a targeted fix. It preserves the performance benefit for the rest of the network and gives us a clean place to investigate.

A practical debugging sequence

When a mixed-precision run fails, changing five settings at once usually hides the cause. We can isolate the issue more efficiently:

1. Run a short FP32 baseline. Confirm that the model, labels, loss, and optimizer are correct before introducing reduced precision.

2. Switch to autocast without other training changes. Keep the same seed policy, batch size, and learning-rate schedule.

3. Start with BF16 if the GPU supports it. Its dynamic range reduces the number of FP16-specific failure modes.

4. Check the first non-finite tensor. Inspect inputs, logits, loss, and gradients rather than only the final metric.

5. Inspect scaler behavior for FP16. Repeated scale reductions indicate overflow, not a normal performance fluctuation.

6. Disable autocast locally around the failing operation. Avoid converting the entire model back to FP32.

7. Compare convergence, not only loss at step zero. Similar initial loss does not guarantee similar training behavior.

The cleanest implementation keeps precision selection configurable. A command-line option or configuration value such as precision="fp32", precision="fp16", or precision="bf16" lets us reproduce failures and compare metrics without maintaining separate training scripts.

Hardware constraints decide the winner

Mixed precision is only useful when the hardware and software stack support the chosen format efficiently. FP16 acceleration is available across a broader range of GPU generations, but BF16 support is a more important dividing line.

BF16 is natively supported and accelerated on modern architectures including NVIDIA Ampere and newer. It should not be assumed to run with the same performance or capability on older Pascal or Volta GPUs. A framework may accept the dtype while silently falling back to slower kernels, or a particular operation may lack an optimized implementation.

Before selecting BF16, we should check:

• GPU architecture and compute capability.
• Installed CUDA runtime and driver compatibility.
• PyTorch version.
• Whether the relevant kernels support BF16.
• Whether the workload is large enough to benefit from Tensor Cores.
• Whether distributed training libraries support the selected dtype consistently.

A simple runtime check can confirm device availability and dtype support, but it cannot replace a short benchmark. The same model can show different results depending on sequence length, batch size, padding ratio, and whether attention kernels are optimized.

Distributed training adds another layer. Gradient communication may use a different dtype from local computation, and gradient buckets can affect memory pressure. Tensor parallel training, sharding, and checkpointing also interact with the precision policy. We should document which tensors are reduced precision and which remain FP32, particularly when resuming from checkpoints.

A checkpoint is another common gotcha. Saving only a reduced-precision model can make later fine-tuning less stable or complicate evaluation. For training resumes, preserve:

• FP32 master parameters where applicable.
• Optimizer state.
• Learning-rate scheduler state.
• GradScaler state for FP16.
• Current training step and random-state information.

The scaler state matters because it records the current scale and its adjustment history. Restarting with a default scale is not always wrong, but it can change the first part of the resumed run and make comparisons less reproducible.

Metrics that tell us whether the optimization worked

A claim such as “FP32 vs FP16 training speed improved” is incomplete without the measurement conditions. We need to separate throughput, memory, and convergence.

For throughput, use a warm-up period before recording timings. GPU operations are asynchronous, so measurements must synchronize the device around the timed region or use a framework-aware profiler. Otherwise, the host may report that an iteration finished before the GPU actually completed it.

For memory, record peak allocated memory and peak reserved memory. Allocated memory reflects tensors currently in use. Reserved memory includes blocks held by the caching allocator. Both are useful, but they answer different questions.

For model quality, compare the same neural network training metrics:

• Training loss at matched optimizer steps.
• Validation loss at matched evaluation points.
• Accuracy, F1, perplexity, BLEU, or the task-specific primary metric.
• Number of non-finite batches.
• Number of skipped optimizer steps.
• Final checkpoint quality after the same compute budget.

The phrase “without loss of model accuracy” should be treated as an empirical result for a particular model and configuration, not as a universal guarantee. Mixed precision is often transparent for standard architectures, but specialized domains and unstable objectives still require validation.

The most informative comparison uses two budgets:

• Matched steps: Does the reduced-precision run converge similarly when both runs process the same number of optimizer updates?
• Matched wall-clock time: Does the faster run reach a better validation metric within the same elapsed time?

These answer different engineering questions. The first isolates numerical behavior. The second measures the practical value of the optimization.

A small table in the experiment log is often enough:

Do not compare FP32 with a mixed-precision run that also has a larger batch size, fewer evaluations, or different gradient accumulation and then attribute the entire improvement to dtype. That is a benchmark design error, not a model result.

A decision framework for real training jobs

FP32 remains the right starting point when we are validating a new implementation, debugging a custom operation, or trying to establish a trustworthy baseline. It is slower and more memory-hungry, but it removes a large class of numerical variables.

FP16 mixed precision is useful when the GPU provides strong FP16 Tensor Core performance and the training stack already has reliable scaler support. It can deliver substantial throughput and activation-memory improvements, but the loss-scaling path must be treated as part of the algorithm.

BF16 is often the cleaner option on NVIDIA Ampere and newer hardware. It retains 16-bit storage and compute benefits while providing a much broader exponent range. The trade-off is lower mantissa precision and hardware dependence. It is not a drop-in assumption for older GPUs.

The choice can be summarized like this:

• Use FP32 to establish correctness and diagnose failures.
• Use FP16 AMP when hardware support is strong and the pipeline handles loss scaling correctly.
• Use BF16 AMP when the GPU supports it and you want a simpler stability profile.
• Keep optimizer states and master weights in FP32 unless the training method explicitly provides another validated policy.
• Measure the actual workload rather than relying on a headline speedup.

The same engineering discipline applies outside image and language models. For example, teams working on educational game-development systems may combine model training with specialized software pipelines, similar to the applied context described in this overview of game-development training with an industry partner. The surrounding application can change, but the implementation rule remains the same: benchmark the complete workload, not only the most attractive kernel.

The best precision mode is the one that improves throughput while preserving a training curve we can explain.

Closing the implementation loop

A reliable mixed-precision migration is incremental. First, make the FP32 model converge. Then add autocast. After that, introduce the scaler for FP16 or switch to BF16 on supported hardware. Finally, benchmark the end-to-end job and inspect the metrics that matter for deployment.

Before treating the change as complete, we should confirm that:

• The FP32 baseline reaches the expected loss and validation quality.
• The reduced-precision run uses the intended dtype on the intended GPU.
• FP16 uses dynamic loss scaling and unscales gradients before clipping.
• BF16 is not silently falling back to an unsupported or inefficient path.
• Peak memory and step time are measured after warm-up.
• Optimizer states remain in FP32 where required.
• Non-finite losses and gradients are detected explicitly.
• Checkpoints include scaler and optimizer state when resuming training.
• Validation metrics are compared under matched experimental conditions.
• Any FP32 fallback is local, documented, and justified by a stability check.

Mixed precision is one of the highest-leverage optimizations available in a PyTorch training loop, but it is not a magic dtype switch. The implementation works when precision becomes a deliberate policy: fast formats for high-throughput arithmetic, FP32 for numerical control, and benchmarks that prove the trade-off on the model we actually run.