Speed Python Pick Two How Cuda Graphs Enable Fast Python Code For Deep Learning

Speed, Python: Pick Two. How CUDA Graphs Enable Fast Python Code for Deep Learning

GLM 5.2 is live! Opus-level intelligence at open-source rates. Pay per token on serverless. Try it today.

Blog

Speed Python Pick Two How Cuda Graphs Enable Fast Python Code For Deep Learning Speed, Python: Pick Two. How CUDA Graphs Enable Fast Python Code for Deep Learning

PUBLISHED 8/29/2023

Table of Contents

Keeping Modern GPUs Busy: CPU/GPU Overlap

CPU Overheads — Where Does All The Time Go?

Feed the (GPU) Beast: Optimizing for CPU Overheads

CUDA Graphs

LLM Inference + CUDA graphs

Real-World CUDA Graphs: LLaMA v2 Inference

LLaMA2–7B + CUDA Graph Inference Performance Results

CUDA Graphs in the Fireworks Inference Platform

Conclusion

Appendix

Table of Contents

This is the second in a series of technical blog posts about the techniques we use for optimization of the high-performance Fireworks Gen AI Platform . See also the previous post about Multi-Query Attention . This post explores how the explosion in GPU speed over the past several years has changed the landscape of performance optimization for deep learning workloads. We examine how the host CPU has become a bottleneck due to this trend and review several techniques for mitigating this. We highlight one technique — CUDA graphs — that balances performance with usability. We examine the effect of CUDA graphs on Large Language Model (LLM) inference workloads and show a 2.3x speedup on a LLaMAv2–7B inference workload . We show that the Fireworks Inference Platform uses CUDA graphs and other aggressive machine and service optimizations to provide best-in-class speed and efficiency for LLM serving. Keeping Modern GPUs Busy: CPU/GPU Overlap

Contemporary deep learning programs are most often written in the Python language using PyTorch as a framework. PyTorch is fundamentally simple: a collection of pre-optimized Tensor operations that the user calls from a Python program. On the other hand, PyTorch has historically provided high performance for deep learning workloads, as described in the PyTorch paper . A key idea introduced in the PyTorch paper is the idea of CPU/GPU overlap : the CPU program dispatches work (kernels) for the GPU to execute, and so long as the CPU program runs faster than the GPU work, high performance is achieved.

A large batch-size training run with high CPU/GPU overlap However, the claims made in the paper about CPU/GPU overlap have become less true over time as GPU speeds have increased at a breakneck pace. Let's examine floating point performance and memory bandwidth for GPU architectures over time: • GP100 (as evaluated in the PyTorch paper) (2016) has 21 TFLOPS half-precision and 730 GB/s memory bandwidth • V100 (2017) has 112 TFLOPS half-precision (TensorCores) and 900 GB/s memory bandwidth • A100 (2020) has 312 TFLOPS half-precision (TensorCores) and 1600–2000 GB/s memory bandwidth • H100 (2022) 989 TFLOPs half-precision (TensorCores) and 3350 GB/s memory bandwidth

NVIDIA GPU Performance Over Time GPU half-precision floating point performance has increased 47x and memory bandwidth has increased 4.6x since the PyTorch paper was written. This speedup has profound implications on the ability of the CPU to stay ahead of the GPU and deliver maximum performance for deep learning workloads. CPU Overheads — Where Does All The Time Go?

What exactly is the CPU doing when you run a PyTorch program? Several layers of overhead exist. First, a PyTorch program has user-written logic that must be executed on the CPU. This logic includes metaprogramming, i.e. defining the structure of the network in Python code based on the hyperparameters. The simplest example of this is a loop over network layers. The best practice when writing PyTorch programs is to push this logic into Module construction so that this overhead is not incurred during model execution. Nonetheless, metaprogramming overheads still exist at runtime in most PyTorch programs. Second, when calling PyTorch operations, several decisions must be made about which compute kernel to call on the device (GPU) based on properties like the device/dtype of the operands or whether autograd recording is enabled. These decisions are made by a component called the dispatcher running on the CPU. Although the dispatcher is written as highly-optimized C++, its execution still contributes to runtime overhead while executing a PyTorch program. Third, GPU memory allocation contributes to runtime overhead. PyTorch uses a sophisticated caching memory allocator to alleviate much of this overhead, but runtime performance may still be negatively affected by allocator activities on the CPU when executing programs with small operations. Finally, CUDA itself has overheads in the driver and kernel launch paths, which can slow down the CPU execution of the program. Feed the (GPU) Beast: Optimizing for CPU Overheads

As GPUs get faster, CPU overheads become more of a problem when executing deep learning programs. Several techniques and tools have emerged to solve this problem.

A batch-size 1 inference run on a HuggingFace Transformers model with very poor CPU/GPU overlap HuggingFace Transformers is a ubiquitous codebase for language models based on the Transformer architecture. However, in practice, Transformers is not highly optimized for inference, as the Python code is written in a way that maximizes flexibility but incurs significant CPU overhead. On the other end of the spectrum is FasterTransformer (FT), written in highly optimized C++ to maximize performance. In practice, FT's optimized C++ code is quite hard to work with, slowing down the development of new features or adding new model architectures (e.g. LLaMA models are not officially supported ). Another approach to performance optimization is automatic code transformation via compilation. Apache TVM is an early example of this approach, which provided an end-to-end stack that compiled a high-level representation of a neural network down to native machine code (e.g. on GPU). However, converting from a flexible PyTorch program to TVM's high-level representation is non-trivial, especially when a program contains control flow (loops, branches) implemented in Python or when a program involves distributed operations. Dynamic shapes (Tensor shapes that change at runtime) introduce additional complications. A newer approach is torch.compile...