PyTorch Profiler

PyTorch Profiler bridges the gap between your Python code and GPU execution — showing which operators take the most time, where memory peaks, and whether your GPUs are being utilized. It’s the first tool to reach for before diving into Nsight.

---

1. When to Use PyTorch Profiler

Situation	Tool
“Which layer is slowest?”	PyTorch Profiler
“Why is there a gap between GPU kernels?”	Nsight Systems
“Why is this specific GEMM only 50% efficient?”	Nsight Compute
“Is data loading the bottleneck?”	PyTorch Profiler (CPU ops visible)

PyTorch Profiler is the starting point — it gives operator-level insight without requiring NVIDIA’s specialized tools.

---

2. Basic Usage

Simple Profile

from torch.profiler import profile, record_function, ProfilerActivity

with profile(
    activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
    record_shapes=True,
    profile_memory=True,
    with_stack=True,
) as prof:
    for step in range(5):
        with record_function("forward"):
            output = model(input)
            loss = criterion(output, target)
        with record_function("backward"):
            loss.backward()
        with record_function("optimizer"):
            optimizer.step()
            optimizer.zero_grad()

# Print top operators by GPU time
print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=15))

Output Table

Name                    CPU total   CUDA total   # Calls
----------------------  ----------  -----------  --------
aten::mm                 12.3ms      45.2ms        96
aten::addmm               3.1ms      18.7ms        48
aten::_softmax            1.2ms       8.3ms        32
aten::layer_norm          0.8ms       6.1ms        32
aten::gelu                0.3ms       2.4ms        32
forward                  18.1ms      82.3ms         5
backward                 22.4ms      98.1ms         5

What to look for:

• aten::mm / aten::addmm = GEMM (linear layers) — should dominate
• If non-GEMM ops (softmax, layernorm) take a large fraction → fusion opportunity
• CPU total >> CUDA total for an op → CPU overhead (Python dispatch, data prep)

---

3. Chrome Trace Export

The most powerful visualization — shows a timeline in Chrome’s trace viewer:

# Export trace
prof.export_chrome_trace("trace.json")

# Open in browser: chrome://tracing → Load → trace.json

The trace shows:

• CPU operators (with Python call stacks if with_stack=True)
• CUDA kernels (with duration and correlation to CPU ops)
• Memory allocations and frees
• Gaps between operations (idle time)

What to Look For in the Trace

Pattern	Meaning	Fix
Long CPU bars before GPU kernels	Python/framework overhead	torch.compile, reduce Python logic
GPU gaps between kernels	Launch overhead or sync stalls	CUDA Graphs, remove .item() / .cpu()
aten::to taking significant time	Device transfers	Minimize CPU↔GPU transfers
aten::copy_ repeated	Unnecessary tensor copies	Use in-place operations where safe

---

4. Scheduled Profiling (Training Loop)

Profile specific iterations to capture steady-state behavior:

from torch.profiler import schedule

with profile(
    activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
    schedule=schedule(
        wait=2,      # skip 2 iterations (warmup)
        warmup=2,    # profile 2 iterations (JIT compilation)
        active=3,    # capture 3 iterations (steady state)
        repeat=1,    # one cycle
    ),
    on_trace_ready=torch.profiler.tensorboard_trace_handler("./tb_logs"),
    record_shapes=True,
    profile_memory=True,
) as prof:
    for step in range(10):
        # Training step...
        prof.step()  # signal end of iteration

View in TensorBoard:

tensorboard --logdir=./tb_logs
# Navigate to "PyTorch Profiler" tab

TensorBoard’s profiler plugin shows:

• Overview: GPU utilization, step time breakdown
• Operator view: sorted by GPU time, call count
• Trace view: interactive timeline (like Chrome trace)
• Memory view: allocation timeline, peak analysis
• GPU kernel view: individual CUDA kernel durations

---

5. Memory Profiling

Track GPU memory allocation to find peaks and leaks:

with profile(
    activities=[ProfilerActivity.CUDA],
    profile_memory=True,
    record_shapes=True,
) as prof:
    output = model(input)
    loss = criterion(output, target)
    loss.backward()

# Memory timeline
print(prof.key_averages().table(sort_by="self_cuda_memory_usage", row_limit=10))

Memory Snapshot (PyTorch 2.1+)

torch.cuda.memory._record_memory_history()

# Run your code
output = model(input)
loss = criterion(output, target)
loss.backward()

torch.cuda.memory._dump_snapshot("memory_snapshot.pickle")
# Visualize at: https://pytorch.org/memory_viz

The memory visualizer shows every allocation and free, with stack traces — invaluable for debugging OOM errors.

---

6. Common Bottleneck Diagnosis

“Training is slow but GPU utilization looks high”

# Check: is it data loading?
print(prof.key_averages().table(sort_by="cpu_time_total"))
# If DataLoader ops dominate CPU time → data loading bottleneck

“GPU utilization is low (< 50%)”

Possible causes:

1. Small batch size → GPU not fully utilized. Increase batch or use gradient accumulation
2. CPU bottleneck → data loading or Python overhead. Profile CPU ops
3. Sync stalls → check for .item(), .cpu(), or print(tensor) in the training loop
4. Unbalanced distributed → one GPU waiting for others

“OOM at a specific step”

# Use memory snapshot to find the allocation spike
torch.cuda.memory._record_memory_history(max_entries=100000)
# Run until OOM, then dump

---

7. Quick Profiling Checklist

# 1. Quick operator breakdown
python -c "
import torch
from torch.profiler import profile, ProfilerActivity
model = ...  # your model
with profile(activities=[ProfilerActivity.CUDA]) as p:
    for _ in range(3): model(input)
print(p.key_averages().table(sort_by='cuda_time_total', row_limit=10))
"

# 2. Chrome trace for visual inspection
# Add prof.export_chrome_trace('trace.json') and open in chrome://tracing

# 3. If specific kernel is slow → switch to Nsight Compute
# ncu --kernel-name "slow_kernel" python train.py

---

Self-Check

• Profile a model and find the top 5 operators by GPU time
• Export and read a Chrome trace — identify GPU idle gaps
• Set up scheduled profiling with warmup/active phases
• Diagnose a data loading bottleneck from profiler output
• Use memory snapshot to find peak allocation
• Explain the profiling workflow: PyTorch Profiler → Nsight Systems → Nsight Compute

---

References

• PyTorch Profiler Tutorial
• PyTorch TensorBoard Profiler
• PyTorch Memory Visualizer