The Complete Guide to GPU Optimization for Deep Learning

When training deep neural networks, the GPU is often the most expensive and underutilized resource. You might think you're maxing out your A100s, but profiling often reveals significant idle time due to memory bottlenecks.

1. Mixed Precision Training

The easiest win for GPU optimization is Automatic Mixed Precision (AMP). By utilizing float16 or bfloat16 for most operations while keeping float32 for accumulations and gradient updates, you can double your throughput and halve your memory usage.

PyTorch AMP Example

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

model = MyModel().cuda()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
scaler = GradScaler()

for data, target in dataloader:
    data, target = data.cuda(), target.cuda()
    
    optimizer.zero_grad()
    
    # Forward pass with autocast
    with autocast():
        output = model(data)
        loss = loss_fn(output, target)
        
    # Backward pass with scaler
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

2. Dataloader Bottlenecks

If your nvidia-smi shows GPU utilization constantly bouncing between 0% and 100%, you are likely CPU-bound in your dataloader. The GPU is waiting for the CPU to load and augment images.

Quick Fixes

• Increase num_workers: Set num_workers to the number of available CPU cores (but test to find the sweet spot, often 4-8).
• pin_memory=True: This allows PyTorch to allocate page-locked memory, making memory transfers to the GPU much faster.
• Avoid heavy transformations on CPU: If possible, move augmentations to the GPU. libraries like NVIDIA DALI or Kornia are excellent for this.

3. Fused Optimizers and Kernels

When you run optimizer.step(), PyTorch typically iterates over every parameter tensor and performs the update element-wise. This creates significant CUDA kernel launch overhead.

Using fused optimizers (like torch.optim.Adam(..., fused=True)) or xFormers can combine these small operations into a single kernel, significantly speeding up training.

4. Activation Checkpointing (Gradient Checkpointing)

If you are running out of memory (OOM), activation checkpointing trades compute for memory. Instead of storing all activations for the backward pass, it discards some of them and recomputes them on the fly during the backward pass.

python
from torch.utils.checkpoint import checkpoint

class LargeModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.layer1 = HeavyLayer()
        self.layer2 = HeavyLayer()
        
    def forward(self, x):
        # Normal execution
        # x = self.layer1(x)
        
        # Checkpointed execution
        x = checkpoint(self.layer1, x)
        x = checkpoint(self.layer2, x)
        return x

Mastering these techniques will save your organization thousands of dollars in compute costs and drastically reduce your iteration time.