Accelerating Transformers with NVIDIA cuDNN 9
Accelerating Transformers with NVIDIA cuDNN 9
According to NVIDIA's technical blog, the NVIDIA CUDA Deep Neural Network library (cuDNN) is a GPU-accelerated library designed to enhance the performance of deep learning primitives. cuDNN is integrated with widely-used deep learning frameworks like PyTorch, TensorFlow, and XLA (Accelerated Linear Algebra), abstracting the complexities of direct GPU programming and allowing developers to focus on model design and training.
Performance Improvements and New Features
Recent updates in cuDNN have focused on improving the performance of scaled dot product attention (SDPA), a crucial primitive in large language models (LLMs). cuDNN has optimized SDPA through flash attention and other techniques, achieving up to 1.2 PFLOPS in FP8 on the NVIDIA H200 Tensor Core GPU. For instance, enabling cuDNN FP8 SDPA for Llama2 70B LoRA fine-tuning resulted in a 1.15x speedup, as measured by NVIDIA's team using the NVIDIA NeMo and Transformer Engine on an 8-GPU H200 node.
Scaled Dot Product Attention Performance
NVIDIA has significantly advanced the fused attention kernel, transitioning from the open-sourced Multihead Attention (fMHA) in the APEX library to a more flexible and faster implementation. This new implementation, which is the default in NVIDIA's Transformer Engine for Hopper architecture GPUs, uses flash attention techniques to optimize performance.
The XLA compiler provides a path to cuDNN SDPA, accessible through the JAX SDPA API or by lowering customized implementations in JAX/PyTorch to cuDNN SDPA. Efforts are underway to integrate a cuDNN-based SDPA implementation in PyTorch eager mode.
SDPA as cuDNN Graphs
SDPA in cuDNN can be defined as a cuDNN graph of tensor operations. These graphs, which are executed by engines within the cuDNN library, encapsulate a series of operations that can be run efficiently on NVIDIA GPUs. This approach allows for flexibility and high performance, utilizing the most efficient algorithms available, such as flash attention.
SDPA Usage Walkthrough
The cuDNN library provides several API entry points for creating and running graphs, including both frontend and backend APIs. The frontend API, available in Python and C++, simplifies the process of constructing and executing cuDNN graphs. An example of using the frontend API to create and run an SDPA graph in Python is provided in the NVIDIA documentation.
Other Notable cuDNN 9 Features
In addition to SDPA improvements, cuDNN 9 introduces several other enhancements:
- Mixed input precision support for matrix multiplications (matmuls) and convolutions
- Improved error reporting with more specific error codes and nested logging levels
- Hardware forward compatibility for future GPU architectures
- Streamlined installation processes
Mixed Input Precision Support
cuDNN 9 supports mixed input precision matmuls and convolutions, allowing operands of different data types (e.g., FP16 and INT8) with online fused type conversion, optimizing both performance and memory usage.
Improved Error Reporting
Enhanced error reporting in cuDNN 9 includes more specific error codes, categorization, and a new function cuDNNGetLastErrorString for retrieving the last error message programmatically.
Hardware Forward Compatibility
cuDNN 9 offers hardware forward compatibility for a subset of the API, enabling the use of future GPU architectures without needing to upgrade the cuDNN library.
Streamlined Installation
Installation of cuDNN 9 has been simplified, with pip support for the Python frontend and streamlined RPM and Debian meta-package installations.
Next Steps
NVIDIA encourages feedback and discussions on the cuDNN forum and GitHub. Developers can download cuDNN to start leveraging these new features and optimizations.
As AI continues to evolve, NVIDIA remains committed to optimizing performance and enhancing user experience, ensuring cuDNN remains a vital tool for deep learning frameworks and graph compilers.
Acknowledgments
The NVIDIA cuDNN team, in collaboration with various teams across the company, contributed the technical content for this post.