Guide to vLLM Using Gemma-7b-it

Introduction

Everyone needs to have faster and reliable inference from the Large Language models. vLLM, a cutting-edge open source framework designed to simplify the deployment and management of large language models with very less throughput. vLLM makes your job easier by offering efficient and scalable tools for working with LLMs. With vLLM, you can manage everything from model loading and inference to fine-tuning and serving, all with a focus on performance and simplicity. In this article we will implement vLLM using Gemma-7b-it model from HuggingFace. Lets dive in.

Learning Objectives

• Learn what vLLM is all about, including an overview of its architecture and why it’s generating significant buzz in the AI community. 
• Understand the importance of KV Cache and PagedAttention, which form the core architecture that enables efficient memory management and fast LLM inference and serving.
• Learn and explore in detail guide to vLLM using Gemma-7b-it
• Additionally, explore how to implement HuggingFace models, such as Gemini, using vLLM.
• Understand the importance of using Sampling Params in vLLM, which helps to tweak the model’s performance.

This article was published as a part of the Data Science Blogathon.

vLLM Architecture Overview

vLLM, short for “Virtual large language model,” is an open-source framework designed to streamline and optimize the use of large language models (LLMs) in various applications. vLLM is a game-changer in the AI space, offering a streamlined approach to handling large language models. Its unique focus on performance and scalability makes it an essential tool for developers looking to deploy and manage language models effectively.

The buzz around vLLM is due to its ability to handle the complexities associated with large-scale language models, such as efficient memory management, fast inference, and seamless integration with existing AI workflows. Traditional methods often struggle with efficient memory management and fast inference, two critical challenges when working with massive datasets and complex models. vLLM addresses these issues head-on, offering a seamless integration with existing AI workflows and significantly reducing the technical burden on developers.

In order to understand how, let’s understand the concept of KV Cache and PagedAttention. 

Understanding KV Cache

KV Cache (Key-Value Cache) is a technique used in transformer models, specifically in the context of Attention mechanisms, to store and reuse the intermediate results of key and value computations during the inference phase. This caching significantly reduces the computational overhead by avoiding the need to recompute these values for each new token in a sequence, thus speeding up the processing time.

How KV Cache Works?

• In transformer models, the Attention mechanism relies on keys (K) and values (V) derived from the input data. Each token in the input sequence generates a key and a value.
• During inference, once the keys and values for the initial tokens are computed, they are stored in a cache.
• For subsequent tokens, the model retrieves the cached keys and values instead of recomputing them. This allows the model to efficiently process long sequences by leveraging the previously computed information.

Math Representation

• Let K_i ​ and V_i ​ be the key and value vectors for token i.
• The cache stores these as K_cache = {K_1 , K_2 ,…, K_n } and V_cache = { V_1 , V_2 ,… , V_n }.
• For a new token t, the attention mechanism computes the attention scores using the query Q_t​ with all cached keys K_cache.

Despite being so efficient, in most of the cases, KV cache is large. For instance, in the LLaMA-13B model, a single sequence can take up to 1.7GB. The size of the KV cache depends on sequence length, which is variable and unpredictable, leading to inefficient memory usage.

Traditional methods often waste 60%–80% of memory due to fragmentation and over-reservation. To mitigate this, vLLM introduces PagedAttention.

What is PagedAttention?

PagedAttention addresses the challenge of efficiently managing memory consumption when handling very large input sequences, which can be a significant issue in transformer models. Unlike the KV Cache, which optimizes the computation by reusing previously computed key-value pairs, PagedAttention further enhances efficiency by breaking down the input sequence into smaller, manageable pages. The concept operates uses these manageable pages and performs attention calculations within these pages.

How it Works?

Unlike traditional attention algorithms, PagedAttention allows for the storage of continuous keys and values in a fragmented memory space. Specifically, PagedAttention divides the KV cache of each sequence into distinct KV blocks.

• In transformer models, the attention mechanism relies on keys (K) and values (V) derived from the input data. Each token in the input sequence generates a key and a value.
• During inference, once the keys and values for the initial tokens are computed, they are stored in a cache.
• For subsequent tokens, the model retrieves the cached keys and values instead of recomputing them. This allows the model to efficiently process long sequences by leveraging the previously computed information.

Math Representation

• B be the KV block size (number of tokens per block)
• K_j be the key block containing tokens from position (j-1)B + 1 to j_B
• V_j be the value block containing tokens from position (j-1)B + 1 to j_B
• q_i be the query vector for token i
• A_ij be the attention score matrix between q_i and K_j
• o_i be the output vector for token i
• The query vector `q_i` is multiplied with each KV block (`K_j`) to calculate the attention scores for all tokens within that block (`A_ij`).
• The attention scores are then used to compute the weighted average of the corresponding value vectors (`V_j`) within each block, contributing to the final output `o_i`.

