vLLMによるLLM推論速度の向上

# はじめに
LLMを個人でも使ってみたい，研究で使う機会がある．そんなときに，問題になるのが推論速度．かく言う私もモデルの推論が遅いために待ち時間が長く，Try & Errorが気軽にできずに困っていました．


# vLLMって何？
[公式ドキュメント](https://docs.vllm.ai/en/stable/)には以下のように書いています．（抜粋）

> vLLM is a fast and easy-to-use library for LLM inference and serving
> vLLM is fast with:
> - Efficient management of attention key and value memory with **PagedAttention**

以下，DeepLによる日本語訳．

> vLLMはLLMの推論と処理のための高速で使いやすいライブラリです．
> vLLMは以下のために高速です：
> - **PagedAttention**によるAttentionのkeyとvalueのメモリの効率的な管理

どうやらPagedAttentionという技術を用いてAttention層のkeyとvalueを効率化することにより，高速化を達成しているようです．この辺りの詳細はまだ勉強できていないので，触れないことにします．
一応，Hugging Faceからモデルを利用する際に用いるTransformersライブラリにもこれに似た機能はあるようですが，その効率と実装の容易さでvLLMが用いられているようです．


# vLLMの実装を見てみる
本章での実装は以下を前提とします．

:::note
- Hugging Faceにアクセスできる
    + アカウントの発行
    + API Tokenの発行 など
- GPUが利用できる
    + `nvidia-smi`でGPUが認識できる
- 環境構築ができている
    + PyTorchやTransformers，vLLMが利用できる
:::

参考までに，私の環境を記載しておきます．

|項目|バージョンなど|追記事項|
|:-:|:-:|:-:|
|OS|Ubuntu 22.04LTS||
|GPU|GeForce RTX 3090|24GB|
|RAM||128GB|
|CPU|Ryzen 9 5950X||
|Docker|27.3.1||
|PyTorch|2.4.0||
|Transformers|4.45.2||
|vLLM|0.6.3.post1||

<details>
    <summary>全体コード</summary>
    
    import os
    import json 
    
    from dotenv import load_dotenv 
    
    import torch 
    import huggingface_hub 
    from vllm import SamplingParams, LLM
    
    
    ###################################################
    ########## Variables and Hyperparameters ##########
    ###################################################
    SEED = 42
    
    MODEL_NAME = 'rinna/japanese-gpt-1b'
    MAX_TOKENS = 256
    ###################################################
    
    # Initial Setup
    load_dotenv()
    huggingface_hub.login(token=os.getenv('HUGGINGFACE_ACCESS_TOKEN'))
    
    
    # Setup the configuration of vLLM
    conf_vllm = SamplingParams(
        presence_penalty=0.1,
        seed=SEED,
        max_tokens=MAX_TOKENS,
    )
    
    
    # Load the model and tokenizer
    llm = LLM(
        model=MODEL_NAME,
        seed=SEED, 
        gpu_memory_utilization=0.95,
        download_dir=os.getenv('HUGGINGFACE_CACHEDIR'),
    )
    
    
    # Load sample prompts
    with open('sample_prompts.json', 'r') as f:
        prompts = json.load(f)
        prompts = [prompt['text'] for prompt in prompts]
    
    
    # Generate outputs
    outputs = llm.generate(
        prompts=prompts,
        sampling_params=conf_vllm,
    )
    
    
    # Print the outputs
    for idx, (prompt, output) in enumerate(zip(prompts, outputs)):
        print(f'### {idx + 1} ###')
        print(f'Prompt: {prompt}')
        print(f'Output: {output.outputs[0].text}')
</details>

モデルのロードはTransformersとほとんど同じように記述できます．

```python: vLLM_model_load
# Load the model and tokenizer
llm = LLM(
    model=MODEL_NAME,
    seed=SEED, 
    gpu_memory_utilization=0.95,
    download_dir='モデルをDLするディレクトリ',
)
```

<details>
    <summary>Transformersでのモデル・トークナイザのロード</summary>

    # Download the model and tokenizer
    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, cache_dir='モデルをDLするディレクトリ')
    model = AutoModel.from_pretrained(MODEL_NAME, cache_dir='モデルをDLするディレクトリ')
