NVIDIA Improves Llama 3.1 405B Performance with TensorRT Model Optimizations

According to the NVIDIA Tech Blog, Meta’s Llama 3.1 405B large-scale language model (LLM) is reaching new levels of performance thanks to NVIDIA’s TensorRT Model Optimizer. These improvements result in up to a 1.44x increase in throughput when running on NVIDIA H200 GPUs.

Outstanding Llama 3.1 405B inference throughput using TensorRT-LLM

TensorRT-LLM has already delivered impressive inference throughput on Llama 3.1 405B since the model was released. This was achieved through a variety of optimizations, including in-flight batching, KV caching, and optimized attention kernels. These techniques accelerate inference performance while maintaining low-precision computing.

TensorRT-LLM adds support for the official Llama FP8 quantization recipe, which computes static and dynamic scaling factors to maintain maximum accuracy. Additionally, custom kernels, such as matrix multiplication in FBGEMM, are optimized via plugins inserted into the network graph at compile time.

Up to 1.44x performance improvement with TensorRT Model Optimizer

NVIDIA’s custom FP8 post-training quantization (PTQ) recipe, available via the TensorRT Model Optimizer library, improves Llama 3.1 405B throughput and reduces latency without sacrificing accuracy. The recipe reduces inference compute overhead by integrating FP8 KV cache quantization with self-attention static quantization.

Table 1 shows the peak throughput performance, showing significant improvements across a range of input and output sequence lengths on an 8-GPU HGX H200 system. The system features eight NVIDIA H200 Tensor Core GPUs, each with 141 GB of HBM3e memory, and four NVLink switches, providing 900 GB/s of GPU-to-GPU bandwidth.

Similarly, Table 2 shows the minimum delay performance using the same input and output sequence lengths.

These results demonstrate that H200 GPUs with TensorRT-LLM and TensorRT Model Optimizer deliver superior performance in both latency-optimized and throughput-optimized scenarios. The TensorRT Model Optimizer FP8 recipe also achieves similar accuracy to the official Llama 3.1 FP8 recipe on the Massively Multitask Language Understanding (MMLU) and MT-Bench benchmarks.

Equipping Llama 3.1 405B with two H200 GPUs using INT4 AWQ

For developers with hardware resource constraints, the INT4 AWQ technique in the TensorRT Model Optimizer compresses models so that the Llama 3.1 405B fits on just two H200 GPUs. This method compresses weights into 4-bit integers while encoding activations using FP16, significantly reducing the memory footprint required.

Tables 4 and 5 show the maximum throughput and minimum latency performance measurements, showing that the INT4 AWQ scheme provides similar accuracy scores as Meta’s Llama 3.1 official FP8 recipe.

Advances in NVIDIA’s TensorRT Model Optimizer and TensorRT-LLM pave the way for improved performance and efficiency in running large language models such as Llama 3.1 405B. These improvements provide developers with more flexibility and cost-effectiveness, whether they have extensive hardware resources or more limited environments.

Image source: Shutterstock