Profile PyTorch Training with the Built-in Profiler

Nsight Systems (Lab I-3) gives you the system view — CUDA API calls, CPU threads, NVTX ranges, kernel launches. It's thorough but heavy, and you need to drop to shell + parse CSVs.

PyTorch Profiler is the lighter cousin. Five lines of Python give you per-operator timings, memory allocations, and Chrome trace exports — all inside your Python script. It's the go-to tool for model-level optimization: "which torch op is eating my time?"

Both matter. They're complementary:

References

• PyTorch Profiler recipe — the official quickstart.
• PyTorch Profiler API docs — torch.profiler.profile, activities, schedules.
• Chrome Trace Event Format — the JSON format PyTorch Profiler exports for chrome://tracing / Perfetto.
• Horace He — Making Deep Learning Go Brrrr — mental model for when to trust op-level times (answer: always bracket with torch.cuda.synchronize()).
• PyTorch TensorBoard Profiler plugin — rich GUI if you install torch_tb_profiler.

What you'll do

1. Wrap a training step with torch.profiler.profile(...); extract top ops
2. Pretty-print the key_averages() table; read the columns correctly
3. Add memory profiling; find the biggest allocators
4. Export a Chrome trace; inspect it to see the scheduler gaps

---

The minimum useful profile:

with torch.profiler.profile(activities=[CPU, CUDA]) as prof:
    for _ in range(5):
        training_step()
prof.key_averages().table(sort_by='cuda_time_total', row_limit=20)

Two gotchas:

• Run a few warmup steps outside the profiler — the first CUDA launch compiles kernels, skewing the numbers.
• torch.cuda.synchronize() at the end — without it, some async kernels may still be in flight and not counted.

The profiler's three timing columns often confuse first-time users:

Optimize by self CUDA time for op-level work, CUDA time for whole-module work.

import torch
import torch.nn as nn
from torch.profiler import profile, ProfilerActivity, record_function

device = torch.device('cuda')
torch.manual_seed(0)

# A small but realistic training setup
model = nn.Sequential(
    nn.Conv2d(3, 64, 7, stride=2, padding=3), nn.BatchNorm2d(64), nn.ReLU(),
    nn.Conv2d(64, 128, 3, stride=2, padding=1), nn.BatchNorm2d(128), nn.ReLU(),
    nn.Conv2d(128, 256, 3, stride=2, padding=1), nn.BatchNorm2d(256), nn.ReLU(),
    nn.AdaptiveAvgPool2d(1), nn.Flatten(),
    nn.Linear(256, 1000),
).to(device)
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
loss_fn = nn.CrossEntropyLoss()

def training_step():
    x = torch.randn(64, 3, 224, 224, device=device)
    y = torch.randint(0, 1000, (64,), device=device)
    out = model(x)
    loss = loss_fn(out, y)
    optimizer.zero_grad(set_to_none=True)
    loss.backward()
    optimizer.step()

# Warmup OUTSIDE the profiler (first CUDA launches compile kernels)
for _ in range(3):
    training_step()
torch.cuda.synchronize()
print('warmed up')

# Profile 5 steps
with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof:
    for _ in range(5):
        with record_function('training_step'):
            training_step()
    torch.cuda.synchronize()

# Extract top ops
top_ops = []
for evt in prof.key_averages():
    top_ops.append({
        'name': evt.key,
        'cpu_time_us': evt.cpu_time_total,
        'cuda_time_us': evt.device_time_total if hasattr(evt, 'device_time_total') else evt.cuda_time_total,
        'self_cuda_time_us': evt.self_device_time_total if hasattr(evt, 'self_device_time_total') else evt.self_cuda_time_total,
        'count': evt.count,
    })
top_ops.sort(key=lambda op: -op['cuda_time_us'])
top_ops = top_ops[:15]

print(f'Captured {len(top_ops)} top ops (by CUDA time):\n')
print(f'{"op":<40s} {"CPU (ms)":>10s} {"CUDA (ms)":>10s} {"count":>6s}')
for op in top_ops[:12]:
    print(f'{op["name"][:40]:<40s} {op["cpu_time_us"]/1000:>10.2f} {op["cuda_time_us"]/1000:>10.2f} {op["count"]:>6d}')