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

Profiling is essential for optimizing AI models, but it can be intimidating for beginners. A new series from Hugging Face aims to lower that barrier, starting with the fundamentals of using torch.profiler in PyTorch.

The series begins with a simple operation: matrix multiplication followed by bias addition. This foundational step serves as a gateway to understanding more complex profiling scenarios.

Setting Up torch.profiler

To profile PyTorch code, you'll use the torch.profiler module. The process involves:

1. Preparing the code to profile (in this case, a function that performs matrix multiplication and addition)
2. Annotating the algorithm with record_function (optional but recommended for easier trace navigation)
3. Wrapping the code with the torch.profiler.profile context manager
4. Exporting the profile data

The profiler generates two key artifacts:

1. A profiler table: Provides a statistical summary, highlighting hotspots and time-consuming operations.
2. A profiler trace: Offers a temporal view of execution, showing when and why operations occur on CPU and GPU.

Interpreting Profiler Output

The profiler table reveals crucial timing information. For a small 64x64 matrix operation, the results were surprising:

Self CPU time total: 2.314ms
Self CUDA time total: 23.104us

The GPU time (23.104 microseconds) is less than 1% of the CPU time. This indicates an overhead-bound algorithm, where most time is spent on setup rather than computation.

Scaling Up for Better GPU Utilization

To address this inefficiency, the authors increased the matrix size to 4096x4096. The results changed dramatically:

Self CPU time total: 4.908ms
Self CUDA time total: 4.495ms

Now both times are in milliseconds, showing much better GPU utilization. The CUDA kernel (ampere_bf16_s16816gemm_...) now dominates the execution time, indicating a shift from overhead-bound to compute-bound operation.

Visualizing the Execution Chain

The profiler trace, viewable in tools like Perfetto UI, provides a visual representation of the execution timeline. It shows the sequence of events from Python function calls down to CUDA kernel launches, helping identify bottlenecks and optimization opportunities.

Key Takeaways

This introductory exploration of PyTorch profiling reveals:

1. How to set up and use torch.profiler
2. The importance of scaling computations to fully utilize GPU capabilities
3. The value of both statistical summaries and temporal traces in understanding performance

Future installments in the series promise to delve into more complex scenarios, including nn.Linear layers, small MLPs, and eventually, large language models. By building profiling skills incrementally, the series aims to make performance optimization accessible to a wider range of PyTorch users.