2024/01/10 12:30 JST 追記。
GPU複数枚を使用して1つのモデルをロードさせる方法が分かったので、「6. GPUx2を試してみる」追記。

「LLM 推論と提供のための高速で使いやすいライブラリ」と言われているvLLMのコードを読みつつ、アレコレ試してみます。

使用するPCはドスパラさんの「GALLERIA UL9C-R49」。スペックは
・CPU: Intel® Core™ i9-13900HX Processor
・Mem: 64 GB
・GPU: NVIDIA® GeForce RTX™ 4090 Laptop GPU(16GB)
・GPU: NVIDIA® GeForce RTX™ 4090 (24GB)
・OS: Ubuntu22.04 on WSL2（Windows 11）
です。

1. 準備

venv環境を作り、

python3 -m venv vllm
cd $_
source bin/activate

pip installして、

pip install vllm

pip listで確認です。

Package                   Version
------------------------- ------------
aioprometheus             23.12.0
aiosignal                 1.3.1
anyio                     4.2.0
attrs                     23.2.0
certifi                   2023.11.17
charset-normalizer        3.3.2
click                     8.1.7
exceptiongroup            1.2.0
fastapi                   0.108.0
filelock                  3.13.1
frozenlist                1.4.1
fsspec                    2023.12.2
h11                       0.14.0
httptools                 0.6.1
huggingface-hub           0.20.2
idna                      3.6
Jinja2                    3.1.2
jsonschema                4.20.0
jsonschema-specifications 2023.12.1
MarkupSafe                2.1.3
mpmath                    1.3.0
msgpack                   1.0.7
networkx                  3.2.1
ninja                     1.11.1.1
numpy                     1.26.3
nvidia-cublas-cu12        12.1.3.1
nvidia-cuda-cupti-cu12    12.1.105
nvidia-cuda-nvrtc-cu12    12.1.105
nvidia-cuda-runtime-cu12  12.1.105
nvidia-cudnn-cu12         8.9.2.26
nvidia-cufft-cu12         11.0.2.54
nvidia-curand-cu12        10.3.2.106
nvidia-cusolver-cu12      11.4.5.107
nvidia-cusparse-cu12      12.1.0.106
nvidia-nccl-cu12          2.18.1
nvidia-nvjitlink-cu12     12.3.101
nvidia-nvtx-cu12          12.1.105
orjson                    3.9.10
packaging                 23.2
pip                       22.0.2
protobuf                  4.25.1
psutil                    5.9.7
pydantic                  1.10.13
python-dotenv             1.0.0
PyYAML                    6.0.1
quantile-python           1.1
ray                       2.9.0
referencing               0.32.1
regex                     2023.12.25
requests                  2.31.0
rpds-py                   0.16.2
safetensors               0.4.1
sentencepiece             0.1.99
setuptools                59.6.0
sniffio                   1.3.0
starlette                 0.32.0.post1
sympy                     1.12
tokenizers                0.15.0
torch                     2.1.2
tqdm                      4.66.1
transformers              4.36.2
triton                    2.1.0
typing_extensions         4.9.0
urllib3                   2.1.0
uvicorn                   0.25.0
uvloop                    0.19.0
vllm                      0.2.7
watchfiles                0.21.0
websockets                12.0
xformers                  0.0.23.post1

2. コードの修正

vLLMのコードを読むと、transformersを内部的に呼び出していました。なので、既存の問合せのコードからどこを修正するのがよいか、という観点から見ていきましょう。
※わたしが理解したかったからだけです、はい。

import

#import torch
#from transformers import AutoModelForCausalLM, AutoTokenizer, TextStreamer

これらはコメントアウトで良し。

from vllm import LLM, SamplingParams

vllmを追加。

from typing import List, Dict
import time

このまま。

modelとtokenizer

model_id = "elyza/ELYZA-japanese-Llama-2-7b-instruct"

# トークナイザーとモデルの準備
model = LLM(
    model=model_id,
    #device_map="auto",
    #torch_dtype="auto",
    dtype="auto",
    #low_cpu_mem_usage=True,
    trust_remote_code=True,
)

cuda前提ですし、accelerateパッケージをインストールしていないので、device_mapは指定できない。同じ理由でlow_cpu_mem_usageも。
あと、torch_dtypeはdtypeという変数名へ修正します。

