Mixed precision training (AMP) in PyTorch - Practice Problems & Coding Challenges

Challenge - 5 Problems

🧠 Conceptual

intermediate

Why use mixed precision training (AMP)?

What is the main benefit of using Automatic Mixed Precision (AMP) in PyTorch training?

A It increases model accuracy by using higher precision calculations everywhere.

B It reduces memory usage and speeds up training by using float16 where possible.

C It automatically tunes hyperparameters during training.

D It converts the model to run on CPU instead of GPU.

💡 Hint

Think about how using smaller number formats affects speed and memory.

AMP uses float16 precision where safe, reducing memory and speeding up training without losing accuracy.

❓ Predict Output

intermediate

Output of AMP training step snippet

What will be the printed loss value type after this AMP training step?

PyTorch

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

model = torch.nn.Linear(2, 1).cuda()
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
scaler = GradScaler()

inputs = torch.tensor([[1.0, 2.0]], device='cuda')
target = torch.tensor([[1.0]], device='cuda')

optimizer.zero_grad()
with autocast():
    output = model(inputs)
    loss = torch.nn.functional.mse_loss(output, target)
print(type(loss))

A <class 'torch.Tensor'> with dtype=torch.float32

B <class 'torch.Tensor'> with dtype=torch.float16

C <class 'torch.Tensor'> with dtype=torch.float64

D RuntimeError due to dtype mismatch

💡 Hint

AMP uses float16 for some ops but loss is usually float32 for stability.

Loss tensors are kept in float32 even inside autocast to maintain numerical stability.

❓ Model Choice

advanced

Choosing model parts for AMP

Which part of a model should NOT be wrapped inside autocast() for AMP training?

A The forward pass of the neural network layers

B The loss calculation function

C The data loading pipeline

D The optimizer step call

💡 Hint

Consider which operations benefit from mixed precision and which do not.

Optimizer step updates weights and should not be inside autocast; autocast is for forward and loss computations.

❓ Hyperparameter

advanced

Effect of GradScaler in AMP

What is the role of GradScaler in PyTorch AMP training?

A It converts all model weights to float16 permanently.

B It automatically adjusts learning rate during training.

C It scales the loss to prevent underflow in gradients during backpropagation.

D It disables gradient computation for certain layers.

💡 Hint

Think about why gradients might vanish when using float16.

GradScaler multiplies loss by a scale factor to keep gradients in a safe range, avoiding underflow in float16.

🔧 Debug

expert

Debugging AMP training error

Given this AMP training snippet, what error will occur and why?

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

model = torch.nn.Linear(2, 1).cuda()
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
scaler = GradScaler()

inputs = torch.tensor([[1.0, 2.0]], device='cuda')
target = torch.tensor([[1.0]], device='cuda')

optimizer.zero_grad()
with autocast():
    output = model(inputs)
    loss = torch.nn.functional.mse_loss(output, target)
loss.backward()
scaler.step(optimizer)
scaler.update()

A RuntimeError: You must call scaler.scale(loss).backward() instead of loss.backward()

B TypeError: optimizer.step() missing required positional argument

C No error, code runs successfully

D RuntimeError: autocast context must wrap optimizer.step()

💡 Hint

Check how gradients are computed when using GradScaler.

When using GradScaler, loss.backward() must be replaced by scaler.scale(loss).backward() to scale gradients properly.