vLLMを利用したLLM推論高速化テクニック

皆さんこんにちは
Acroquestのデータサイエンスチーム「YAMALEX」のチームリーダ、@tereka114です。
YAMALEXチームでは、コンペティションへの参加や自社製品開発、技術研究などに日々取り組んでいます。

大規模言語モデル（通称：LLM）は近年、非常に注目される技術となりました。
ただ、7Bや13Bといった巨大モデルのパラメータは推論時間も長時間で計算時間の面からも運用が非常に難しいです。
しかし、vLLMを使えば、高速化できます。本記事では、推論をどこまで高速化できるのかを検討したいと思います。

※本記事はLLM・LLM活用のAdvent Calendar 24日目の記事です。
qiita.com

vLLMとは？
vLLMによる高速化実践
その他
- CPU Offload
最後に

vLLMとは？

vLLMは大規模言語モデル（LLM）の推論とサービングを高速かつ効率的に行うためのライブラリです。
計算コストの高いTransformerで利用されるAttentionについて、Paged Attentionと呼ばれるアルゴリズムで計算することにより、高速化を実現しています。

それ以外のvLLMの良いポイントとして、
1. 有名どころのモデルはカスタム実装不要で計算が可能
vLLMは、Huggingfaceの主要モデルの多くについて、推論をサポートしているため、簡単に扱えることです。
主要なモデルであれば、モデル特有の部分の実装は不要です。
また、量子化手法であるGPTQ, AWQに加え、最新版であればBitsandBytesにも対応しており、、巨大モデルを推論できる可能性が高いです。

2. リソース管理が容易
複数のGPUが存在していた場合もライブラリ側がコントロールするため、マシンのリソースを殆ど意識することなく、使える点が便利です。
また、CPUへのオフロード機能が備わっているため、GPUリソースで足りない場合にもvLLMを使って高速化できる見込みがあります。

github.com

vLLMによる高速化実践

vLLMを利用しない場合

まずは、Huggingfaceを利用した実装を使います。
今回は質疑応答データセットであるJSQuADテストデータ、4482件のデータに対して推論を実行します。

huggingface.co

モデルは多言語で学習され、日本語に対応しているQwen2.5の7Bを利用します。

huggingface.co

from datasets import load_dataset
from tqdm.auto import tqdm
from transformers import AutoTokenizer, AutoModelForCausalLM
 
dataset = load_dataset("shunk031/JGLUE", name="JSQuAD")
model_name = "Qwen/Qwen2.5-7B-Instruct"
df = dataset["validation"].to_pandas()
df["assistant_content"] = df.apply(lambda x: "\nQuestion:" + x["question"] + "\nContext: " + x["context"] + "\nOutput:", axis=1)

tokenizer = AutoTokenizer.from_pretrained(model_name)
prompts = [tokenizer.apply_chat_template([{"role": "system", "content": "Contextを参考にして、問題に回答してください。"}, 
                                          {"role":"user", "content": content}], tokenize=False) for content in df["assistant_content"].to_list()]
 
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    device_map="auto",
)
 
inputs = tokenizer(prompts, return_tensors="pt", padding=True)
 
generated_tokens = []
for text in tqdm(texts):
    inputs = tokenizer(text, return_tensors="pt", padding=False)
    generated_tokens.append(model.generate(
        input_ids=input_id.cuda().reshape(1, -1),
        attention_mask=attention_mask.cuda().reshape(1, -1),
        max_new_tokens=100
    ))

24h経過しても終了しませんでしたが、進捗具合から、おおよそ、92h必要になるようです。
これを基準にvLLMを使って高速化します。

vLLMを利用する

vLLMを利用する場合、次のとおりです。
vLLM特有の部分は非常にシンプルで、基本、LLMクラスで初期化して、generateで生成します。

from datasets import load_dataset
from vllm import LLM, SamplingParams
from transformers import AutoTokenizer, AutoModelForCausalLM

dataset = load_dataset("shunk031/JGLUE", name="JSQuAD")
model_name = "Qwen/Qwen2.5-7B-Instruct"
df = dataset["validation"].to_pandas()
df["assistant_content"] = df.apply(lambda x: "\nQuestion:" + x["question"] + "\nContext: " + x["context"] + "\nOutput:", axis=1)

tokenizer = AutoTokenizer.from_pretrained(model_name)
prompts = [tokenizer.apply_chat_template([{"role": "system", "content": "Contextを参考にして、問題に回答してください。"}, 
                                          {"role":"user", "content": content}], tokenize=False) for content in df["assistant_content"].to_list()]
 
llm = LLM(
    model=model_name,
    gpu_memory_utilization=0.95,
    dtype='half',
    enforce_eager=True
)
 
sampling_params = SamplingParams(
        temperature=0.0,
        top_p=0.9,
        max_tokens=1000
    )
outputs = llm.generate(prompts, sampling_params, use_tqdm=True)

