Mixed Precision Training with PyTorch AMP

Training modern deep learning models often pushes the limits of your hardware, both in terms of computation time and memory. One effective technique to alleviate these pressures is mixed precision training. This approach combines the use of lower-precision floating-point numbers (like 16-bit floats) for certain parts of your model with higher-precision numbers (like 32-bit floats) for others, aiming to speed up training and reduce memory usage without a significant loss in model accuracy.

What is Mixed Precision Training?

Traditionally, most neural network training has been performed using 32-bit floating-point numbers (FP32 or single-precision). While FP32 offers a wide dynamic range and good precision, computations with it can be slower and require more memory compared to lower-precision formats.

Mixed precision training intelligently uses a combination of:

• FP16 (Half-Precision): Uses 16 bits to represent numbers. Operations in FP16 are significantly faster on modern GPUs, especially those with specialized hardware like NVIDIA's Tensor Cores. FP16 also halves the memory required for storing weights, activations, and gradients.
• FP32 (Single-Precision): Continues to be used for operations that require higher precision to maintain numerical stability and model accuracy, such as weight updates or certain sensitive layers.

The goal is to reap the benefits of FP16 (speed, memory) while mitigating its potential downsides, such as a smaller representable range which can lead to overflow (numbers becoming too large) or underflow (gradients becoming zero).

PyTorch AMP: torch.cuda.amp

PyTorch provides convenient tools for automatic mixed precision (AMP) training primarily through the torch.cuda.amp module. This module automates most of the process, making it relatively straightforward to integrate into your existing training scripts. The two main components you'll interact with are autocast and GradScaler.

torch.cuda.amp.autocast

The autocast context manager is the workhorse for selecting which operations run in FP16 and which remain in FP32. When you enable autocast for a region of your code (typically the forward pass), it automatically casts inputs of eligible PyTorch operations to FP16.

Which operations are eligible?

• Operations that benefit most from FP16, like convolutions and matrix multiplications (linear layers), are typically run in FP16.
• Operations that might suffer from numerical instability in FP16, such as reductions (like sum) or normalization layers (like BatchNorm), are often kept in FP32 to maintain precision.

autocast handles these conversions on the fly. For example:

# model and data are on CUDA
model = MyModel().cuda()
input_data = torch.randn(N, C, H, W, device="cuda")

# Enable autocast for the forward pass
with torch.cuda.amp.autocast():
    output = model(input_data)
    loss = loss_fn(output, target) # loss computation also under autocast

# Gradients from this loss will be FP16
# loss.backward() # (we'll see how GradScaler modifies this)

Inside the autocast block, operations like model(input_data) will internally use FP16 for many computations if input_data is a CUDA tensor and the operations are deemed safe and efficient in FP16 by PyTorch. The loss tensor produced will also likely be FP16.

torch.cuda.amp.GradScaler

While autocast handles the forward pass, using FP16 for gradients during the backward pass can lead to underflow. Gradients, especially for deep networks or small parameter updates, can become very small. If these small values fall below the minimum representable positive number in FP16, they become zero, effectively halting learning for those parameters.

GradScaler helps prevent this by scaling the loss before the backward pass. Here's the process:

1. Scale Loss: The loss computed (potentially in FP16 under autocast) is multiplied by a large scaling factor. This inflates the magnitude of the subsequent gradients.
2. Backward Pass: The backward() call is performed on this scaled loss. The resulting gradients are also scaled. Because they are larger, they are less likely to underflow in FP16.
3. Unscale Gradients: Before the optimizer updates the model weights, GradScaler un-scales the gradients by dividing them by the same scaling factor. This restores them to their correct magnitude.
4. Optimizer Step: The optimizer then uses these un-scaled (but now numerically stable) gradients to update the weights.
5. Update Scale Factor: GradScaler dynamically adjusts the scaling factor. If gradients overflow (become inf or NaN) during a step, it means the scaling factor was too high, so it's reduced for the next iteration, and the optimizer step for the current iteration is skipped. If no overflows occur for a certain number of steps, the scaling factor may be increased to further improve precision.

Integrating AMP into Your Training Loop

Let's see how to modify a standard PyTorch training loop to use AMP.

Typical Training Loop (FP32):

import torch
import torch.nn as nn
import torch.optim as optim

# Assume model, data_loader, loss_fn are defined
# model = MyModel().cuda()
# optimizer = optim.Adam(model.parameters(), lr=1e-3)
# loss_fn = nn.CrossEntropyLoss()

# for epoch in range(num_epochs):
#     for input_batch, target_batch in data_loader:
#         input_batch, target_batch = input_batch.cuda(), target_batch.cuda()

#         optimizer.zero_grad()

#         outputs = model(input_batch)
#         loss = loss_fn(outputs, target_batch)

#         loss.backward()
#         optimizer.step()
#     print(f"Epoch {epoch+1} completed.")

