Transformers高速化ライブラリvLLMのAsyncLLMEngineを利用した非同期高速文章生成

概要

先日までKaggleのAIMOコンペ(数学の問題をLLMに解かせて正答率を競う)に参戦していました。結果は初のチーム参加でメンバーに助けられつつ運もあり、なんとか銀メダルを取れました！これでMasterにリーチがかかりましたが、金メダルは未だ取れる気がしないので遠い道のりです……。

www.kaggle.com

このコンペについて、近い内に同様のコンペが開催予定なこともあり上位解法があまり出ていない状態なので、どのような手法が良かったのかまだわかっていないのですが、とりあえず公開されている情報を元にすると、

LLMとしてはほぼほぼ全員が数学問題に特化したLLMであるDeepseek-Math-7Bを利用している
LLMが出力したPythonコードを実行するインタープリターを実装することで、LLMのハルシネーションによる計算ミスを防ぐパイプラインが有力であった
LLMの出力を比較的高い温度で複数回サンプリングすることで、多様な解答を得つつ最も回答数が多いものを選択することで、精度を高めることができる

という点がベースラインとなっているコンペでした。

(2024/7/11 追記) 1st ソリューション公開されました。H100x8を使ったフルパラメータチューニングという常人には真似できない凄まじい解法！どうやらHFとMistral AIの人たちのチームらしい。

www.kaggle.com

その中で、最後の点について、通常のHugging Face Transformersではなく、vLLMというライブラリを利用することで文章生成を高速化し、サンプリング回数を増やすことができるというディスカッションがありました。

www.kaggle.com

全くこのvLLMの存在を知らなかったのですが、使ってみるとたしかに速い。更に、非同期推論用のラッパーであるAsyncLLMEngineを利用することで、GPUを遊ばせることなくフルに動かすことができ、通常のTransformersの最大20倍以上の速度で文章生成を行うことができました。

ただ、vLLMのドキュメントが少なく、特に日本語情報は殆どなさそう。更に、AsyncLLMEngineの使い方に至っては公式ドキュメントにもわかりやすい動作例がないため、少しvLLMの使い方をまとめることとします。

まとめ

ここに動作サンプル(Kaggle Notebook)を用意しましたので、単に使うだけならこれをご参照ください。

www.kaggle.com

いま私のソリューションは書いている途中なので、これもできたら共有します。

(20240707追記)ソリューション公開しました。

www.kaggle.com

今後、KaggleでTransformersを使うときには、推論時にvLLMを使うことで大幅に時間短縮が可能になるかもしれませんので、試す価値アリだと思います。

vLLMについて

Geminiに作らせた要約を一部訂正したものです。

高速: 最先端のサービングスループット、PagedAttentionによる効率的なアテンションキーと値メモリの管理、受信リクエストの連続バッチング、CUDA/HIPグラフと量子化による高速なモデル実行により、高速な処理を実現します。

柔軟性と使いやすさ: 人気のHugging Faceモデルとのシームレスな統合、並列サンプリング、ビームサーチなどのさまざまなデコーディングアルゴリズムを備えた高スループットサービング、分散推論用のテンソル並列サポート、ストリーミング出力、OpenAI互換APIサーバー、NVIDIA GPU、AMD GPU、Intel CPUとGPUのサポート、接頭辞キャッシュサポート（実験的）、マルチloraサポート(実験的)など、柔軟で使いやすい設計になっています。

「人気のHugging Faceモデルとのシームレスな統合」ということで、ほとんどHuggingFaceのTransformersと同じAPIを持ちつつ、裏側でPagedAttention等の高速化技術を使って推論速度を上げてくれるような機能も持っており、最小限の変更でノートの実行時間を減らすこともできます。

公式Exampleより抜粋。一番シンプルなLLM APIを使うことで対応しているモデルであれば、"facebook/opt-125m"のようにHuggingFace Hubのパス指定で当該のモデルを呼び出し、HF transformersのgenerateメソッドっぽく文章生成をすることができます。

from vllm import LLM, SamplingParams
sampling_params = SamplingParams(temperature=0.8, top_p=0.95)

llm = LLM(model="facebook/opt-125m")
outputs = llm.generate(prompts, sampling_params)

近いライブラリというと、Deepspeedが挙げられるかもしれません。

インストール

活発にアップデートされているライブラリなのですぐにインストール方法が変わりそうですが、とりあえず20240706現在、0.4.0post1をKaggleでインストールする方法を載せておきます。

Kaggle DatasetにWhlがあるので、これを使ってインストールします

!pip uninstall -y torch
!pip install --no-index --find-links=/kaggle/input/vllm-whl -U vllm
!pip install /kaggle/input/vllm-additional-packages/grpcio-1.60.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl

一般的なインストールのガイダンスは次を参照してください。

