高速なLLM用推論ライブラリ「vLLM」を試す

Huggingface Transformersに代わる高速ライブラリとして、vLLMというのが公表されているとのGigazineの記事がありました。とても分かりやすく動作原理やその効果を説明してくれていて、興味深く読ませてもらいました。

大規模言語モデルの出力スピードを最大24倍に高めるライブラリ「vLLM」が登場、メモリ効率を高める新たな仕組み「PagedAttention」とは？

vLLMは、HuggingFace Transformersでサポートされているアーキテクチャーのモデルをサポートしています。

というわけで、くわしい解説はGigazine記事にお任せして、実際に自分の環境で試食してみます。

Docker環境で試す

• vLLMのドキュメントに従って、Docker環境でAPIサーバーをたててvicuna-7b-v1.3を試してみます。

Quickstart — vLLM

$ docker run --gpus all -it --rm --shm-size=8g -p 8000:8000 \
  nvcr.io/nvidia/pytorch:22.12-py3
# pip uninstall torch
# pip install vllm

# python -m vllm.entrypoints.openai.api_server --host "0.0.0.0" \
  --model lmsys/vicuna-7b-v1.3

動作確認

import openai
import os

openai.api_key = "EMPTY"
openai.api_base = "http://localhost:8000/v1"

response= openai.Completion.create(
    model="lmsys/vicuna-7b-v1.3",
    prompt="Please suggest 10 names for a new Japanese airline.",
    max_tokens=500)

print(response['choices'][0]['text'])

Name must be unique, easy to pronounce and remember in English and Japanese. One additional requirement is that the name must not be already taken by an existing Japanese airline.

1. Flying Sakura
2. Sky Wind
3. Sea Breeze
4. Sunrise One
5. Celestial Nights
6. Sky Horizon
7. Garden Air
8. Serene Skies
9. Ocean Wings
10. Cherry Blossoms Soar</s>

動作は確かに早い印象もありますが、私の環境(RTX 3090)では、今のところ量子化オプションが見当たらないため7bモデルしか使えないこともあり、まだオモチャにして遊ぶには時期尚早な印象でした。

今後の開発に期待です！

現場からは以上です。