This in return allows the flexible memory management: 

• Removing the need for contiguous memory allocation by eliminating internal and external fragmentation.
• KV blocks can be allocated on demand as the KV cache expands.
• Physical blocks can be shared across multiple requests and sequences, reducing memory overhead.

Gemma Model Inference Using vLLM

Lets implement the vLLM framework using the Gemma-7b-it model from HuggingFace Hub.

Step1: Installation of the Module

In order to get started, let’s begin by installing the module. 

!pip install vllm

Step2: Define LLM

First, we import the necessary libraries and set up our Hugging Face API token. We only need to set HuggingFace API token only for few models that requires permission. Then, we initialize the google/gemma-7b-it model with a maximum length of 2048 tokens, ensuring efficient memory usage with torch.cuda.empty_cache() for optimal performance.

import torch,os
from vllm import LLM

os.environ['HF_TOKEN'] = "<replace-with-your-hf-token>"

model_name = "google/gemma-7b-it"
llm = LLM(model=model_name,max_model_len=2048)

torch.cuda.empty_cache()

Step3: Sampling Parameters Guide in vLLM

SamplingParams is similar to the model keyword arguments in the Transformers pipeline. This sampling parameters is essential to achieve the desired output quality and behavior.

• temperature: This parameter controls the randomness of the model’s predictions. Lower values make the model output more deterministic, while higher values increase randomness.
• top_p: This parameter limits the selection of tokens to a subset whose cumulative probability is above a certain threshold (p). To simply gets consider top_p to be 0.95, then the model considers only the top 95% probable next words, which helps in maintaining a balance between creativity and coherence, preventing the model from generating low-probability, and often irrelevant, tokens.
• repetition_penalty: This parameter penalizes repeated tokens, encouraging the model to generate more varied and less repetitive outputs.
• max_tokens: Max tokens determine the maximum number of tokens in the generated output.

from vllm import SamplingParams

sampling_params = SamplingParams(temperature=0.1,
                      top_p=0.95,
                      repetition_penalty = 1.2,
                      max_tokens=1000
                  )

Step4: Prompt Template for Gemma Model

Each open-source model has its own unique prompt template with specific special tokens. For instance, Gemma utilizes <start_of_turn> and <end_of_turn> as special token markers. These tokens indicate the beginning and end of a chat template, respectively, for both user and model roles.

def get_prompt(user_question):
    template = f"""
    <start_of_turn>user
    {user_question}
    <end_of_turn>
    <start_of_turn>model
    """
    return template

prompt1 = get_prompt("best time to eat your 3 meals")
prompt2 = get_prompt("generate a python list with 5 football players")

prompts = [prompt1,prompt2]

Step5: vLLM inference

Now that everything is set, let the LLM generate the response to the user prompt. 

from IPython.display import Markdown

outputs = llm.generate(prompts, sampling_params)

display(Markdown(outputs[0].outputs[0].text))
display(Markdown(outputs[1].outputs[0].text))

Once the outputs is executed, it returns the processed prompts result that contains speed and output i.e., token per second. This speed benchmarking is beneficial to show the difference between vllm inference and other. As you can observe in just 6.69 seconds, we generated two user prompts. 

Step6: Speed benchmarking

Processed prompts: 100%|██████████| 2/2 [00:13<00:00, 6.69s/it, est. speed input: 3.66 toks/s, output: 20.70 toks/s]

Output: Prompt-1

Output: Prompt-2

Conclusion

We successfully executed the LLM with reduced latency and efficient memory utilization. vLLM is a game-changing open-source framework in AI, providing not only fast and cost-effective LLM serving but also facilitating the seamless deployment of LLMs on various endpoints. In this article we explored guide to vLLM using Gemma-7b-it.

Click here to access the documentation.

Key Takeaways

• Optimisation of the LLM in memory is very critical and with vLLM one can one easily achieve faster inference and serving. 
• Understanding the basics of Attention mechanism in depth, can lead to understand how beneficial PagedAttention mechanisms and KV cache is about. 
• Implementation of vLLM inference on any HuggingFace model is pretty straight forward and requires very less lines of code to achieve it. 
• Sampling Params in vLLM is very important to be defined, if one needs the right response back from vLLM.

Frequently Asked Questions

Reference

• Paper: https://arxiv.org/abs/2309.06180
• vLLM blog: https://blog.vllm.ai/2023/06/20/vllm.html

The media shown in this article is not owned by Analytics Vidhya and is used at the Author’s discretion.

Tarun R Jain

Data Scientist at AI Planet || YouTube- AIWithTarun || Google Developer Expert in ML || Won 5 AI hackathons || Co-organizer of TensorFlow User Group Bangalore || Pie & AI Ambassador at DeepLearningAI