この実装で、281秒です。前回の92時間といったスケールとは大きく異なり、現実的に計算できる時間になってますね。

vLLMに加えてAWQを利用して量子化する

vLLMはAWQやGPTQをはじめとしたLLMの量子化テクニックを使った推論が可能です。
AWQ/GPTQを利用することで、GPUメモリの消費量を大きく削減できるため、限られたリソースの中で、巨大なモデルを高速に推論できるメリットがあります。
AWQでのvLLMの推論を利用するには次のとおりです。事前にAWQに量子化すると読み込みのみの時間になるため、おすすめします。

huggingface.co

llm = LLM(
    model=model_name,
    quantization="awq",
    gpu_memory_utilization=0.95,
    dtype='half',
    enforce_eager=True,
)

こちらは、360秒です。vLLMをAWQなしで実施する場合と比較して、遅くはなりました。
しかし、GPUメモリが削減されることもあり、巨大モデルを有効に利用できます。

Auto Prefix Caching

「以下の質問に回答してください。」など、先頭に必ず含まれるプロンプトが存在する場合、そのプロンプト分の計算結果を使い回すことで計算時間を削減することが可能です。
この機能を実現しているのがAuto Prefix Cachingです。
例えば、特定のタスクに関してバッチで解く場合、同じプロンプトになる箇所が含まれるため、計算量の削減が期待できます。

実装も非常に簡単で、LLMの初期化のオプションを追加するのみです。

llm = LLM(
    model=model_name,
    gpu_memory_utilization=0.95,
    dtype='half',
    enforce_eager=True,
    enable_prefix_caching=True
)

今回はOne-Shot Sampleを想定して、先頭に問題と答えのサンプルを載せて精度向上を試みます。
この先頭のサンプルが毎回の推論で同じになるので、この部分の高速化を本オプションで期待できます。

HEAD_PROMPT = "\n例\nQuestion: 新たに造られた語のことを新語または何という？\nContext:造語 [SEP] 造語（ぞうご）は、新たに語（単語）を造ることや、既存の語を組み合わせて新たな意味の語を造ること、また、そうして造られた語である。新たに造られた語については、新語または新造語とも呼ばれる。\nOutput: 新造語"
 
prompts = [tokenizer.apply_chat_template([{"role": "system", "content": "Contextを参考にして、問題に回答してください。" + HEAD_PROMPT}, 
                                          {"role":"user", "content": content}], tokenize=False) for content in df["assistant_content"].to_list()]

データのサンプル

<|im_start|>system\nContextを参考にして、問題に回答してください。\n例\nQuestion: 新たに造られた語のことを新語または何という？\nContext:造語 [SEP] 造語（ぞうご）は、新たに語（単語）を造ることや、既存の語を組み合わせて新たな意味の語を造ること、また、そうして造られた語である。新たに造られた語については、新語または新造語とも呼ばれる。\nOutput: 新造語<|im_end|>\n<|im_start|>user\n\nQuestion:日本で梅雨がないのは北海道とどこか。\nContext: 梅雨 [SEP] 梅雨（つゆ、ばいう）は、北海道と小笠原諸島を除く日本、朝鮮半島南部、中国の南部から長江流域にかけての沿海部、および台湾など、東アジアの広範囲においてみられる特有の気象現象で、5月から7月にかけて来る曇りや雨の多い期間のこと。雨季の一種である。\nOutput:<|im_end|>\n

同じ条件でenable_prefix_cachingのTrue/Falseで次のような性能変化があります。
296s→189s

その他

CPU Offload

vLLMはGPUメモリに乗り切らない場合、CPUのオフロード機能を利用した推論が可能です。
例えば、この機能により、AWQを利用しても載らない24GB程度のGPUメモリでQwen2.5-72Bなどの大きなモデルの推論ができるので役立ちます。
実装は次のとおり、パラメータの設定を変更するのみで対応ができます。ただし、全てをGPUに乗せて推論するほどの性能は出ません。試しにQwen2.5の7Bで実施すると1200秒（約4-5倍）の計算時間が必要となりました。

llm = LLM(
    model=model_name,
    quantization="awq",
    gpu_memory_utilization=0.95,
    dtype='half',
    enforce_eager=True,
    cpu_offload_gb=8,
    swap_space=1,
)

最後に

本記事ではvLLMの高速化実践のテクニックをご紹介しました。
推論に関する有用な機能が様々使えますので、vLLMを有効に活用して快適なLLMライフを送りましょう。

Acroquest Technologyでは、キャリア採用を行っています。

Azure OpenAI／Amazon Bedrock等を使った生成AIソリューションの開発

ディープラーニング等を使った自然言語／画像／音声／動画解析の研究開発

マイクロサービス、DevOps、最新のOSSやクラウドサービスを利用する開発プロジェクト

書籍・雑誌等の執筆や、社内外での技術の発信・共有によるエンジニアとしての成長

少しでも上記に興味を持たれた方は、是非以下のページをご覧ください。
www.wantedly.com