vLLMでGGUFフォーマットのモデルファイルを利用する

最近体調崩していて、久しぶりの投稿。リハビリ的な内容です。

導入

ここ最近、ローカルLLM界隈ではLlama.cppやollamaなどがデファクトの推論エンジンになってきているように思います。
これらはGGUFというファイルフォーマットのモデルファイルに対応しています。

GGUFはLlama.cpp用に拡張性高く設計されたファイルフォーマットで、量子化されたモデルも内包することができます。（この表現は微妙に間違ってそうですが。。。）
そのため、非量子化モデルを量子化＆GGUFフォーマットで保管・huggingface上に公開、というケースも多くあります。

ただ、GGUFをサポートしているパッケージは限定的で、個人的にvLLMや最近も取り上げたSGLangでは利用することができません。

vLLMでもGGUFの読み込みをサポートする意見が多く見られていたのですが、最近Nightly版でGGUFの読み込みができるようになったようです。

まだStable版には実装されていないのですが、次のバージョンあたりで来そうなので先に試してみます。

この機能はまだ安定板にマージされていないため、仕様が変更される可能性があります。
参考にされる場合はご注意ください。

検証はDatabricks on AWS上で行いました。
DBRは15.4ML LTS(Beta)、インスタンスタイプはg5.xlargeです。

Step1. パッケージインストール

ノートブックを作成して、vLLMのNightly版をインストール。
現時点で通常のvLLM(現時点のバージョンは0.5.4)では動作しないため、必ずNightly版をインストールしましょう。

%pip install https://vllm-wheels.s3.us-west-2.amazonaws.com/nightly/vllm-0.5.4-cp38-abi3-manylinux1_x86_64.whl
dbutils.library.restartPython()

Step2. モデル読込と推論

こちらの内容を、モデルや一部のパラメータを変更して実行します。

今回は以下のモデルを利用しました。
Llama3.1を日本語ファインチューニングした非常に優秀なモデルです。

かつ、量子化&GGUFフォーマットで保管された以下のものを利用させていただきます。

from huggingface_hub import hf_hub_download
from vllm import LLM, SamplingParams

def run_gguf_inference(model_path):
    PROMPT_TEMPLATE = "<|start_header_id|>system<|end_header_id|>\n\n{system_message}<|eot_id|><|start_header_id|>user<|end_header_id|>\n\n{prompt}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
    system_message = "あなたは親切なAIアシスタントです"  # noqa: E501
    # Sample prompts.
    prompts = [
        "日本の総理大臣は誰？",
        "Databricksとは何ですか？",
    ]
    prompts = [
        PROMPT_TEMPLATE.format(system_message=system_message, prompt=prompt)
        for prompt in prompts
    ]
    # Create a sampling params object.
    sampling_params = SamplingParams(temperature=0, max_tokens=128)

    # Create an LLM.
    llm = LLM(
        model=model_path,
        tokenizer="HODACHI/Llama-3.1-8B-EZO-1.1-it",
        gpu_memory_utilization=0.7,
        max_model_len=8128,
    )

    outputs = llm.generate(prompts, sampling_params)
    # Print the outputs.
    for output in outputs:
        prompt = output.prompt
        generated_text = output.outputs[0].text
        print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")


if __name__ == "__main__":
    repo_id = "mmnga/Llama-3.1-8B-EZO-1.1-it-gguf"
    filename = "Llama-3.1-8B-EZO-1.1-it-Q4_K_M.gguf"

    model = hf_hub_download(repo_id, filename=filename)
    run_gguf_inference(model)

実行結果

Prompt: '<|start_header_id|>system<|end_header_id|>\n\nあなたは親切なAIアシスタントです<|eot_id|><|start_header_id|>user<|end_header_id|>\n\n日本の総理大臣は誰？<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n', Generated text: '現在の日本の総理大臣（内閣総理大臣）は岸田文雄です。ただし、情報はリアルタイムで更新されていないため、最新情報については、2023年4月以降の変動が存在する可能性があります。最新情報については、official websiteやニュースサイトをご参照ください。'
Prompt: '<|start_header_id|>system<|end_header_id|>\n\nあなたは親切なAIアシスタントです<|eot_id|><|start_header_id|>user<|end_header_id|>\n\nDatabricksとは何ですか？<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n', Generated text: 'Databricksは、Apache Sparkを中心としたデータプラットフォームを提供するアメリカの企業です。2013年に創業され、現在では大規模なデータ分析と機械学習を支援するクラウドベースのサービスとして広く利用されています。\n\nDatabricksの主な機能には、次の点が含まれます。\n\n1. **Apache Sparkの統合**: DatabricksはSparkをベースに開発されており、Sparkの最新バージョンを直接サポートしています。ユーザーはSparkのパワーを最大限に活用できる'

正しくGGUFファイルからモデルをロードし、vLLMで推論することができました！

まとめ

vLLMでGGUF読込機能を試してみました。
まだ一部の基盤モデル(LlamaやMistralなど)しかサポートしてなさそうなのですが、主要なモデルのGGUFファイルだけでも読み込めるのは助かります。
また、vLLMがサポートするということは、SGLangでも使えるようになるんじゃないかと期待しています。
(SGLangのモデル読込はvLLMの機能を利用しているはずなので)

LLMの世界でもモデルファイルの相互運用性がどんどん高まるとよいなと思います。