Mixed precision training (AMP) in PyTorch - Cheat Sheet & Quick Revision

Recall & Review

beginner

What is mixed precision training in deep learning?

Mixed precision training uses both 16-bit and 32-bit floating point numbers during model training to speed up computation and reduce memory use, while keeping model accuracy.

beginner

What does AMP stand for in PyTorch?

AMP stands for Automatic Mixed Precision. It is a PyTorch feature that automatically manages mixed precision training to make it easier and safer.

beginner

Why use mixed precision training? List two benefits.

1. Faster training because 16-bit operations are quicker on modern GPUs.
2. Less GPU memory used, allowing bigger models or larger batches.

intermediate

How does PyTorch AMP help prevent underflow or overflow during training?

PyTorch AMP uses a technique called loss scaling. It multiplies the loss by a big number before backpropagation to keep gradients in a good range, then scales them back down.

intermediate

Show a simple PyTorch code snippet to enable AMP during training.

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

What is the main purpose of using mixed precision training?

A Make training more accurate by using 64-bit floats

B Increase model size without changing speed

C Speed up training and reduce memory use

D Avoid using GPUs

Mixed precision training speeds up training and reduces memory by using 16-bit floats where possible.

In PyTorch AMP, what does the GradScaler do?

A Scales the model weights

B Scales the loss to prevent gradient underflow

C Scales the input data

D Scales the learning rate

GradScaler multiplies the loss to keep gradients in a good numeric range during backpropagation.

Which PyTorch context manager is used to enable mixed precision operations?

A torch.cuda.amp.autocast()

B torch.cuda.amp.scale()

C torch.enable_grad()

D torch.no_grad()

torch.cuda.amp.autocast() enables automatic mixed precision for operations inside its block.

What data types are mainly used in mixed precision training?

A float16 and float32

B float64 and float128

C int32 and int64

D int8 and float8

Mixed precision training uses float16 (half precision) and float32 (single precision) floats.

What happens if you do not use loss scaling in mixed precision training?

A Training will be faster

B Memory usage will increase

C Model accuracy will always improve

D Gradients might become zero due to underflow

Without loss scaling, small gradient values can underflow to zero in float16, harming training.

Explain how mixed precision training works and why it is useful.

Think about how using smaller numbers can help speed and memory, but also what problem it causes.

Describe the steps to enable AMP in a PyTorch training loop.

Focus on the order of operations inside the training loop.

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