先日書いた記事、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. 環境構築
python3 -m venv longwriter
cd $_
パッケージのインストール。
pip install vllm
2. 流し込むコード
以下の内容を /path/to/query4vllm.py として保存します。
import sys
import argparse
from vllm import LLM, SamplingParams
from typing import List, Dict
import time
parser = argparse.ArgumentParser()
parser.add_argument("--model-path", type=str, default=None)
parser.add_argument("--no-chat", action='store_true')
parser.add_argument("--no-use-system-prompt", action='store_true')
parser.add_argument("--max-model-len", type=int, default=32768)
parser.add_argument("--tensor-parallel-size", type=int, default=1)
parser.add_argument("--gpu-memory-utilization", type=float, default=0.2)
parser.add_argument("--max-tokens", type=int, default=4096)
args = parser.parse_args(sys.argv[1:])
model_id = args.model_path
if model_id == None:
exit
is_chat = not args.no_chat
use_system_prompt = not args.no_use_system_prompt
max_new_tokens = args.max_tokens
tensor_parallel_size = args.tensor_parallel_size
max_model_len = args.max_model_len
gpu_memory_utilization = args.gpu_memory_utilization
model = LLM(
model=model_id,
dtype="auto",
trust_remote_code=True,
tensor_parallel_size=tensor_parallel_size,
max_model_len=max_model_len,
gpu_memory_utilization=gpu_memory_utilization
)
tokenizer = model.get_tokenizer()
DEFAULT_SYSTEM_PROMPT = "あなたは誠実で優秀な日本人のアシスタントです。"
def q(
user_query: str,
history: List[Dict[str, str]]=None
) -> List[Dict[str, str]]:
generation_params = SamplingParams(
temperature=0.8,
top_p=0.95,
top_k=40,
max_tokens=max_new_tokens,
repetition_penalty=1.1
)
start = time.process_time()
messages = ""
if is_chat:
messages = []
if use_system_prompt:
messages = [
{"role": "system", "content": DEFAULT_SYSTEM_PROMPT},
]
user_messages = [
{"role": "user", "content": user_query}
]
else:
user_messages = user_query
if history:
user_messages = history + user_messages
messages += user_messages
if is_chat:
prompt = tokenizer.apply_chat_template(
conversation=messages,
add_generation_prompt=True,
tokenize=False
)
else:
prompt = messages
input_ids = tokenizer.encode(
prompt,
add_special_tokens=True,
)
print("--- prompt")
print(prompt)
print("--- output")
outputs = model.generate(
sampling_params=generation_params,
prompt_token_ids=[input_ids],
)
print(outputs)
output = outputs[0]
print(output.outputs[0].text)
if is_chat:
user_messages.append(
{"role": "assistant", "content": output.outputs[0].text}
)
else:
user_messages += 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. 聞いてみる
pythonコマンド叩きます。
CUDA_VISIBLE_DEVICES=0 python -i /path/to/query4vllm.py \
--model-path THUDM/LongWriter-glm4-9b \
--max-model-len 126608 \
--max-tokens 1048576 \
--gpu-memory-utilization 0.99
config.json · THUDM/LongWriter-glm4-9b を見ると、
"seq_length": 1048576,
とあるので、max-tokensを1M (1048576) にしています。
それでは聞いてみましょう。
history = q("ドラえもんとのび太を主人公にした長編物語を、1000000文字以上の分量で緻密な表現で記述してください。")
prompt tokens = 76
output tokens = 9962 (41.575348 [tps])
total time = 239.613149 [s]
VRAM使用量は、gpu_memory_utilizationを0.99と指定したので23.3GB以上は使用されません。はい。
4. まとめ
上記で生成されたトークン数は9962。wc -c でカウントしたところ、39,968文字でした。つまり、400字詰め原稿用紙で換算すると、約100枚。これを239.6秒、約4分で生成したことになるので 1分あたり25枚です。
おまけに同じプロンプトで聞いた推論結果をもう一つ貼り付けておきます。
こちらは、28240文字(7127トークン)を165秒。
prompt tokens = 76
output tokens = 7127 (43.057704 [tps])
total time = 165.522061 [s]