tokenizer = model.get_tokenizer()

LLMクラスにget_tokenizerというメソッドがあり、 中身をみると（いろいろとしてますが）AutoTokenizerを呼び出しているだけでした。
チャットテンプレートを展開する際にAutoTokenizerが必要なので、参照しやすいようにtokenizer変数に代入しておきます。

生成のためのパラメータ

DEFAULT_SYSTEM_PROMPT = "あなたは誠実で優秀な日本人のアシスタントです。"

# generation params
max_new_tokens = 1024
generation_params = SamplingParams(
    #do_sample = True,
    temperature = 0.8,
    top_p = 0.95,
    top_k = 40,
    #max_new_tokens = max_new_tokens,
    max_tokens = max_new_tokens,
    repetition_penalty = 1.1
)

SmplingParamsクラスの各種パラメータを設定します。do_sampleはないのでコメントアウト、あとmax_new_tokens ではなくmax_tokensなので修正します。

推論の生成

def q(
    user_query: str,
    chat_history: List[Dict[str, str]]=None
):
    start = time.process_time()
    # messages
    messages = [
        {"role": "system", "content": DEFAULT_SYSTEM_PROMPT},
    ]
    user_messages = [
        {"role": "user", "content": user_query}
    ]
    if chat_history:
        user_messages = chat_history + user_messages
    messages += user_messages

チャット履歴を保持するために、配列を作っています。ここは変わらず。

    # generateion prompts
    prompt = tokenizer.apply_chat_template(
        conversation=messages,
        add_generation_prompt=True,
        tokenize=False
    )

先ほど代入したtokenizer変数はここで使用しています。これまでどおり、apply_chat_templateメソッドを使えます。ですので、ここも変わらずですね。

    # 推論
    outputs = model.generate(
        prompt,
        generation_params
    )

推論部分が変わっています。
ここでの引数は、apply_chat_templateメソッドの返却値であるpromptと、上記の「生成のためのパラメータ」で設定した変数samling_paramsの2つを指定しています。

vllm/vllm/entrypoints/llm.pyを確認するとgenerateメソッドには4つの引数が定義されています。全部、Noneと定義されていますが、promptsもしくはprompt_token_idsのいずれかは指定しなければなりません。

• prompts: prompt_token_idsを指定すればNoneと指定可能。

• sampling_params: リクエストのサンプリングパラメータ。指定しなければデフォルト値が設定される。

• request_id: リクエストのuid。

• prompt_token_ids: promptのtoken id配列。指定しなければprompt変数から勝手にコンバートする。

ですので、generateの呼び出し方としては、
(1) プロンプトのencodeは私にやらせろ。
(2) プロンプトのencodeは委せた。
のいずれかが可能です。

