Last Updated on 2024-07-31 by Clay
介紹
HuggingFace 所提供的 TGI (Text Generation Inference) 是一個專門用來佈署、加速 LLM 推理服務的框架，以下是它的架構圖：
在最外層、最左側，我們可以看到在 API 服務層級，路徑是開到 <url>/generate 這裡，如果資料從 API 中輸入，會先進入一個緩衝區（Buffer），再由 Batcher 去將其傳入後台模型服務；根據官方文件的說明，他們有使用 vLLM 所實作的 PagedAttention，其內核是直接引用 vLLM 的。
關於 vLLM 的加速推理機制，可以參考：
使用方法
最推薦的使用方法自然是使用官方包好的 Docker Image 來啟動，這樣也能保證環境的配置是正確的。
model=/data/teknium--OpenHermes-2.5-Mistral-7B/
volume=$PWD/data/
docker run \
--gpus device=1 \
--shm-size 10g \
-p 8080:80 \
-v $volume:/data/ \
ghcr.io/huggingface/text-generation-inference:2.0.4 \
--model-id $model
有個看官方範例比較不明白的是，$model 需要填寫的是放到容器內的位置，仔細想想是理所當然的，但卻很容易混淆。
總之我們的 TGI 服務開啟之後，我們就可以使用 curl 指令來直接發送請求。
CURL 發送請求
curl 127.0.0.1:8080/generate \
-X POST \
-d '{"inputs":"What is Deep Learning?","parameters":{"max_new_tokens":20}}' \
-H 'Content-Type: application/json'
Output:
{"generated_text":"\n\nDeep learning is a subset of machine learning in artificial intelligence (AI) that has networks called"}
使用 Python 發送請求
當然，我們也可以使用任何我們熟悉的語言發送請求，比如說 Python：
import httpx
url = "http://127.0.0.1:8080/generate"
headers = {
"Content-Type": "application/json"
}
data = {
"inputs": "What is Deep Learning?",
"parameters": {
"max_new_tokens": 20
}
}
response = httpx.post(url, json=data, headers=headers)
# Check response
if response.status_code == 200:
print("Response received successfully:")
print(response.json())
else:
print(f"Failed to get a response. Status code: {response.status_code}")
print(response.text)
Output:
Response received successfully:
{'generated_text': '\n\nDeep learning is a subset of machine learning in artificial intelligence (AI) that has networks called'}
使用 huggingface_hub 的 InferenceClient 發送請求
from huggingface_hub import InferenceClient
client = InferenceClient(model="http://127.0.0.1:8080")
client.text_generation(prompt="How are you today?")
Output:
'\n\nI’m doing well, thank you.\n\nWhat’s your name?\n'
如果想要串流式生成的話...
from huggingface_hub import InferenceClient
client = InferenceClient(model="http://127.0.0.1:8080")
for token in client.text_generation("How are you today?", max_new_tokens=50, stream=True):
print(token, end="")
Output:
I’m doing well, thank you.
What’s your name?
My name is John.
Where are you from?
I’m from the United States.
What do you do?
結尾
在使用 TGI 之前我一直是 vLLM 派的，不過可能是採樣方式不太一樣的緣故，經常在輸出格式上 TGI 更符合我微調後的需求。
總之 TGI 會是一個我用來作加速推理框架的候補，因為其啟用方式跟發送請求都非常直覺且容易，應該會被我拿來測試新模型（前提是新模型若是新架構的話 TGI image 的更新速度需要夠快才型呢）。