Mixed precision training (AMP) in PyTorch - Model Pipeline Trace

Model Pipeline - Mixed precision training (AMP)

This pipeline shows how mixed precision training uses both 16-bit and 32-bit numbers to speed up training while keeping accuracy. It helps the model learn faster and use less memory.

Data Flow - 5 Stages

1Data Loading

1000 rows x 28 x 28 pixels→Load images and labels from dataset→1000 rows x 28 x 28 pixels

Image of handwritten digit '5' with label 5

↓

2Preprocessing

1000 rows x 28 x 28 pixels→Normalize pixel values to 0-1 range→1000 rows x 28 x 28 pixels

Pixel value 150 normalized to 0.59

↓

3Feature Engineering

1000 rows x 28 x 28 pixels→Convert images to tensors for model input→1000 rows x 1 x 28 x 28 tensor

Tensor shape for one image: [1, 1, 28, 28]

↓

4Model Training with AMP

Batch of 32 tensors [32, 1, 28, 28]→Train model using mixed precision (float16 and float32) with automatic scaling→Model weights updated, loss scalar

Loss value 0.45 after first batch

↓

5Metrics Calculation

Model predictions and true labels→Calculate accuracy and loss→Accuracy scalar, loss scalar

Accuracy 0.82, Loss 0.45

Training Trace - Epoch by Epoch

Loss
0.7 |*       
0.6 | *      
0.5 |  *     
0.4 |   *    
0.3 |    *   
0.2 |     *  
0.1 |      * 
     --------
     1 2 3 4 5 Epochs

Epoch	Loss ↓	Accuracy ↑	Observation
1	0.65	0.75	Loss starts high, accuracy moderate as model begins learning
2	0.48	0.83	Loss decreases, accuracy improves with mixed precision speeding training
3	0.35	0.89	Model converges faster due to AMP, loss lowers steadily
4	0.28	0.92	Stable training, accuracy nearing high performance
5	0.22	0.94	Final epoch shows good convergence with low loss and high accuracy

Prediction Trace - 6 Layers

Layer 1: Input Layer

Layer 2: Convolutional Layer (float16)

Layer 3: Activation (ReLU)

Layer 4: Fully Connected Layer (float32)

Layer 5: Softmax

Layer 6: Prediction

Model Quiz - 3 Questions

Why does mixed precision training use both float16 and float32 numbers?

A To speed up training and save memory while keeping accuracy

B To make the model smaller but slower

C To increase the model size for better learning

D To avoid using GPUs during training

Key Insight

Mixed precision training balances speed and accuracy by using faster float16 computations where possible and precise float32 where needed. This helps models train faster and use less memory without losing performance.

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