Zenn
Open3

vLLMを試してみる

LLM
vLLM
shukawamshukawam

vLLM

必要なライブラリをダウンロードする

pip install vllm

雑に、7Bのモデル（SakanaAI/EvoLLM-JP-A-v1-7B）を実行してみる

https://huggingface.co/SakanaAI/EvoLLM-JP-A-v1-7B

from vllm import LLM, SamplingParams

prompts = [
    "Kubernetesってなんですか？"
]

# LLM の推論環境に与えるパラメータ
sampling_params = SamplingParams(
    temperature=0.8,
    max_tokens=128
)

llm = LLM(
    model="SakanaAI/EvoLLM-JP-A-v1-7B"
)

outputs = llm.generate(
    prompts=prompts,
    sampling_params=sampling_params
)

for output in outputs:
    print(output.outputs[0].text)

OOM...

# ... omit
OutOfMemoryError: CUDA out of memory. Tried to allocate 224.00 MiB. GPU 0 has a total capacity of 21.98 GiB of which 146.44 MiB is free. Including non-PyTorch memory, this process has 21.82 GiB memory in use. Of the allocated memory 21.55 GiB is allocated by PyTorch, and 5.07 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.
# ... omit

メモリのアロケートを再設定する

# 追加
import os

os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"

モデルの最大シーケンス長とGPUメモリキャッシュ(KV cache)が一致していない模様

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

gpu_memory_utilization の調整(0.8 -> 0.9)では乗り切らなさそうなので、量子化してメモリ節約してみる。AWQ は vLLM では最適化されていないらしいけど、とりあえずの目的は、メモリフットプリントの削減なのでよしとする

Please note that AWQ support in vLLM is under-optimized at the moment. We would recommend using the unquantized version of the model for better accuracy and higher throughput. Currently, you can use AWQ as a way to reduce memory footprint. As of now, it is more suitable for low latency inference with small number of concurrent requests. vLLM’s AWQ implementation have lower throughput than unquantized version.

https://docs.vllm.ai/en/latest/quantization/auto_awq.html

pip install autoawq

この辺のコード足してみる

from awq import AutoAWQForCausalLM
from transformers import AutoTokenizer

model_path = "SakanaAI/EvoLLM-JP-A-v1-7B"
quantization_path = "SakanaAI/EvoLLM-JP-A-v1-7B-awq"
quantization_config = {
    "zero_point": True,
    "q_group_size": 128,
    "w_bit": 4,
    "version": "GEMM"
}

# モデルのロード
model = AutoAWQForCausalLM.from_pretrained(
    model_path=model_path,
    **{"low_cpu_mem_usage": True, "use_cache": False}
)

tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)

# 量子化
model.quantize(tokenizer, quant_config=quantization_config)

# 量子化したモデルを保存する
model.save_quantized(quantization_path)
tokenizer.save_pretrained(quantization_path)

AWQ での量子化は、今回のモデルだと 20 分くらいかかった

参考: 実行時のログ出力

Fetching 12 files: 100%|██████████| 12/12 [00:00<00:00, 148470.94it/s]

Loading checkpoint shards: 100%|██████████| 3/3 [00:03<00:00,  1.12s/it]
Repo card metadata block was not found. Setting CardData to empty.
WARNING:huggingface_hub.repocard:Repo card metadata block was not found. Setting CardData to empty.
AWQ: 100%|██████████| 32/32 [21:09<00:00, 39.66s/it]
('SakanaAI/EvoLLM-JP-A-v1-7B-awq/tokenizer_config.json',
 'SakanaAI/EvoLLM-JP-A-v1-7B-awq/special_tokens_map.json',
 'SakanaAI/EvoLLM-JP-A-v1-7B-awq/tokenizer.model',
 'SakanaAI/EvoLLM-JP-A-v1-7B-awq/added_tokens.json',
 'SakanaAI/EvoLLM-JP-A-v1-7B-awq/tokenizer.json')

モデルパス修正して再度推論

prompts = [
    "Kubernetesってなんですか？"
]

llm = LLM(
    model=quantization_path,
    quantization="AWQ"
)

outputs = llm.generate(
    prompts=prompts,
    sampling_params=sampling_params
)

for output in outputs:
    print(output.outputs[0].text)
# => Kubernetesは、Open Sourceのコンテナオーケストレーションシステムであり、コンテナ化されたアプリケーションのデプロイ、スケーリング、ヘルス管理を容易にします。 Kubernetesは、コンテナオーケストレーションの標準となり、多くのクラウドプロバイダー、コンテナランチャー、および各種ツー
shukawamshukawam

最終的なスクリプトは以下の通り

import os
from vllm import LLM, SamplingParams
from awq import AutoAWQForCausalLM
from transformers import AutoTokenizer

os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"

model_path = "SakanaAI/EvoLLM-JP-A-v1-7B"
quantization_path = "SakanaAI/EvoLLM-JP-A-v1-7B-awq"
quantization_config = {
    "zero_point": True,
    "q_group_size": 128,
    "w_bit": 4,
    "version": "GEMM"
}

# モデルのロード
model = AutoAWQForCausalLM.from_pretrained(
    model_path=model_path,
    **{"low_cpu_mem_usage": True, "use_cache": False}
)

tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)

# 量子化
model.quantize(tokenizer, quant_config=quantization_config)

# 量子化したモデルを保存する
model.save_quantized(quantization_path)
tokenizer.save_pretrained(quantization_path)

import torch
print(torch.cuda.device_count())

# PyTorchのキャッシュ解放
torch.cuda.empty_cache()

llm = LLM(
    model=quantization_path,
    quantization="awq_marlin"
)

prompts = [
    "Kubernetesってなんですか？"
]

# LLM の推論環境に与えるパラメータ
sampling_params = SamplingParams(
    temperature=0.8,
    max_tokens=512
)

outputs = llm.generate(
    prompts=prompts,
    sampling_params=sampling_params
)

for output in outputs:
    print(output.outputs[0].text)