Training Loop with PyTorch AMP:

import torch
import torch.nn as nn
import torch.optim as optim
from torch.cuda.amp import autocast, GradScaler # Import AMP components

# Assume model, data_loader, loss_fn are defined
# model = MyModel().cuda()
# optimizer = optim.Adam(model.parameters(), lr=1e-3)
# loss_fn = nn.CrossEntropyLoss()

scaler = GradScaler() # Initialize GradScaler

# for epoch in range(num_epochs):
#     for input_batch, target_batch in data_loader:
#         input_batch, target_batch = input_batch.cuda(), target_batch.cuda()

#         optimizer.zero_grad()

#         # Forward pass with autocasting
#         with autocast():
#             outputs = model(input_batch)
#             loss = loss_fn(outputs, target_batch)

#         # Scale loss and call backward() on scaled loss
#         scaler.scale(loss).backward()

#         # Unscale gradients and call optimizer.step()
#         scaler.step(optimizer)

#         # Update the scale for next iteration
#         scaler.update()
#     print(f"Epoch {epoch+1} completed with AMP.")

The main changes are:

1. Initializing GradScaler().
2. Wrapping the forward pass and loss calculation in the with autocast(): block.
3. Using scaler.scale(loss).backward() instead of just loss.backward().
4. Using scaler.step(optimizer) instead of optimizer.step().
5. Calling scaler.update() after scaler.step(optimizer).

The optimizer.zero_grad() call can be placed before the autocast block or at the beginning of the loop as usual. PyTorch recommends setting gradients to None instead of zeroing them for a small performance benefit, which can be done via optimizer.zero_grad(set_to_none=True).

Visualizing the Impact

The primary benefits of AMP are reduced training time and memory usage. While exact numbers vary by model, GPU, and batch size, the improvements can be substantial.

This chart illustrates potential reductions in training time and peak memory usage when switching from standard FP32 training to Automatic Mixed Precision (AMP).

Connection to Your TensorFlow Experience

If you've used TensorFlow, you might be familiar with tf.keras.mixed_precision. The concept is very similar:

• TensorFlow uses a Policy (e.g., mixed_precision.set_global_policy('mixed_float16')) to define how layers should handle mixed precision, somewhat analogous to how autocast implicitly determines types.
• TensorFlow's LossScaleOptimizer wraps an existing optimizer to perform loss scaling, akin to PyTorch's GradScaler.

Both frameworks aim to simplify the adoption of mixed precision by automating type casting and loss scaling. The underlying principles are the same, though the specific API calls and implementation details differ. PyTorch's autocast and GradScaler offer a flexible way to apply mixed precision, often requiring only a few lines of code modification.

Important Considerations

• Hardware Support: AMP with FP16 provides the most significant speedups on NVIDIA GPUs with Tensor Core support (Volta, Turing, Ampere architectures and newer). While it can run on other GPUs, the performance gains might be less pronounced. Some newer GPUs also support BFloat16 (BF16), which autocast can also target if dtype=torch.bfloat16 is specified. BF16 has a similar range to FP32 but less precision, often offering a good balance without requiring loss scaling as frequently as FP16.
• Batch Normalization: Layers like BatchNorm typically maintain their weights and perform computations in FP32 even within an autocast block to ensure stability. PyTorch handles this automatically.
• Gradient Clipping: If you use gradient clipping, it should be applied after gradients are un-scaled by GradScaler but before optimizer.step().

  # scaler.scale(loss).backward()
  # # Unscale the gradients of optimizer's assigned params in-place
  # scaler.unscale_(optimizer)
  # # Clip gradients after unscaling
  # torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)
  # scaler.step(optimizer)
  # scaler.update()
  

• Saving and Loading Checkpoints: If you save and load checkpoints, remember to also save and restore the state of the GradScaler:

  # Saving
  checkpoint = {
      'model_state_dict': model.state_dict(),
      'optimizer_state_dict': optimizer.state_dict(),
      'scaler_state_dict': scaler.state_dict(),
      # ... other things
  }
  torch.save(checkpoint, 'my_checkpoint.pth')
  
  # Loading
  # checkpoint = torch.load('my_checkpoint.pth')
  # model.load_state_dict(checkpoint['model_state_dict'])
  # optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
  # scaler.load_state_dict(checkpoint['scaler_state_dict'])
  

• Numerical Stability: While AMP is designed to be, in rare cases, some specific operations or model architectures might still exhibit numerical issues. If you suspect problems, you can try disabling autocast for certain parts of your model or explicitly casting specific operations to FP32.

By leveraging PyTorch's torch.cuda.amp, you can often achieve significant training speedups and memory savings with minimal code changes, allowing you to train larger models or iterate faster on your existing ones. This is a valuable tool for any TensorFlow developer transitioning to PyTorch and looking to optimize their training workflows.