docs.vllm.ai

AsyncLLMEngineについて

今回フォーカスするのは、非同期推論用のAPIであるAsyncLLMEngineです。

とりあえず公式ドキュメントの記載は以下にあります。

docs.vllm.ai

非同期推論なので、ここに載っているExampleをそのまま使うことはできず、async/awaitを使ったコードを書く必要があります。

そこで、簡単に動作確認ができるサンプルを作成しました。このサンプルは、2つのプロンプトに対して100回ずつサンプリングを行い、その結果をファイルに書き出すというものです。

from vllm import AsyncLLMEngine, SamplingParams
from vllm.engine.arg_utils import AsyncEngineArgs

import asyncio
import time
import uuid

st = time.time()


example_inputs = [
    {
        "prompt": "About 200 words, please give me some tourist information about Tokyo.",
        "temperature": 0.9,
    },
    {
        "prompt": "About 200 words, please give me some tourist information about Osaka.",
        "temperature": 0.9,
    },
]


async def gen(engine, example_input, id):
    results_generator = engine.generate(
        example_input["prompt"],
        SamplingParams(temperature=example_input["temperature"],max_tokens=300, min_tokens=200,),
        id,
    )

    final_output = None

    async for request_output in results_generator:
        final_output = request_output

    prompt = final_output.prompt
    text_output = [output.text for output in final_output.outputs]
    return text_output[0]


async def main():
    engine = AsyncLLMEngine.from_engine_args(
        AsyncEngineArgs(
            model="/kaggle/input/llama-2/pytorch/7b-chat-hf/1",
            dtype="half",
            enforce_eager=True,
            gpu_memory_utilization=0.99,
            swap_space=3,
            max_model_len=1024,
            kv_cache_dtype="fp8_e5m2",
            tensor_parallel_size=2,
            disable_log_requests=True
        )
    )

    results = []

    for example_input in example_inputs:
        tasks = []
        for i in range(100):
            tasks.append(asyncio.create_task(gen(engine, example_input, uuid.uuid4())))

        res = [await task for task in tasks]

        results.append(res)

    with open("async_res.txt", "w") as f:
        for r in results:
            f.writelines(r)


if __name__ == "__main__":
    asyncio.run(main())

    print("Async vLLM inference time: ", time.time() - st)

パラメータについて

AsyncEngineArgsというクラスでLLMの設定、プロンプトとともに渡すSamplingParamsでサンプリングの設定を行うことが可能です。

詳しくは公式ドキュメントを参照いただければと思いますが、とりあえずKaggle上でT4x2を利用する際に、私は次のような設定を利用していました。

engine = AsyncLLMEngine.from_engine_args(
    AsyncEngineArgs(
        model=MODEL_NAME,
        dtype="half",  # ここは指定しなくてもturing世代なら自動的にfp16になるらしい。より新しいGPUならbfloat16をデフォルトでは使ってくれる
        enforce_eager=True,  # Pytorchのeager modeを強制するかどうか。これは参考にしたノートブックがTrueにしていたのでそのまま利用している。
        gpu_memory_utilization=0.99,  # KVキャッシュ含めGPUのメモリをどこまで使うか。Kaggleであれば0.99でも問題ないが、もし映像出力も行っているGPUを使う場合はここを減らさないとマシン自体が落ちるので注意。
        swap_space=3,  # デフォルトは4だが、メインメモリのOOMで落ちたことがあったので3にしている
        kv_cache_dtype="fp8_e5m2",  # auto, fp8, fp8_e5m2, fp8_e4m3から選べる。fp8_e5m2のほうが<s>仮</s>指数部が大きいので<s>精度が高いらしい</s>数値のレンジが大きい。 (20240708)訂正。fp8_e5m2を使っているというディスカッションがあったのでそれに則ってこちらを使っているが、どっちが良いかはケース次第？
        tensor_parallel_size=2,  # ここの数を変えるだけで複数GPUを利用してくれる。
        max_model_len=1024,
        disable_log_requests=True,  # これをTrueにしないとログがめちゃくちゃ出る
    )
)

AIMOコンペでは、コードフェンス内のPythonコードをcode.pyというファイルに書き出して、そのコードをsubprocessで実行するという手法が取られていたので、stop_wordsによる生成停止を行います。

Transformersと似たようにリスト型でstop_wordsが可能ですが、include_stop_str_in_outputをTrueにすることでstop_wordsを含めた文章を返してくれるようになり、Transformersの標準的な出力と互換性を保つことができます。

SamplingParams(
    temperature=0.9,
    top_p=1,
    max_tokens=2048,
    stop=[
        "```output",
        "```python",
        "```\nOutput",
        ")\n```",
        "``````output",
    ],
    include_stop_str_in_output=True,
)