(2)は示したコードなので、ここでは(1)のケースのコードを紹介します。encodeメソッドの引数return_tensorsを指定すると、演算子エラー(tensor + listはできない）が発生しますので、削除が必要です。また、引数prompt_token_idsに渡すとき、変数を[…] で囲ってください。そうしないと、こちらもまた「intには len()は使えない」とエラーになります。

    input_ids = tokenizer.encode(
        prompt,
        add_special_tokens=False,
        #return_tensors="pt"
    )
    # 推論
    outputs = model.generate(
        sampling_params=generation_params,
        prompt_token_ids=[input_ids],
    )

横道にそれたので、戻ります。

なお、generateメソッドの返却型はRequestOutputクラスのリストです。

推論結果の出力

    output = outputs[0]
    print("--- prompt")
    print(output.prompt)
    print("--- output")
    print(output.outputs[0].text)

generateメソッドの返却値は List[RequestOutput] ですので、このコードですと、outputs[0]と書かなければ参照できません。毎回 outputs[0] と書くのが面倒なのと読みにくいので、outputに代入しています。

なお、output.promptとしてprintされるのはgenerateメソッドのprompt引数を指定した場合のみです。prompts_token_idsを指定してpromptを指定しない場合、Noneと表示されますので、ご注意を。

デコードされた出力結果は、output.outputs[0].textに格納されています。

    user_messages.append(
        {"role": "assistant", "content": output.outputs[0].text}
    )
    end = time.process_time()
    ##
    input_tokens = len(output.prompt_token_ids)
    output_tokens = len(output.outputs[0].token_ids)

入力tokenはprompt_token_idsに、出力tokenはtoken_idsに代入されています。

    total_time = end - start
    tps = output_tokens / total_time
    print(f"prompt tokens = {input_tokens:.7g}")
    print(f"output tokens = {output_tokens:.7g} ({tps:f} [tps])")
    print(f"   total time = {total_time:f} [s]")
    return user_messages

ここは変わらずです。

コード全体はこちら

コメントアウトなどを除くと、こんな感じです。
ほとんど変わらずに書けました。よかった、よかった。

from vllm import LLM, SamplingParams
from typing import List, Dict
import time

model_id = "elyza/ELYZA-japanese-Llama-2-7b-instruct"

# トークナイザーとモデルの準備
model = LLM(
    model=model_id,
    dtype="auto",
    trust_remote_code=True,
)
tokenizer = model.get_tokenizer()

DEFAULT_SYSTEM_PROMPT = "あなたは誠実で優秀な日本人のアシスタントです。"

# generation params
max_new_tokens = 1024
generation_params = SamplingParams(
    temperature = 0.8,
    top_p = 0.95,
    top_k = 40,
    max_tokens = max_new_tokens,
    repetition_penalty = 1.1
)

def q(
    user_query: str,
    chat_history: List[Dict[str, str]]=None
):
    start = time.process_time()
    # messages
    messages = [
        {"role": "system", "content": DEFAULT_SYSTEM_PROMPT},
    ]
    user_messages = [
        {"role": "user", "content": user_query}
    ]
    if chat_history:
        user_messages = chat_history + user_messages
    messages += user_messages
    # generateion prompts
    prompt = tokenizer.apply_chat_template(
        conversation=messages,
        add_generation_prompt=True,
        tokenize=False
    )
    # 推論
    outputs = model.generate(
        prompt,
        sampling_params=generation_params,
    )
    # debug
    #print(outputs)
    output = outputs[0]
    print("--- prompt")
    print(output.prompt)
    print("--- output")
    print(output.outputs[0].text)
    user_messages.append(
        {"role": "assistant", "content": output.outputs[0].text}
    )
    end = time.process_time()
    ##
    input_tokens = len(output.prompt_token_ids)
    output_tokens = len(output.outputs[0].token_ids)
    total_time = end - start
    tps = output_tokens / total_time
    print(f"prompt tokens = {input_tokens:.7g}")
    print(f"output tokens = {output_tokens:.7g} ({tps:f} [tps])")
    print(f"   total time = {total_time:f} [s]")
    return user_messages

3. 試してみる

聞いてみる

聞いてみましょう。

chat_history = q("ドラえもんとはなにか")

質問と推論の間に、気になるINFOメッセージが出力されています。「WSLを検知したから、pin_memory=Falseにした。パフォーマンス落ちるよ」だと…。

INFO 01-09 23:57:53 llm_engine.py:275] # GPU blocks: 882, # CPU blocks: 512
WARNING 01-09 23:57:53 cache_engine.py:96] Using 'pin_memory=False' as WSL is detected. This may slow down the performance.
INFO 01-09 23:57:53 model_runner.py:501] Capturing the model for CUDA graphs. This may lead to unexpected consequences if the model is not static. To run the model in eager mode, set 'enforce_eager=True' or use '--enforce-eager' in the CLI.
INFO 01-09 23:57:53 model_runner.py:505] CUDA graphs can take additional 1~3 GiB memory per GPU. If you are running out of memory, consider decreasing `gpu_memory_utilization` or enforcing eager mode.
INFO 01-09 23:57:55 model_runner.py:547] Graph capturing finished in 2 secs.

ここでは気にせず、推論結果を確認します。

秒間55.2トークン。これは、、、速い。
transformersの推論にかかる時間は、以下のように同じモデルで秒間22.9～24.3トークンでしたので2倍は速い。

GPUリソース使用状況

0GBから13.1GB、13.1GBから21.1GBと2段階でメモリ確保しています。

ログの時間から推測するに、1つめはLLM engineの初期化処理、2つめはCUDA graphsの処理で使用しているようです。

