Speeding up Your ML Workload: Pytorch Compile and... Aishwariya C, Priyanka N & Kavya G & G Chander

### Overview This presentation examines methods for accelerating machine learning workloads in PyTorch, transitioning from traditional eager execution to graph-based execution via `torch.compile` introduced in PyTorch 2.0. The speaker highlights that while eager mode offers dynamic execution, flexibility, and straightforward debugging, it introduces substantial overhead in production. Under eager execution, computations are dispatched dynamically line-by-line without a global picture of the workload, frequently leaving the GPU underutilized while waiting on Python frame evaluation and small kernel dispatches. `torch.compile` solves this by capturing whole computation graphs and compiling optimized, backend-specific machine code without requiring framework rewrites.

### Key Topics - **Eager Mode Bottlenecks**: Analysis using PyTorch Profiler on a Granite model demonstrates that GPU idle gaps occur due to CPU-bound Python dispatch overhead, repetitive memory-bound kernel launches, and frequent data movement between memory and registers. - **Graph-Based Compilation**: How `torch.compile` delivers JIT compilation without framework redesigns via a decorator or wrapper (`torch.compile(model)`). - **Core Internal Stages**: Examination of the compilation pipeline: Graph Capture (TorchDynamo and AOTAutograd), Graph Lowering, and Graph Compilation (TorchInductor). - **Handling Graph Breaks and Guards**: Mechanisms for falling back to the CPython interpreter for unsupported operations and using guards to reuse compiled artifacts across runs. - **Optimizations and Backend Execution**: Techniques such as operator fusion (e.g., fusing matrix multiplication and addition into a single Triton kernel), constant folding, dead-code elimination, and buffer reuse. - **IBM Spyre Hardware and Torch-Spyre**: Compilation considerations for IBM's Spyre dataflow AI accelerator architecture.

### Technical Details - **PyTorch Profiler Observations**: Tracing showed instances where 131 ms elapsed between kernel calls solely due to CPU evaluation of Python frames, and scenarios where seven tiny kernels (2–7 µs each) ran across 41 ms with significant memory transfer overhead. A 73 µs GEMM kernel only represented 2.2% of compute time. - **Internal Components**: - **TorchDynamo**: Inspects Python bytecode via the Frame Evaluation API, extracts PyTorch operations into an FX graph, and rewrites bytecode. If unsupported calls (e.g., Python built-in `print`) occur, it issues a graph break, delegating execution back to the CPython interpreter without crashing. - **Guards**: Tracks invariants (shapes, input datatypes, target hardware/devices) to reuse compiled artifacts; changing conditions triggers recompilation. - **AOTAutograd**: Uses fake tensor inputs based on initial execution shapes to generate joint forward and backward graphs, then applies a min-cut algorithm to minimize the intermediate tensor states stored between forward and backward passes. - **TorchInductor**: The default compiler backend, generating Triton code for GPUs and OpenMP C++ code for CPUs. Other backends mentioned include TVM, TensorRT, and XLA. - **Benchmarks**: An experiment on MobileNetV2 with ImageNet showed a 32% reduction in latency and a 48% increase in throughput. CPU benchmarks demonstrated up to 1.5x speedups, and official PyTorch benchmarks on NVIDIA A100 GPUs showed up to 76% faster execution on TorchBench. - **IBM Spyre Accelerator**: An AI accelerator featuring 32 ring-connected cores with mixed-precision SIMD systolic arrays, no hardware cache, and scratchpad memory. Because data movement across scratchpads must be explicitly planned prior to execution, graph compilation via Torch-Spyre is essential.

### Takeaways - `torch.compile` bridges the gap between eager prototyping ease and graph execution performance with a single-line integration. - Eliminating graph breaks and enabling aggressive operator fusion are critical for minimizing memory traffic and maximizing GPU compute saturation. - Getting started with contributions can begin by profiling standard models to report graph breaks, improving TorchInductor passes with Triton, or implementing custom backends for novel hardware accelerators.