Hex-LLM: A New LLM Serving Framework Designed for Efficiently Serving Open LLMs on Google Cloud TPUs

In the rapidly evolving world of artificial intelligence, large language models (LLMs) have become essential tools for a variety of applications, ranging from natural language understanding to content generation. While the capabilities of these models continue to expand, efficiently serving and deploying them remains a challenge, particularly when it comes to balancing cost, throughput, and latency. Recent advancements by Google and the introduction of Hex-LLM, a specialized serving framework, offer promising solutions for efficiently deploying open LLMs from Hugging Face on Google TPUs.

Hex-LLM: A Game-Changer for Serving Open LLMs on TPUs

Hex-LLM is Vertex AI’s in-house LLM serving framework that is designed and optimized for Google’s Cloud TPU hardware, which is available as part of AI Hypercomputer. It provides a high-performance, low-cost solution for deploying open-source models from Hugging Face. Developed to address the challenges of serving large models at scale, Hex-LLM stands out due to its advanced optimization techniques, which allow it to handle significant workloads with impressive efficiency.

Key Features and Innovations of Hex-LLM

To efficiently serve LLMs on TPUs, Hex-LLM integrates a variety of key features and optimization techniques, which significantly enhance performance:

1. Token-Based Continuous Batching: One of the standout features of Hex-LLM is token-based continuous batching. This method allows for efficient utilization of TPU resources by processing incoming tokens in a continuous stream. By handling requests in this manner, Hex-LLM maximizes throughput, significantly reducing the cost per token served. This approach ensures that no TPU cycles are wasted, resulting in an overall boost in efficiency.
2. XLA-Optimized PagedAttention Kernels: Hex-LLM employs XLA (Accelerated Linear Algebra) optimized PagedAttention kernels, which are crucial for managing the attention mechanism of transformer models. These kernels are tailored to exploit the full potential of TPU hardware, minimizing the latency and computational load associated with the attention calculations. By leveraging XLA-optimized kernels, Hex-LLM achieves low-latency inference, which is essential for applications requiring real-time or near-real-time responses.
3. Tensor Parallelism: Another critical feature of Hex-LLM is tensor parallelism, which enables the distribution of model computations across multiple TPU cores. This parallelism is particularly beneficial for serving large models like Llama 2 70B, as it allows for the workload to be split effectively, ensuring that the TPUs operate at peak efficiency without being bottlenecked by single-threaded tasks.
4. Dynamic LoRA Adapters and Quantization: Hex-LLM supports the use of Dynamic Low-Rank Adaptation (LoRA) adapters, which offer a flexible way to fine-tune models for specific tasks without retraining the entire model. Additionally, Hex-LLM supports quantization techniques, including BNB (Billion-scale Neural Basis) and AWQ (Adaptive Weight Quantization), allowing models to run with lower precision, thereby reducing memory usage and increasing inference speed without compromising performance.

Integration with Hugging Face Hub

Hex-LLM integrates directly with the Hugging Face Hub, allowing developers to easily load and serve models from the extensive library of open LLMs available. This seamless integration simplifies the process of deploying models on Google TPUs, making it more accessible for those who may not have extensive experience with TPU infrastructure. By directly pulling models from Hugging Face, users can quickly experiment with different LLMs and deploy them in production environments without the need for extensive manual configuration.

Performance Metrics: Speed and Cost

The performance of Hex-LLM is impressive, particularly when serving large models. For instance, Hex-LLM achieves a throughput of 1510 output tokens per second for Llama 2 70B in int8 precision on a single TPU v5e-8, with an approximate cost of $9.60 per hour. This translates to a latency of 26 milliseconds per token, which is remarkable for a model of this size. These metrics demonstrate that Hex-LLM is not only capable of serving large models with high efficiency but also does so at a cost that is feasible for many applications.

Availability in Vertex AI Model Garden

Hex-LLM is available as part of the Vertex AI Model Garden, a platform that offers a wide variety of pre-trained models and tools for machine learning. By including Hex-LLM in the Model Garden, Google provides users with a straightforward way to access and deploy open LLMs on TPUs, complete with the optimizations offered by the Hex-LLM framework. This availability ensures that users can leverage the power of TPUs for LLM deployment without needing to set up the infrastructure from scratch.

Conclusion

Hex-LLM represents a significant step forward in the efficient serving of open LLMs, particularly for users looking to deploy large models on Google TPUs. With features like token-based continuous batching, XLA-optimized PagedAttention kernels, tensor parallelism, and direct integration with Hugging Face, Hex-LLM offers a powerful and cost-effective solution for LLM deployment. While its current status as a closed-source framework may limit its accessibility, the performance gains and cost reductions it provides make it an attractive option for organizations seeking to leverage the power of large language models in their applications.

Check out the Details here and LInkedIn Post. All credit for this research goes to the researchers of this project. Also, don’t forget to follow us on Twitter and join our Telegram Channel and LinkedIn Group. If you like our work, you will love our newsletter.. Don’t Forget to join our 50k+ ML SubReddit

[Upcoming Event- Oct 17 202] RetrieveX – The GenAI Data Retrieval Conference (Promoted)