INFO 01-09 23:58:57 llm_engine.py:70] Initializing an LLM engine with config: model='elyza/ELYZA-japanese-Llama-2-7b-instruct', tokenizer='elyza/ELYZA-japanese-Llama-2-7b-instruct', tokenizer_mode=auto, revision=None, tokenizer_revision=None, trust_remote_code=True, dtype=torch.float16, max_seq_len=4096, download_dir=None, load_format=auto, tensor_parallel_size=1, quantization=None, enforce_eager=False, seed=0)
INFO 01-09 23:59:10 llm_engine.py:275] # GPU blocks: 882, # CPU blocks: 512
WARNING 01-09 23:59:10 cache_engine.py:96] Using 'pin_memory=False' as WSL is detected. This may slow down the performance.
INFO 01-09 23:59:10 model_runner.py:501] Capturing the model for CUDA graphs. This may lead to unexpected consequences if the model is not static. To run the model in eager mode, set 'enforce_eager=True' or use '--enforce-eager' in the CLI.
INFO 01-09 23:59:10 model_runner.py:505] CUDA graphs can take additional 1~3 GiB memory per GPU. If you are running out of memory, consider decreasing `gpu_memory_utilization` or enforcing eager mode.
INFO 01-09 23:59:12 model_runner.py:547] Graph capturing finished in 2 secs.

4. まとめ

数行の修正でこれまでのコードが使用でき、スピードアップが図れるのであれば、これを採用しないなんて選択肢はないです。はい。

5. おまけ

pinning memoryとは

CPUのメモリ領域がページングされないようにする（＝ピン留め）機能なのですが、WSLではPinning memoryはサポートされていません…。Windowsの配下でLinuxが動いているから、メモリ管理の機構として難しいというか無理なんでしょうね。おそらく。

generateメソッドの返却値の例

引数にpromptを指定せずにprompt_token_idsを指定した場合の、変数の内容です。promptがNoneになっているのがわかるかと思います。
あと、見やすいように成形しています。こまったときは、printfデバッグ、ということで。