</details>

vLLMで推論する際はリスト形式でプロンプトを与えることで，**入力した順番**の出力が得られます．

```python: vLLM_generate
# Generate outputs
outputs = llm.generate(
    prompts=prompts,
    sampling_params=conf_vllm,
)
```

<details>
    <summary>Transformersでのモデルの推論</summary>

    # Generate outputs
    outputs = []
    with torch.no_grad():
        for prompt in tqdm(prompts):
            input_ids = tokenizer.encode(prompt, add_special_tokens=False, return_tensors='pt')
            input_ids = input_ids.to(model.device)
            output_ids = model.generate(
                input_ids, 
                max_length=MAX_TOKENS,
                min_length=MAX_TOKENS,
                repetition_penalty=1.5,
                pad_token_id=tokenizer.pad_token_id,
                bos_token_id=tokenizer.bos_token_id,
                eos_token_id=tokenizer.eos_token_id,
                bad_words_ids=[[tokenizer.unk_token_id]],
            )
            output = tokenizer.decode(output_ids.tolist()[0])
            outputs.append(output)
</details>

非常に単純で，Transformersでの実装から少し変更を加えるだけなので，導入がしやすくてよいですね．


# 推論速度
気になる推論速度です．
使用したモデルは[rinna/japanese-gpt-1b](https://huggingface.co/rinna/japanese-gpt-1b)で，日本語対応の1Bサイズモデルです．推論に必要なvRAM（GPUメモリのサイズ）は大体モデルサイズの2倍から5倍程度[^1][^2]のようなので，上記のモデルであれば5GB程度のvRAMで動かせる試算です．

[^1]: https://huggingface.co/spaces/hf-accelerate/model-memory-usage
[^2]: https://note.com/kan_hatakeyama/n/nfb1ca0ef25f9

実験条件は次の通りです．

:::note
- プロンプト
    + 言語は日本語
    + 全部で24文です
- 大量のプロンプトでの実験は24文をコピーして入力します
    + 24文を用意するので力尽きました．許して下さい．
    + コピーしているのでvLLMに有利な実験になっていると思われます
- repetition penaltyを与えています
    + そのままでは単語を繰り返すだけの意味のない文章が出力されます
    + 今回は生成速度を見るだけなので与えなくてもよいのですが，せっかくなので意味のある文章を出力させます
- 時間計測
    + 1回だけ実施しています
        * 複数回の平均を取るほうが望ましいと思いますが，簡単な実験なのでご容赦下さい
    + 推論時間は推論に要した時間です
        * 時間の計測はtqdmによるものです
        * 精度については疑問が残りますが，手軽さを優先しました
    + 実行時間は推論も含めた全体の実行に要した時間です
        * モデルのロードが遅ければ意味がないので，全体の実行に要する時間も計測しています
        * 時間の計測はUbuntuにあるtimeコマンドによるものです
- モデルはあらかじめダウンロードしています
:::

というわけで，実験結果を以下に示します．

|プロンプト数|vLLM - 推論時間|vLLM - 実行時間|Transformers - 推論時間|Transformers - 実行時間|（参考）vLLM - 出力トークン速度|
|:-:|:-:|:-:|:-:|:-:|:-:|
|24|0 sec|25.526 sec|54 sec|58.629 sec|1511.58 toks/sec|
|48|1 sec|26.025 sec|1 min 49 sec|1 min 54.897 sec|1864.11 toks/sec|
|96|1 sec|26.719 sec|3 min 40 sec|3 min 44.691 sec|3951.25 toks/sec|
|144|1 sec|26.750 sec| --- | --- |4474.41 toks/sec|
|192|2 sec|27.373 sec| --- | --- |4915.15 toks/sec|

上記の表を見るとvLLMには以下の特徴があるようです．

- 一度に入力するプロンプト数が多いほど推論速度が上昇する
    + 上限はあるが比較的すぐに速度が上昇する
- モデルのロードには（Transformersと比べると）時間がかかる

実験設定がガバガバなので疑義が残るところではありますが，これまでの利用経験からも概ねこのような特徴があるかなと思います．
ちなみに，今回の実験では考慮していませんが，入力トークン長が長いと比較的速度が上がりやすいようです．


# 【補足】CPUオフロード
最後に追加実験です．
折角ならより性能の良いモデルを使いたいですよね．しかし，大きなサイズのモデルはvRAMの制限が厳しく扱いにくい．そんな時に使うCPUオフロードですが，一般にCPUオフロードをすると推論速度が遅くなり，場合によっては使い物にならないこともあります．
vLLMならそれも解決してくれるかもしれません．ということで，こちらを検証します．
以下に該当する部分のコードを示します．

```python: vllm_cpu_offload
# Load the model and tokenizer
llm = LLM(
    model=MODEL_NAME,
    seed=SEED, 
    gpu_memory_utilization=0.10,    # 2GB程度のみ利用可能（今回のモデルの推論に約4GB必要）
    cpu_offload_gb=16,              # RAMを16GB割り当て
    download_dir=os.getenv('HUGGINGFACE_CACHEDIR'),
)
```

vLLMにおいてCPUオフロードを使う際は，上記のように引数を設定するだけで良いです．また，今回の実験ではGPUだけで収まり切らないようにするために，利用可能なvRAMを再現します．（概ねモデルの半分がCPU RAMに乗るようにします．）

以下に実験結果を示します．

|プロンプト数|vLLM - 推論時間|vLLM - 実行時間|（参考）vLLM - 出力トークン速度|
|:-:|:-:|:-:|:-:|
|24|13 sec|47.615 sec|101.04 toks/sec|
|48|24 sec|58.412 sec|131.24 toks/sec|
|96|14 sec|47.668 sec|408.16 toks/sec|
|144|16 sec|50.800 sec|524.31 toks/sec|
|192|22 sec|55.587 sec|521.12 toks/sec|

今回の結果では10分の1程度の推論速度となりました．


# 最後に
いかがでしたでしょうか？
自分の勉強がてら記事を書いてみましたが，どなたかの助けになれば幸いです．


<details>
    <summary>与えたプロンプトとその結果の一覧</summary>

    ### 1 ###
    Prompt: 機械学習とディープラーニングの違いは何ですか？
    Output:  主に機械学習とディープラーニングは一見違いがありませんが2つは大きく異なる点があります。機械学習は、サンプルの学習データを使用して生物学的な学習からコンピューターが自動でアイデアを生み出したり、自動で因果関係を抽出するデータサイエンスのフレームワークです。 予測モデルは一般的に、時間の経過とともに実行されます。
    ### 2 ###
    Prompt: Pythonはどのような言語ですか？
    Output:  RubyとTypeScriptは一部のコード自動生成でRSpecを使っています。
    ### 3 ###
    Prompt: 量子コンピュータの仕組みについて教えてください。
    Output:  主にどういった研究が行われているのでしょうか。
    ### 4 ###
    Prompt: AIによる画像認識の原理を説明してください。
    Output:  営業スキルがあればすぐ採用してもらえると思っていませんか?残念ながら、そんなことはありません。採用が決まると、最短即日で勤務し、育成を施すため、必要なスキルを備えている人がこれまで組織からは十把一絡げで選ばれてきました。営業スキル不要の「AIアシスタント」が誕生しようとしています。
    ### 5 ###
    Prompt: 自然言語処理の応用例をいくつか教えてください。
    Output:  主に機械翻訳サービスへの応用ですが、研究分野としては2つあります。1つは文章を生成するハイパーリンクのサンプルの作成と講義での実演です。もう一つはこれまで経験してきた内容をまとめることと学会発表なので、合わせて3点が今現在の取り組みになります。
    ### 6 ###
    Prompt: データサイエンスとは何ですか？
    Output:  主に機械学習と人工知能研究における用語で、人々は2つの方法で言葉を使っています: データサイエンスとデータサイエンティスト これは、データを使用して生物学、物理学からコンピューターサイエンスまでに及ぶ組織における有効な意思決定の決定論に関与する技術です。
    ### 7 ###
    Prompt: 最も有名なオープンソースの機械学習ライブラリは何ですか？
    Output:  RとPythonはこれらのケースではあまり役立ちません。理由は、これらの言語にあるすべてのコンパイラやサンプルのコードが、他のライブラリを必要としないためです。したがって、これらのコードが望みどおりの結果を生成していない場合は、サンプル・コードが不要であるために処理できないコードを処理しようとしないでください。
    ### 8 ###
    Prompt: クラウドコンピューティングとは何ですか？
    Output:  主にマイクロソフトが提供するクラウドコンピューティングサービス。インターネットを利用して、個々のコンピュータやサービスを必要なときに必要なだけ利用するコンピューティング環境。
    ### 9 ###
    Prompt: 強化学習の基本的な概念を教えてください。
    Output:  主に囚人のジレンマの強化学習で自動化を試みている。 なぜ完全合成か? 一つにはサンプルの数が少ないことにより、生物学的な学習からどういう法則を導き出して強化学習を行うか?が確立されておらず、サンプル数 が多い中で実際の学習をシミュレーションしたいためです。 当然至極ではあるのですが、実際に生物的な学習を大規模で継続するのは難しく、簡単に可能な強化学習をサンプ ル数によって追体験するのではなく、位置精度の向上を含めたある意味過酷な強化学習をシミュレーションするための手法であ り、これがゴールです。
    ### 10 ###
    Prompt: 教師あり学習と教師なし学習の違いは何ですか？
    Output:  RNN 探索とは、何の説明なしに使用できる2つ以上の入力レベルを使用して学習するハイパーパラメータのサンプリングです。 最近の講義では、学習器に対して追加のルール ジェネレータを使用して明示的なハイパーパラメータ ライセンスを適用する方法について解説しました。詳細は 教師なし学習 サンプルについてを参照してください。
    ### 11 ###
    Prompt: LLaMAのアーキテクチャの特徴は何ですか？
    Output:  RARA ®は、8ビット システムの自動並列2電源システムの急速充電モジュールです。 8mb413T は、DC 5 V 4.4 A / 2.5 A、出力電流と信号を同時にサポートします。
    ### 12 ###
    Prompt: トランスフォーマーモデルの基本的な仕組みは何ですか？
    Output:  主にトランスフォームトランスフォームと呼ぶケースでは、ハウジングの後ろ2つは、プラスチック部品、磁気シール、および4つの行で構成されます。プラスチックの開口部を指で押し出すと、これまで経験してきたように場所は自動でシフトします。ほとんどのメーカーは、シールの行を「密封」、またはエポキシすき間で接着しています。 3番目の行は、トランスフォームトランスフォームと呼ばれる場合、これは前腕と肩全体を包み込んで、かつてなかったほどの強さと精度の武器を与える余分な肉を挽きます。これには特に長すぎます。各トランスフォーマーモデルの中心には、金属のプライマーがあり、手で「密封」するように見えます。
    ### 13 ###
    Prompt: データの前処理において最も重要なステップは何ですか？
    Output:  RARA 探索とは、偶発性分類を可能にする2つ以上の因子レベルにある既知のハイパーパラメータ値を用いる回帰モデルで、目的地の設定を必要とせずに追加のルールを用いて組織等級クラスと被接種治療または治療必要年齢の開始年齢を予測できます。 RARA 探索は至適化を経ているので、速度違反手続きのような誤り修正を含む種々のディペンダブル手続きに対してもマルチパラメータ予測として応用できます。
    ### 14 ###
    Prompt: GPUの性能が機械学習にどのように影響するのか説明してください。
    Output:  RNN を用いて機械学習モデルを構築している自動運転車のドライバが、完全自動運転中に現れる多くの疑問を提起しました。 最近の講義では、学習から生成されたデータを用いて、実世界で自律的なロボットが子どもを追跡する方法を示します。 強力な学習アルゴリズムを使って、自律的なロボットは遠隔操作を経ることなくリモート地上ユニットを攻撃することもできます。 この講義では、自動運転車とRNNの複雑な関係をより詳しく説明します。
    ### 15 ###
    Prompt: ChatGPTのトレーニング方法について説明してください。
    Output:  RARA クレジットとは、8,860円を自動計算で割り、その2割を生成時に使用できるポイントです。
    ### 16 ###
    Prompt: ディープラーニングモデルの過学習を防ぐ方法は何ですか？
    Output:  RNN 探索とは、モデルを使用して分類を可能にするか、または視覚化する場合に学習中に使用できる最良の方法です。 これは、思考を使用して動的な学習アプローチを使用して追加のルール ジェネレータを可能にする技術です。 これにより、ニューラル ネットワーク 内の潜在的な学習率を増やすことが可能です。
    ### 17 ###
    Prompt: バッチサイズとは何ですか？
    Output:  RARA ® バッテリーケースでは、ハウジング素材の主要な材料のひとつはポリカーボネートである413Tです。 最近のバッテリーの例は振動からリチウムイオン電池を守る保護機能が備わっていますが、当社はほとんどのメーカーと競合しています。 強力なバッテリー密度は一般的にバッテリーケースの下に配置されているため、現代の主流バッテリーはリチウムイオン電池を含む必要があります。 バッテリーのケースは高出力であり、より大きなスペースを占有するのではなく、位置を占有すべきです。 同時に、高品質はより耐腐食性であり、およびさまざまな部品との容易な互換性、より長期的に続いてきれいに、容易かつ長期的に行われるべきです。
    ### 18 ###
    Prompt: グラディエントディセントとは何ですか？
    Output:  主に頭蓋の動きが研究に使用 されます。人々は2つ以上の椎骨を矯正し、動かすための4つの体位をとるためにさまざまなリハビリテーションを行います。これはどういう意味ですか?はい、これは単に被接種治療なので、何のためにも使用 されません。実際の効果としては、患者が健康に、いつまでも至福であると身体が認めているときです。多くの人々がメンタルヘルスに問題を抱えているため、頭蓋機能にも損傷が見られます。
    ### 19 ###
    Prompt: 多層パーセプトロンの基本的な構造はどのようになっていますか？
    Output:  RARA 探索とはどのような過程で使われているのですか?
    ### 20 ###
    Prompt: 自然言語処理でよく使われる手法を教えてください。
    Output:  主に機械学習を用いて自然言語処理をやっております。
    ### 21 ###
    Prompt: PyTorchとTensorFlowの違いは何ですか？
    Output:  RとPythonはこれらのスクリプト言語よりも自動化レベルが高いです なぜ?
    ### 22 ###
    Prompt: 分散学習とは何ですか？
    Output:  RARA 分散学習は、分散している自動化ジョブと個々の教師によって生成されたジョブの合計値を用いることにより、学習のスループットと精度を高めることです。これまで組織における有効な分散学習のモデルについては明らかになっていませんでしたが、今回のリリースにより、実際の学習率、クラス サンプル、機械学習戦略を競合他社の速度と比較できるようになったため、分散学習を用いる組織は、簡単に成果を決定付けることができるようになります。
    ### 23 ###
    Prompt: 事前学習済みモデルを使用する利点は何ですか？
    Output:  RARA 事前学習済みモデルを使用して、自動化は2つの方法で実行できます。1つ目の方法は、登録済みモデルを使用した、正確なモデルの動的な学習です。2つ目の方法は、カスタムモデルと登録済みモデルが連携するデータ ソースを備えたエンティティの学習です。
    ### 24 ###
    Prompt: モデルの推論を最適化する方法を教えてください。
    Output:  RARA 探索とは、モデル 分類に使用できる2つ以上の正規表現を使用して、モデル クラスからサンプルのカテゴリを選択する技術です。たとえば、「小児皮膚」という組織名を持つクラスと、「中学生」または「若年年齢」というクラスも参照することにより、「小児皮膚」というカテゴリのカルマンフィルタを作成する場合があります。
</details>