Mixed precision training (AMP) in PyTorch

Introduction

Mixed precision training helps your model learn faster and use less memory by mixing small and normal numbers smartly.

Training deep learning models on GPUs with limited memory.

Speeding up training without losing model accuracy.

Running large models that otherwise don't fit in GPU memory.

Reducing electricity and hardware costs during training.

Syntax

PyTorch

import torch
from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()

for data, target in dataloader:
    optimizer.zero_grad()
    with autocast():
        output = model(data)
        loss = loss_fn(output, target)
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

autocast() runs operations in mixed precision automatically.

GradScaler helps keep training stable by scaling gradients.

Examples

This runs the forward pass and loss calculation in mixed precision.

PyTorch

with autocast():
    output = model(data)
    loss = loss_fn(output, target)

This scales the loss, does backpropagation, updates weights, and adjusts scaling.

PyTorch

scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

Sample Model

This code trains a simple model for 3 epochs using mixed precision. It prints the loss each epoch.

PyTorch

import torch
import torch.nn as nn
import torch.optim as optim
from torch.cuda.amp import autocast, GradScaler

# Simple model
class SimpleModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(10, 1)
    def forward(self, x):
        return self.linear(x)

# Data
x = torch.randn(20, 10).cuda()
y = torch.randn(20, 1).cuda()

model = SimpleModel().cuda()
optimizer = optim.SGD(model.parameters(), lr=0.1)
loss_fn = nn.MSELoss()
scaler = GradScaler()

model.train()
for epoch in range(3):
    optimizer.zero_grad()
    with autocast():
        output = model(x)
        loss = loss_fn(output, y)
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()
    print(f"Epoch {epoch+1}, Loss: {loss.item():.4f}")

Important Notes

Mixed precision works best on GPUs with Tensor Cores (like NVIDIA RTX series).

Always use GradScaler to avoid problems with very small gradients.

Check your model's accuracy to ensure mixed precision does not reduce it.

Summary

Mixed precision training speeds up learning and saves memory by using smaller numbers where possible.

Use autocast() for automatic mixed precision during forward pass.

Use GradScaler to keep training stable when using mixed precision.

Practice

1. What is the main benefit of using mixed precision training (AMP) in PyTorch?

easy

A. It replaces the need for GPUs during training.

B. It guarantees 100% accuracy in model predictions.

C. It speeds up training and reduces memory usage by using smaller number formats.

D. It automatically tunes hyperparameters for better performance.

Solution

1. Step 1: Understand mixed precision training purpose

   Mixed precision training uses smaller floating point numbers (like float16) to speed up computation and reduce memory use.

2. Step 2: Identify the main benefit

   This leads to faster training and less memory consumption without losing much accuracy.

3. Final Answer:

   It speeds up training and reduces memory usage by using smaller number formats. -> Option C

4. Quick Check:

   Mixed precision = faster training + less memory [OK]

Hint: Mixed precision means faster and lighter training [OK]

Common Mistakes:

• Thinking AMP guarantees perfect accuracy
• Believing AMP removes need for GPUs
• Confusing AMP with hyperparameter tuning

2. Which of the following is the correct way to use automatic mixed precision (AMP) in a PyTorch training loop?

easy

A. with torch.autocast(device_type='cuda', dtype=torch.float16): output = model(input)

B. with torch.autocast(device_type='cpu'): output = model(input)

C. with torch.autocast(device='cpu'): output = model(input)

D. with torch.autocast(dtype=torch.int32): output = model(input)

Solution

1. Step 1: Recall correct autocast usage

   The correct syntax requires specifying device_type (usually 'cuda') and dtype (usually torch.float16) for mixed precision.

2. Step 2: Check each option

   with torch.autocast(device_type='cuda', dtype=torch.float16): output = model(input) correctly uses device_type='cuda' and dtype=torch.float16. B uses wrong device_type='cpu', C uses invalid parameter 'device', D uses unsupported dtype.

3. Final Answer:

   with torch.autocast(device_type='cuda', dtype=torch.float16): output = model(input) -> Option A

4. Quick Check:

   Correct autocast needs device_type and dtype [OK]

Hint: Always specify device_type='cuda' and dtype=torch.float16 [OK]

Common Mistakes:

• Omitting device_type parameter
• Using CPU device for autocast
• Using unsupported dtypes like int32

3. Consider this PyTorch code snippet using AMP:

scaler = torch.cuda.amp.GradScaler()
for input, target in dataloader:
    optimizer.zero_grad()
    with torch.cuda.amp.autocast():
        output = model(input)
        loss = loss_fn(output, target)
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

