PyTorch Performance Profiling & Optimization: From Profiler to Kernel Fusion

Note: This is a review article (no first-hand experimental data), based on public practice and performance engineering.

Experiment Objective

"Training is slow" is the most common PyTorch complaint, but where it is slow (GPU compute? communication? CPU starving data?) cannot be answered without profiling. Goal: a standard PyTorch performance optimization flow — use torch.profiler to locate the bottleneck (GPU time vs communication vs data loading) → govern by bottleneck class (kernel optimization/fusion/DataLoader/memory) → finish with compilation (torch.compile). After reading you can answer: what is slow here, what to optimize first, and how to verify the improvement.

Code Walkthrough

Part 1: Locating bottlenecks with the profiler

from torch.profiler import profile, ProfilerActivity

with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
             schedule=torch.profiler.schedule(wait=1, warmup=2, active=5),
             record_shapes=True) as prof:
    for step in range(8):
        train_one_step()
        prof.step()

# Summary: GPU total time / communication time / idle time in one view
print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=15))

Key points:

1. warmup+active is the standard posture: skip compile/warmup steps, only measure steady state — wait=1, warmup=2, active=5 is a common template.
2. Watch three numbers: cuda_time_total (is GPU busy), cpu_time_total (is CPU the bottleneck), and the communication-op share (NCCL time) — classify before optimizing; 90% of "slowness" is data loading or communication, not kernels.
3. record_shapes for shape-related analysis: shape-dependent recompilation/fragmentation is visible from here.

Part 2: Governing by bottleneck — DataLoader and communication

# Common fixes for data-loading bottlenecks (DataLoader config)
DataLoader(dataset, batch_size=64,
           num_workers=16,          # multi-process prefetch
           prefetch_factor=4,       # prefetched batches per worker
           pin_memory=True,         # pinned memory, faster H2D copy
           persistent_workers=True) # avoid worker restarts per epoch

Key points:

• Data loading is the #1 hidden bottleneck: low GPU utilization + high CPU utilization + a "DataLoader" share → that's it. The four-piece set (multi-process/prefetch/pin_memory/persistent workers) works immediately in most cases.
• Communication bottleneck signals: normal kernel share but slow end-to-end, high NCCL time → check ch04 storage (data remote) and ch03 network (IB/ROCE), or tune ch11's communication overlap (gradient buckets).
• Kernel bottleneck: a single op dominating cuda_time → only then do op-level work (fusion, Triton custom kernels — linking ch09's Triton article).

Principles

Four-layer bottleneck classification

Layer	Signal	Governance
data loading	CPU high, GPU low, DataLoader share	multiprocess/prefetch/pin_memory
communication	high NCCL time share	gradient buckets/topology/network tuning
kernel	one op dominates cuda_time	fusion/quantization/custom kernels
memory	OOM or forced small batch	gradient accumulation/activation recompute/ZeRO

Core method: the profiler first attributes "slow" to one of the four layers, then pick the matching governance — do not start optimizing kernels; most cases are in the first three layers. This is "observability first" (ch02) extended into the single training process.

torch.compile and operator fusion

model = torch.compile(model, mode="reduce-overhead")  # or default/max-autotune

• torch.compile is "free fusion optimization": compiles and fuses the op graph (same roots as ch08's XLA/torch compilers); small models/dynamic shapes gain little or negative, large-model training gains are clear (5%-30%).
• Mode selection: reduce-overhead cuts Python dispatch; max-autotune searches per-kernel configs (longer compile, higher gains); production training commonly uses default.
• Verification: compare with the profiler — same load, how much cuda_time or end-to-end throughput improved; validate with data, not "feels faster."

Memory governance order

On OOM, try in increasing cost: smaller batch + gradient accumulation → torch.utils.checkpoint (activation recompute) → mixed precision (bf16/FP8) → ZeRO/FSDP sharding (ch09) — tune hyperparameters first, then framework-level means; never jump to the most expensive option.

Results

(Estimated, not measured; typical large-model training)

Item	Typical gain	Difficulty
DataLoader four-piece	+10%-40% (if data is the bottleneck)	low
torch.compile	+5%-30%	low-mid
kernel fusion/custom	+10%-50% (kernel-bound)	high
gradient accumulation/recompute	effective batch enlargement	low

Engineering Takeaways

1. Profile before optimizing: the profiler is step one; the bottleneck class decides the direction — "classify before governing" saves 90% of wasted work.
2. DataLoader four-piece is the best-ROI first move: most "GPU not saturated" is data starvation; check this before kernels.
3. torch.compile is worth trying by default: large-model training usually gains clearly; verify with the profiler before keeping it.
4. Memory issues governed in increasing cost: hyperparameters → checkpoint recompute → precision → sharding; never jump straight to ZeRO.

Next article preview: single-process optimization done; back to the scale dimension — "Long-Context Pretraining: Position Encoding, Efficiency, and Stability."