Posted by Manish Gupta, Staff Software Engineer, Google Research
AI-driven technologies have become an integral part of our daily lives, offering immense potential to enhance our access to information and improve productivity. These applications heavily rely on large language models (LLMs), which are memory-intensive and require specialized hardware accelerators to deliver high computing power efficiently. In this blog post, we will explore how we can address the computational challenges of LLMs by optimizing memory usage.
The majority of an LLM’s memory and compute resources are consumed by weights in matrix multiplication operations. One effective way to reduce memory consumption is by using narrower data types. For example, storing weights in 8-bit integer (U8 or S8) data type can reduce the memory footprint by 4x compared to single-precision (F32) and 2x compared to half-precision (F16) or bfloat16 (BF16). Previous research has shown that utilizing S8 weights and F16 input (preserving higher precision) in matrix multiplications can significantly improve efficiency without sacrificing accuracy. This technique, known as weight-only quantization, requires efficient implementation of matrix multiplication with mixed inputs, such as half-precision input multiplied with 8-bit integer.
Hardware accelerators, including GPUs, support a fixed set of data types, making it necessary to transform mixed-input matrix multiplication operations to match the hardware operations. In this blog, we focus on mapping mixed-input matrix multiplication onto the NVIDIA Ampere architecture. We present software techniques that address data type conversion and layout conformance to efficiently map mixed-input matrix multiplication onto hardware-supported data types and layouts. Our results demonstrate that the additional software work has minimal overhead and allows us to achieve performance close to the peak hardware capabilities. The software techniques we describe are available in the open-source NVIDIA/CUTLASS repository.
One key aspect of modern AI hardware accelerators, like Google’s TPU and NVIDIA’s GPU, is their ability to multiply matrices natively in the hardware using specialized processing elements called Tensor Cores. In this blog, we specifically focus on NVIDIA Ampere Tensor Cores, which perform the matrix-multiply-accumulate (mma) operation. The mma operation in Tensor Cores supports mixed-precision, allowing for the combination of different input data types. However, mixed-input matrix multiplication, which involves mixing input data types, is not directly supported by the hardware and needs to be implemented in software.
To address the challenges of mixed-input matrix multiplication, we examine a specific example of F16 user input and U8 model weights (F16 * U8). We discuss two key challenges: data type conversion and layout conformance. The mma operation requires two input matrices with the same data type, so a data type conversion from U8 to F16 is necessary to utilize the hardware-supported mixed-precision Tensor Cores. Additionally, the layout of the input matrices needs to conform to the hardware specifications within a warp, which is a group of 32 threads.
To tackle the data type conversion challenge, we implement FastNumericArrayConvertor, which operates on 4xU8 in 32-bit registers without unpacking individual 1xU8 values. This optimized implementation reduces the number of instructions and improves conversion speed. Moreover, it utilizes permute bytes and packed arithmetic operations to achieve the conversion, resulting in a 1.6x reduction in the number of operations compared to conventional approaches.
For layout conformance, we introduce FragmentShuffler, which rearranges the data within registers using shuffle operations. This enables the use of wider bitwidth load operations, increases shared memory bandwidth utilization, and reduces the overall number of operations. We leverage the ldmatrix instruction provided by the NVIDIA Ampere architecture, which allows for efficient movement of data from shared memory to registers, aligning with the mma matrix A and B requirements. FragmentShuffler ensures that the layout after the load conforms to the U8*U8 mma operation.
By implementing these optimized software strategies, we can further reduce the overhead of data type conversion and layout conformance. FastNumericArrayConvertor and FragmentShuffler significantly improve the efficiency of these operations, resulting in better performance. These strategies have been incorporated into the NVIDIA/FasterTransformer library and are available for use.
In conclusion, we have explored the challenges of mixed-input matrix multiplication in LLMs and presented software techniques to address these challenges on the NVIDIA Ampere architecture. By optimizing data type conversion and layout conformance, we can effectively utilize hardware-supported mixed-precision Tensor Cores and achieve efficient matrix multiplication with mixed inputs. These software strategies enable us to maximize performance while minimizing overhead.
Source link