What is the role of scaler.scale(loss).backward() here?

medium

A. It scales the loss to prevent underflow during backward pass with mixed precision.

B. It converts the loss to float32 before backward pass.

C. It disables gradient computation for the loss.

D. It resets the gradients to zero.

Solution

1. Step 1: Understand GradScaler's purpose

   GradScaler scales the loss to avoid very small gradient values (underflow) during mixed precision backward pass.

2. Step 2: Analyze scaler.scale(loss).backward()

   This scales the loss before calling backward(), ensuring stable gradient computation.

3. Final Answer:

   It scales the loss to prevent underflow during backward pass with mixed precision. -> Option A

4. Quick Check:

   GradScaler scales loss to keep gradients stable [OK]

Hint: GradScaler scales loss to keep gradients stable [OK]

Common Mistakes:

• Thinking it converts loss to float32
• Believing it disables gradients
• Confusing it with zeroing gradients

4. You have this training loop using AMP but encounter a runtime error related to gradients:

scaler = torch.cuda.amp.GradScaler()
for input, target in dataloader:
    optimizer.zero_grad()
    with torch.cuda.amp.autocast():
        output = model(input)
        loss = loss_fn(output, target)
    loss.backward()
    scaler.step(optimizer)
    scaler.update()

What is the error and how to fix it?

medium

A. Error: optimizer.zero_grad() should be called after backward().

B. Error: loss.backward() should be replaced with scaler.scale(loss).backward() to scale gradients.

C. Error: autocast() should be outside the loop.

D. Error: scaler.step() should be called before backward().

Solution

1. Step 1: Identify missing scaler scaling

   The code calls loss.backward() directly without scaling, which causes gradient underflow errors in AMP.

2. Step 2: Correct usage with GradScaler

   Replace loss.backward() with scaler.scale(loss).backward() to properly scale gradients before backward pass.

3. Final Answer:

   loss.backward() should be replaced with scaler.scale(loss).backward() to scale gradients. -> Option B

4. Quick Check:

   Use scaler.scale(loss).backward() with AMP [OK]

Hint: Always scale loss before backward with GradScaler [OK]

Common Mistakes:

• Calling loss.backward() without scaling
• Calling zero_grad() after backward
• Misplacing autocast context

5. You want to train a model on a GPU using mixed precision to save memory. Which of the following code snippets correctly combines AMP with gradient clipping to avoid exploding gradients?

hard

A. with torch.cuda.amp.autocast():
output = model(input)
loss = loss_fn(output, target)
scaler.scale(loss).backward()
optimizer.zero_grad()
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
scaler.step(optimizer)
scaler.update()

B. with torch.cuda.amp.autocast():
output = model(input)
loss = loss_fn(output, target)
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
optimizer.step()

C. with torch.cuda.amp.autocast():
output = model(input)
loss = loss_fn(output, target)
scaler.scale(loss).backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
scaler.step(optimizer)
scaler.update()

D. with torch.cuda.amp.autocast():
output = model(input)
loss = loss_fn(output, target)
scaler.scale(loss).backward()
scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
scaler.step(optimizer)
scaler.update()

Solution

1. Step 1: Understand AMP with gradient clipping

   When using AMP, gradients must be unscaled before clipping to get correct values.

2. Step 2: Check code correctness

   with torch.cuda.amp.autocast(): output = model(input) loss = loss_fn(output, target) scaler.scale(loss).backward() scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) scaler.step(optimizer) scaler.update() correctly calls scaler.unscale_(optimizer) before clipping gradients, then steps optimizer and updates scaler.

3. Step 3: Identify errors in other options

   with torch.cuda.amp.autocast(): output = model(input) loss = loss_fn(output, target) loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) optimizer.step() lacks scaler usage and uses loss.backward() directly. with torch.cuda.amp.autocast(): output = model(input) loss = loss_fn(output, target) scaler.scale(loss).backward() torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) scaler.step(optimizer) scaler.update() clips gradients before unscale, which is incorrect. with torch.cuda.amp.autocast(): output = model(input) loss = loss_fn(output, target) scaler.scale(loss).backward() optimizer.zero_grad() torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) scaler.step(optimizer) scaler.update() calls optimizer.zero_grad() after backward, which is wrong.

4. Final Answer:

   Correctly unscales gradients before clipping and uses scaler properly. -> Option D

5. Quick Check:

   Unscale before clip_grad_norm with scaler [OK]

Hint: Unscale gradients before clipping with GradScaler [OK]

Common Mistakes:

• Clipping gradients before unscale
• Not using scaler with backward
• Calling zero_grad after backward