Profiling in PyTorch: A Beginner's Guide to torch.profiler

Profiling in PyTorch: A Beginner's Guide to torch.profiler

Understanding torch.profiler for Performance Optimization

Profiling is the essential first step in optimizing Large Language Models (LLMs) and deep learning pipelines, as it allows developers to identify whether a model is overhead-bound or compute-bound. The torch.profiler module provides two primary artifacts to diagnose performance: the profiler table, which offers a statistical summary of "what" is taking the most time, and the profiler trace, which provides a temporal view of "when" and "why" operations occur across the CPU and GPU lanes.

Identifying Bottlenecks: Overhead-Bound vs. Compute-Bound

Performance bottlenecks are often revealed by comparing the time spent on the CPU versus the GPU in the profiler table.

  • Overhead-Bound: When Self CPU time total is significantly higher than Self CUDA time total (e.g., CPU time in milliseconds while GPU time is in microseconds), the algorithm is overhead-bound. This indicates the CPU spends more time preparing and launching kernels than the GPU spends executing them. This is common with small matrix multiplications.
  • Compute-Bound: When Self CPU time total and Self CUDA time total are both in the millisecond range and comparable, the algorithm is compute-bound. This is the ideal state for high-performance computing, where the GPU is the primary bottleneck.

Increasing the workload size (e.g., moving from 64x64 to 4096x4096 matrices) typically shifts a model from an overhead-bound regime to a compute-bound one.

Analyzing the CPU and GPU Dispatch Chain

PyTorch operations follow a specific chain of dispatch from Python calls down to CUDA kernels. A typical matrix multiplication and addition sequence looks like this:

ProfileStep $\rightarrow$ record_function (user annotation) $\rightarrow$ aten::matmul (ATen-level dispatch) $\rightarrow$ aten::mm (2D matrix-matrix multiply backend) $\rightarrow$ cudaLaunchKernel.

Key CPU Lane Indicators

  • Cold-Start Overhead: The first ProfileStep is often wider than subsequent steps due to workspace allocations, cuBLAS heuristics, and lazy module loading. This can be mitigated using warmup iterations.
  • CUDA Occupancy Queries: The presence of cudaOccupancyMaxActiveBlocksPerMultiprocessor before a cudaLaunchKernel indicates a "heavyweight" kernel (like GEMM or convolution). The CPU is querying the driver to determine the optimal block size based on hardware capacity.
  • Resource-Light Kernels: Elementwise or reduction kernels typically lack an occupancy query because their resource footprint is fixed and small.
  • Synchronization: A long cudaDeviceSynchronize at the end of a trace is often the profiler flushing events; its duration reflects the time the CPU waited for the GPU to finish pending work.

Key GPU Lane Indicators

  • Activity Buffer Requests: Gaps between kernels or an initial offset between CPU and GPU lanes are often caused by the profiler allocating or refilling its own event buffers.
  • Runtime Variance: Identical kernels can exhibit different runtimes across steps due to GPU clock fluctuations, thermals, power management, or driver housekeeping.

The Impact of torch.compile

Using torch.compile transforms eager PyTorch code into optimized graphs via TorchDynamo, AOTAutograd, and TorchInductor. Profiling a compiled function reveals several architectural changes:

Dispatcher-Level Fusion

For an operation like torch.add(torch.matmul(x, w), b), torch.compile performs operator fusion at the graph level, replacing the separate aten::add and aten::mm calls with a single aten::addmm call. However, this is fusion at the dispatcher level, not the kernel level. The GPU still executes a Memcpy DtoD (to seed the destination buffer with the bias) followed by a GEMM kernel with a bias-add epilogue.

Runtime Architecture and Overhead

Compiled regions introduce a specific CPU-side hierarchy:

  1. TorchDynamo Cache Lookup: Verifies that input shapes, dtypes, and devices match the cached compilation. This occurs on every call.
  2. Torch-Compiled Region: The wrapper entering the compiled version.
  3. AOTDispatcher Runtime Wrapper Prologue: Handles tensor metadata and view tracking.
  4. Call CompiledFxGraph: Executes the generated code (identified by a content hash).

For very small operations, torch.compile can actually increase CPU overhead because the cost of walking the Dynamo $\rightarrow$ AOTAutograd $\rightarrow$ Inductor stack exceeds the savings from fusion. This overhead amortizes as the number of operations in the model increases.

Trace Reading Summary Table

Observation Likely Meaning
Self CPU $\gg$ Self CUDA Overhead-bound; increase batch size or fuse ops
cudaOccupancy... before launch Heavyweight, adaptively-launched kernel (GEMM/Conv)
aten::matmul $\rightarrow$ aten::bmm Batched matrix multiplication on 3D+ tensors
Torch-Compiled Region Execution is within a torch.compile optimized block
Memcpy DtoD before GEMM Bias copy for an addmm epilogue in compiled mode

Sources