>>> print(outputs)
[ RequestOutput(
    request_id = 0,
    prompt = None,
    prompt_token_ids = [ 1,
    518,
    25580,
    29962,
    3532,
    14816,
    29903,
    6778,
    13,
    30641,
    30371,
    30366,
    30449,
    235,
    173,
    163,
    31525,
    30499,
    232,
    135,
    173,
    31701,
    30371,
    30325,
    30346,
    30313,
    30199,
    30310,
    30373,
    30255,
    30369,
    30203,
    30279,
    30499,
    30427,
    30267,
    13,
    29966,
    829,
    14816,
    29903,
    6778,
    13,
    13,
    30335,
    30281,
    30914,
    30723,
    30389,
    30364,
    30449,
    30371,
    30353,
    30412,
    518,
    29914,
    25580,
    29962 ],
    prompt_logprobs = None,
    outputs = [ CompletionOutput(
        index = 0,
        text = '  承知しました。ドラえもんとは、藤子・F・不二雄による日本の漫画です。1970年代後半から1980年代前半にかけて「月刊コロコロコ ミック」や「小学館の幼稚園」等の雑誌に連載されました。主人公の少年・のび太の活躍を通じて、冷静かつ客観的に現代社会の問題点 を浮き彫りにすることが特徴です。ただし、実在の人物や企業との copyright infringement (著作権侵害) を防ぐため、登場する人物や企業は必ず変形やモチーフに改変を加えています。',
        token_ids = [ 29871,
        29871,
        233,
        140,
        194,
        31043,
        30326,
        30441,
        30326,
        30366,
        30267,
        30335,
        30281,
        30914,
        30723,
        30389,
        30364,
        30449,
        30330,
        30804,
        30319,
        30290,
        29943,
        30290,
        30413,
        30685,
        31322,
        30353,
        30787,
        30332,
        30325,
        30346,
        30199,
        233,
        191,
        174,
        31046,
        30499,
        30427,
        30267,
        29896,
        29929,
        29955,
        29900,
        30470,
        30690,
        31220,
        232,
        144,
        141,
        30412,
        30513,
        29896,
        29929,
        29947,
        29900,
        30470,
        30690,
        30658,
        232,
        144,
        141,
        30353,
        30412,
        30807,
        30466,
        30481,
        30534,
        232,
        139,
        141,
        30459,
        30378,
        30459,
        30378,
        30459,
        30627,
        30317,
        30305,
        30482,
        31111,
        30481,
        30446,
        30415,
        31161,
        30199,
        232,
        188,
        191,
        234,
        171,
        157,
        31179,
        30482,
        31184,
        30199,
        236,
        158,
        148,
        235,
        173,
        143,
        30353,
        31692,
        235,
        191,
        140,
        30566,
        30553,
        30441,
        30326,
        30366,
        30267,
        30888,
        30313,
        30539,
        30199,
        31022,
        30470,
        30290,
        30199,
        31298,
        30654,
        30199,
        31704,
        235,
        189,
        144,
        30396,
        30768,
        31115,
        30466,
        30330,
        232,
        137,
        186,
        236,
        160,
        156,
        30412,
        30773,
        31915,
        235,
        169,
        182,
        30210,
        30353,
        31928,
        30690,
        30564,
        30437,
        30199,
        232,
        152,
        146,
        236,
        164,
        143,
        30940,
        30396,
        233,
        184,
        177,
        30538,
        232,
        192,
        174,
        30453,
        30353,
        30427,
        30332,
        30589,
        30364,
        30458,
        31141,
        232,
        193,
        183,
        30499,
        30427,
        30267,
        30366,
        30955,
        30326,
        30330,
        31525,
        30505,
        30199,
        30313,
        30834,
        31111,
        231,
        191,
        132,
        31564,
        30364,
        30199,
        3509,
        1266,
        297,
        1341,
        292,
        882,
        313,
        235,
        148,
        154,
        30732,
        233,
        171,
        172,
        231,
        193,
        184,
        232,
        177,
        182,
        29897,
        29871,
        30396,
        236,
        155,
        181,
        31907,
        30366,
        30954,
        30330,
        31451,
        31045,
        30427,
        30332,
        30313,
        30834,
        31111,
        231,
        191,
        132,
        31564,
        30449,
        31641,
        31761,
        31786,
        31305,
        31111,
        30761,
        30656,
        30185,
        30423,
        30353,
        31264,
        31786,
        30396,
        30666,
        30914,
        30466,
        30298,
        30441,
        30427,
        30267,
        2 ],
        cumulative_logprob = -95.97567919455469,
        logprobs = None,
        finish_reason = stop
    ) ],
    finished = True
) ]

6. GPUx2を試してみる

CUDA_VISIBLE_DEVICESでGPU_IDを複数指定しても1枚しか使用しなかったので、7Bモデルで試していたのですが、LLM初期化処理vllm/vllm/engine/llm_engine.pyを眺めていて、tensor_parallel_sizeパラメータを指定すればよいことがわかりました。

コードの修正

弊環境は2枚なのでtensor_parallel_size = 2 としました。

model_id="elyza/ELYZA-japanese-Llama-2-13b-instruct"

# トークナイザーとモデルの準備
model = LLM(
    model=model_id,
    dtype="auto",
    trust_remote_code=True,
    tensor_parallel_size=2,
    max_model_len=256
)

また、max_model_lenのデフォルトがELYZAのこのモデルは4k(4096)となっています。が、このままですと以下のように「KVキャッシュ用のメモリが確保できん」とLLM初期化処理中にエラーとなってしまったため、適当な値を指定しています。

ValueError: The model's max seq len (4096) is larger than the maximum number of tokens that can be stored in KV cache (32). Try increasing `gpu_memory_utilization` or decreasing `max_model_len` when initializing the engine.

CUDA_VISIBLE_DEVICESも忘れずに

export CUDA_VISIBLE_DEVICES=0

と1つしか指定していないと、

ValueError: The number of required GPUs exceeds the total number of available GPUs in the cluster.

「tensor_parallel_sizeの値との整合が合わん！けしからん！」とエラーになりますので、ご注意を。

聞いてみる

chat_history = q("ドラえもんとはなにか")

続きを聞きます。

chat_history = q("続きを教えてください", chat_histor

秒間26.3～27.3トークンで、7Bだけでなく13Bでも速くなっています。
transformersの推論にかかる時間は、以下のように同じモデルで秒間14.1～16.1トークンでした。ダブルスコアとはなりませんでしたが、それでも1.6～1.8倍は速い。

関連