WSL2でvLLMをコードを読みつつ試してみる
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倍は速い。
関連
この記事が気に入ったらサポートをしてみませんか?