Profiling & Optimization: PyTorch Profiler, FLOPs, Memory

🎯 Learning Objectives

• Profile PyTorch models to identify bottlenecks
• Measure FLOPs and parameter counts
• Optimize memory usage and training throughput
• Use profiling results to guide optimization

📋 Prerequisites

• PyTorch training loop
• GPU architecture basics

---

1. 📖 Core Content

1.1 Why Profile?

1.2 PyTorch Profiler

python
Run Inline
import torch.profiler as profiler
with profiler.profile(
    activities=[
        profiler.ProfilerActivity.CPU,
        profiler.ProfilerActivity.CUDA,
    ],
    schedule=profiler.schedule(wait=1, warmup=1, active=3),
    on_trace_ready=profiler.tensorboard_trace_handler('./logs'),
    record_shapes=True,
    profile_memory=True,
) as prof:
    for step, data in enumerate(dataloader):
        train_step(data)
        prof.step()
        if step > 10:
            break
# Print top 10 CUDA operations by time
print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10))

1.3 Key Profiling Metrics

1.4 Memory Optimization

python
Run Inline
# Gradient checkpointing (trade compute for memory)
model = torch.utils.checkpoint.checkpoint_sequential(model, segments=4)
# In-place operations (careful with autograd)
torch.relu(x, inplace=True)  # Avoids allocating new tensor
# Batch normalization tracking
model.eval()  # Uses running stats, not batch stats
# Delayed parameter initialization
with torch.no_grad():
    for param in model.parameters():
        param.requires_grad = False  # Freeze layers
# Empty cache when needed
torch.cuda.empty_cache()

1.5 Common Bottlenecks

---

📝 Practice Questions

Q1

<strong>Q1

<strong>Q1

<strong>Q1

<strong>Q1

<strong>Q1</strong>: Profiling shows GPU utilization at 45% during training. What could be wrong?

Low GPU utilization (45%) means the GPU is idle >50% of the time. Common causes:

1. Data loading bottleneck: CPU can't prepare batches fast enough
   • Fix: Increase num_workers, enable pin_memory, use SSD
2. CPU preprocessing: Transforms/augmentations are CPU-heavy
   • Fix: Move some preprocessing offline, use GPU transforms
3. Small batch size: GPU finishes computation quickly, waits for next batch
   • Fix: Increase batch size, use gradient accumulation
4. Model too small: Not enough computation to saturate GPU
   • Fix: Increase model size, add more layers
5. Frequent synchronizations: Operations that sync CPU/GPU (.item(), .numpy(), print loss every step)
   • Fix: Accumulate metrics, sync less frequently

Target: GPU utilization > 90% for efficient training. Q2

<strong>Q2

<strong>Q2

<strong>Q2

<strong>Q2

<strong>Q2

<strong>Q2

<strong>Q2

<strong>Q2

<strong>Q2

<strong>Q2</strong>: A model with 100M parameters trains at batch_size=32 but OOMs at batch_size=64. Estimate the memory per sample.

Memory for batch_size=32: > GPU memory (e.g., 24GB) Memory per sample ≈ 24GB / 32 = 750MB Memory for batch_size=64: 64 × 750MB = 48GB

But this is approximate — activations scale linearly with batch size. Weights (100M × 4 bytes = 400MB) are constant regardless of batch size.

To fit batch_size=64 on the same GPU:

• Use gradient accumulation: batch=32, accumulation=2 (effective 64)
• Use mixed precision (reduces activation memory by ~40%)
• Use gradient checkpointing (reduces activation memory but increases compute by ~20%)

With gradient accumulation, effective batch size = 64 but per-step batch = 32. Q3

<strong>Q3

<strong>Q3

<strong>Q3

<strong>Q3</strong>: A model's forward pass takes 50ms, backward pass takes 100ms, and data loading takes 30ms per batch. What's the optimization priority?

Total time per step without overlap: 50 + 100 + 30 = 180ms (assuming sequential) With DataLoader prefetch (overlapping data loading with computation): 50 + 100 = 150ms (data loading overlaps)

Optimization priority:

1. Backward pass (100ms, 67% of time): This is dominated by weight gradient computation. Options: smaller model, mixed precision (2-3× faster backward)
2. Forward pass (50ms, 33% of time): Smaller model, efficient architecture
3. Data loading (30ms, already overlapped): Only if GPU utilization is still low

Always optimize the biggest bottleneck first. Improving data loading when it's already overlapped gives zero speedup. Q4: Compare torch.compile, torch.jit.script, and eager mode for inference performance.

torch.compile (PyTorch 2.0+) is the modern recommendation:

• Speed: 2-4× inference speedup, 1.5× training speedup
• Ease of use: Single line change: model = torch.compile(model)
• Debugging: Better error messages than torch.jit
• Dynamic shapes: Handles variable-length inputs well

For maximum performance: torch.compile + AMP + optimized batch size.

</details> * * * ## 🔗 Cross-References - **Next**: [Model Deployment](/notes/04-degree-electives-bsda5013-dl-practice-week08-08-model-deployment) - **Previous**: [Mixed Precision](/notes/04-degree-electives-bsda5013-dl-practice-week06-06-mixed-precision) - **Video**: BSDA5013 Week 7 transcripts [Join Discord](https://discord.gg/gE2m4Qrdqv) [Previous**Mixed Precision Training**](/notes/04-degree-electives-bsda5013-dl-practice-week06-06-mixed-precision)[Next**Model Deployment**](/notes/04-degree-electives-bsda5013-dl-practice-week08-08-model-deployment)