Last Updated on 2023-12-14 by Clay
介紹
vLLM 是加州柏克萊分校所開發的一種大型語言模型（Large Language Model, LLM）加速推理框架。它主要是利用 PagedAttention 機制提高了 GPU VRAM 的使用率，並且這一方法無須更改模型的架構。
PagedAttention 是啟發於作業系統中經典的虛擬記憶體和分頁技術，故命名為 Paged（分頁）。它的架構中最重要的實現便是能夠在『非連續的記憶體區塊』上儲存 transformer 模型注意力機制中『連續的 key 和 value』。
PagedAttention 將每個序列的 KV Cache 分割成一個個的 Block（每個 Block 中 token 數量是固定的），接著再把 Block 中的值與 Q 批次做運算。
並且這一作法取得了良好的效果。
我們可以看到在不同的模型和 GPU 條件下，vLLM 框架展現了壓倒性的推理速度。
使用方式
首先，自然是安裝 vllm 套件，如果需要使用 AWQ 量化，則需要額外安裝 autoawq。
pip3 install vllm autoawq話說我本來是想要測試 Mistral 系列微調後的 OpenHermes，並且由於我本地端的 VRAM 不夠，所以想要使用量化模型，結果沒想到怎麼樣載入都會報錯。一查之下，才發現 vllm 的 AWQ 量化模型尚未支持 Mistral 系列。
所以我只好先測試 Llama 了；Mistral 系列的測試暫時還是用 CPU offloading 來稍微撐一下。
from vllm import LLM, SamplingParams
# Sample prompts.
prompts = [
"Hello, my name is",
"The president of the United States is",
"The capital of France is",
"The future of AI is",
]
# Create a sampling params object.
sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
# Create an LLM.
llm = LLM(model="TheBloke/Llama-2-7b-Chat-AWQ", quantization="AWQ")
# Generate texts from the prompts. The output is a list of RequestOutput objects
# that contain the prompt, generated text, and other information.
outputs = llm.generate(prompts, sampling_params)
# Print the outputs.
for output in outputs:
prompt = output.prompt
generated_text = output.outputs[0].text
print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
Output:
Prompt: 'Hello, my name is', Generated text: " Sherry and I'm a 35-year-old woman from" Prompt: 'The president of the United States is', Generated text: ' a member of the executive branch of the federal government. The president serves as the' Prompt: 'The capital of France is', Generated text: ' Paris, which is known for its stunning architecture, art museums, historical' Prompt: 'The future of AI is', Generated text: ' exciting and uncertain. Here are some potential developments that could shape the field'
這樣的推理，包含模型載入時間一共才 10 秒！如果把模型去掉，只生成這短短幾個 tokens 的時間甚至只需要花 0.3 秒！
當然，只有生這一點點 tokens 的時間很難評估是否真的推理速度極快、所以我又測試了生成 100 tokens 的速度 —— 發現也只需要 2.3 秒。
這真的很快啊！希望能找到更多的時間來好好研究下 vLLM 的實作細節，包含目前 Mistral 系列模型的技術難點。
與此同時，官方也好心地提供了量化的簡易腳本（需要安裝 autoawq）：
from awq import AutoAWQForCausalLM
from transformers import AutoTokenizer
model_path = 'lmsys/vicuna-7b-v1.5'
quant_path = 'vicuna-7b-v1.5-awq'
quant_config = { "zero_point": True, "q_group_size": 128, "w_bit": 4, "version": "GEMM" }
# Load model
model = AutoAWQForCausalLM.from_pretrained(model_path, **{"low_cpu_mem_usage": True})
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
# Quantize
model.quantize(tokenizer, quant_config=quant_config)
# Save quantized model
model.save_quantized(quant_path)
tokenizer.save_pretrained(quant_path)
不過，官方的文件裡面也說明了，目前的 vLLM 雖然支援 AWQ，但是尚未進行優化，所以推理速度會比未量化的更慢一些。也就是說，我上面測試的那個 AWQ 模型，實際上未量化版本只會更快。