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.

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