Why Mixed precision training (AMP) in PyTorch? - Purpose & Use Cases

The Big Idea

What if you could train your AI models twice as fast without buying new hardware?

The Scenario

Imagine training a deep learning model on a large dataset using only full 32-bit precision numbers. The training takes hours or even days, and your computer's memory fills up quickly, slowing everything down.

The Problem

Using only full precision means your computer uses more memory and runs slower. It can't handle bigger models or larger batches efficiently. Also, training can be unstable or crash if the numbers get too big or too small.

The Solution

Mixed precision training uses both 16-bit and 32-bit numbers smartly. It speeds up training and saves memory without losing accuracy. This lets you train bigger models faster and more reliably.

Before vs After

✗ Before

for data, target in loader:
    optimizer.zero_grad()
    output = model(data)
    loss = loss_fn(output, target)
    loss.backward()
    optimizer.step()

✓ After

scaler = torch.cuda.amp.GradScaler()
for data, target in loader:
    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 It Enables

Mixed precision training unlocks faster, more memory-efficient model training, making complex AI projects possible on common hardware.

Real Life Example

A researcher trains a large image recognition model on a single GPU. Using mixed precision, training time is cut in half, allowing quicker experiments and better results.

Key Takeaways

Training with only full precision is slow and memory-heavy.

Mixed precision smartly combines 16-bit and 32-bit numbers for speed and stability.

This approach enables faster, larger, and more efficient AI model training.

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