newsfilter.io
Conference Presentation, Keynote

Making GPUs Actually Fast: A Deep Dive into Training Performance

Jane StreetCorwin, Savant Diaz, Sylvain De Wecker
Watch on YouTube →
  • CPUs are approximately 1,000 times faster than GPUs for single floating-point operations (1 nanosecond vs. 1 microsecond) but lack the massive parallelism required for matrix multiplication.
  • For 2048x2048 matrix multiplication, GPUs achieve a 100x speedup over CPUs (200 microseconds vs. 28 milliseconds) by parallelizing output elements across threads.
  • The standard PyTorch training loop involves a data loader, forward pass, loss calculation, gradient computation (loss.backward()), and optimizer step.
  • GPU utilization is inefficient when the "blue bar" (CUDA hardware line) has gaps, indicating idle time caused by CPU-GPU synchronization.
  • Synchronizations occur when copying data back to the CPU for metrics, checking torch.isnan (due to boolean coercion), or holding CPU memory stable during asynchronous CPU-to-GPU transfers.
  • Synchronizations stall the pipeline because the CPU cannot launch new kernels while waiting for the GPU to finish the transfer or for memory to be pinned.
  • Training loop performance can be improved by computing metrics entirely on the GPU via kernels and deferring data transfer until the end of the epoch.
  • Asynchronous isNan checks can be implemented using background threads or pipelines to prevent the main thread from blocking on data synchronization.
  • Memory pinning (pin_memory=True) combined with non_blocking=True allows Direct Memory Access (DMA) for asynchronous CPU-to-GPU transfers without CPU involvement.
  • GPU architecture consists of Streaming Multiprocessors (SMs); an H100 contains 132 SMs that cannot efficiently communicate with each other.
  • Memory hierarchy speeds on an H100 are: Shared Memory (26 TB/s) > L2 Cache (7 TB/s) > Global Memory (3 TB/s).
  • Efficient kernel design prioritizes data residency in shared memory to minimize traffic from global memory.
  • Three primary bottlenecks in GPU kernel performance are compute limits, memory bandwidth, and kernel launch overhead (5–10 microseconds per launch).
  • Kernel fusion combines multiple small operations (e.g., LayerNorm, ReLU, addition) into a single kernel to amortize memory loads and reduce launch overhead.
  • torch.compile automatically traces model operations to identify fusion opportunities, though it may create "graph breaks" around custom operators, dictionary comprehensions, or hidden synchronization points.
  • torch.compile relies on Triton templates which may achieve 80–90% of NVIDIA's hand-optimized kernel performance, but fails to parallelize Python for-loops effectively.
  • Custom CUDA kernels can resolve scenarios where automated fusion fails, such as summing thousands of independent tensors inside a Python loop.
  • Custom kernel strategy involves assigning one tensor to one SM, using thread indexing for partial sums, and performing reductions via shared memory and warp-level operations.
  • A custom kernel reduced the time for summing multiple tensors from 20 milliseconds to 30 microseconds, a nearly 1,000x performance improvement.
  • Development hierarchy at Jane Street: Pure PyTorch > Triton (Python DSL, 80-90% performance) > CUDA C++ (for absolute performance or hardware-specific features).
  • PTX (PTX Intermediate Representation) is used sparingly, typically only for injecting specific instructions into existing CUDA kernels.