vLLM Sleep Modeよるモデルのゼロリロード切り替え機能の検証

こんにちは。NTTドコモビジネスの露崎です。本ブログではvLLMの本家コミュニティのブログで紹介されたvLLMのモデルのゼロリロード切り替え機能の概要に加えて本機能をContainerベースで検証した結果について紹介します。

はじめに

こんにちは。NTTドコモビジネスの露崎です。本ブログではvLLMの本家コミュニティのブログで紹介されたvLLMのモデルのゼロリロード切り替え機能について紹介します。

vLLMはカリフォルニア大学バークレー校のSky Computing Labが開発しOSSで公開している推論高速化ソフトウェアです。vLLMでは量子化やバッチ処理を始めとしたさまざまな高速化テクニックを実装しており、Large Language Model(LLM)の推論を高速化できます。また、vLLMはモデルリポジトリであるhuggingfaceで公開されているさまざまなモデルに対応しており、Meta社のLlamaを始めとするオープンウェイトモデルを高速に実行する基盤ソフトウェアとして世界的に利用されています。

モデルのゼロリロード切り替え機能が取り組む課題

本ブログではvLLMを用いたLLM推論の提供において、同一のGPU上で複数のモデルを利用したい場合のモデル切り替えのオーバヘッドを解決する機能を紹介します。

vLLMはLLMの推論を高速に実行し、ユーザに結果を素早く提供できる有用性の高いソフトウェアです。 vLLMは高速な推論を実現するためにLLMの計算用キャッシュ(a.k.a. KV Cache) を事前にできるだけ確保する(デフォルトでは利用可能なメモリの90%)という仕組みを持っています。 この動作は単一のモデルを提供する際には特に問題となることはありません。しかし、GPUと言う貴重な計算リソースをある1つのモデルのために占有してしまうと言うことでもあります。

この特性のため、以下のように複数モデルを使い分けるユースケースにおいては、利用するモデル毎にGPUリソースを割り当てる必要があり、コストを増大させることにつながります。

  • 質問からテキストベースで回答するLLMと画像の説明や分類を説明するVisual Launguage Model (VLM)を使い分けるシステム
  • 複数の異なるLLMの出力結果を組み合わせて最終的な回答を得るエージェント

この対策として、vLLMが提供している公式Containerイメージを利用し、需要やタイミングに応じて使うモデルをContainerの起動、終了によって切り替えると言った運用は可能です。 しかし、ファイルシステムからのモデルのロードや、起動時に実行されるCuda Graph構築による最適化のオーバヘッドのため、例えばgemma-3-4b-itのような比較的小さいモデルであっても1分程度の起動待ちが発生すると言う課題があります。

vLLM Sleep Mode

前述の課題を解決する機能の1つがvLLMのSleep Modeです。Sleep Modeは起動中のvLLMに対して /sleep APIを呼び出すことでロード済みのモデルをCPUのメモリに待避する機能です。待避されたメモリは /wake_up APIを呼び出すことでGPUメモリへリロードできます。 /wake_up ではCPUからGPUへメモリ転送を実施するためファイルシステムからモデルを読み込むより高速にロードが実行でき、CUDA Graphの構築のオーバヘッドを回避できるため、vLLMのプロセスを最初から起動するよりも迅速にAPIが利用できるようになります。

NOTE
Sleep Modeを使うためには環境変数 VLLM_SERVER_DEV_MODE=1 を設定する必要があります。

Sleep Level

/sleep APIには引数でスリープ時のレベルを設定できます。(例: /sleep?level=1)

現在は2段階のスリープレベルがあり、それぞれ以下の通りです。

  • Level 1: モデルの重みをCPU RAMへオフロードし、KV CacheをGPUメモリから廃棄する
  • Level 2: モデルの重みとKV CacheをGPUメモリから破棄し、スリープ時のCPU RAMの使用量を最小限にする

公式ドキュメントではLevel 1は同一のモデルを再度利用する際に有効であり、Level 2は異なるモデルを利用する場合に有効である、と解説しています。

動作確認

vLLMの本家コミュニティのブログ公式ドキュメントではPython API、及び、マルチプロセスでの実例が紹介されていますが、実際の運用を考慮するとContainerベースでもこの機能が実行できるかを確認したかったため、現在のStableの最新版であるv0.11.0で動作確認を実施した結果を紹介します。

環境構成

  • ハードウェア

    • NVIDIA GeForce RTX 4090 24GB
  • ソフトウェア

    • Ubuntu 22.04.5 LTS
    • NVIDIA Driver 580.95.05
    • docker-ce 5:28.5.1-1
    • vLLM v0.11.0

利用モデル

実験1. API利用状況の確認

vLLMの起動

まずはSleep Modeが現在のStableであるv0.11.0で利用できるのかを確認します。 以下のコマンドを実行してモデルを起動します。リポジトリへのアクセス権等の設定については必要に応じてHuggingFaceのアカウントで作成したトークンを環境変数 HF_TOKEN へ設定してください。

MODEL="google/gemma-3-4b-it"

sudo docker run --runtime nvidia --gpus all \
    -v ~/.cache/huggingface:/root/.cache/huggingface \
    -p 8000:8000 \
    --ipc=host \
    --env "HUGGING_FACE_HUB_TOKEN=$HF_TOKEN" \
    --env "VLLM_SERVER_DEV_MODE=1" \
    vllm/vllm-openai:v0.11.0 \
    --model ${MODEL} --enable-sleep-mode

このコマンドでvLLM v0.11.0を起動できますが、起動時に以下がログ出力されます。

WARNING 10-29 18:58:56 [api_server.py:966] SECURITY WARNING: Development endpoints are enabled! This should NOT be used in production!

この後の手順でAPIのサンプルを提示しますが、 /sleep APIは現在、通常のChat API同様のAPIエンドポイントによって受け付けられるため、素のままのAPIを利用すると任意のユーザがモデルの停止ができる状態になります。このため、プロダクションでの利用は現在推奨されていないようです。

Sleepによるモデルの停止

次にSleep ModeのAPIを呼び出します。Sleep ModeのAPIは以下の通りREST API経由で実行します。

curl -X POST 'localhost:8000/sleep?level=1'

これを実行すると以下のようなログが出力され、GPUメモリが解放されます。

(EngineCore_DP0 pid=165) INFO 10-29 19:36:16 [cumem.py:228] CuMemAllocator: sleep freed 20.25 GiB memory in total, of which 8.63 GiB is backed up in CPU and the rest 11.62 GiB is discarded directly.
(EngineCore_DP0 pid=165) INFO 10-29 19:36:16 [gpu_worker.py:117] Sleep mode freed 20.23 GiB memory, 1.21 GiB memory is still in use.
(EngineCore_DP0 pid=165) INFO 10-29 19:36:16 [executor_base.py:189] It took 1.695459 seconds to fall asleep.
(APIServer pid=1) INFO:     172.17.0.1:36138 - "POST /sleep?level=1 HTTP/1.1" 200 OK

今回利用した google/gemma-3-4b-it ではvLLM起動時に23GB程度占有していたメモリが、 /sleep 呼び出し後には1GB程度まで減っていることを確認しました。

NOTE
Sleep Mode中のモデルに対して推論リクエストを実施するとv0.11.0の時点では400 Errorとなり、vLLMプロセスがクラッシュして停止します。

WakeUpによるモデルのリロード

モデルのリロードを実施するには以下のコマンドを実行します。

curl -X POST 'localhost:8000/wake_up'

これを実行すると以下のようなログが出力されます。

(APIServer pid=1) INFO 10-29 19:43:03 [api_server.py:1016] wake up the engine with tags: None
(EngineCore_DP0 pid=165) INFO 10-29 19:43:03 [executor_base.py:205] It took 0.368304 seconds to wake up tags {'kv_cache', 'weights'}.
(APIServer pid=1) INFO:     172.17.0.1:58818 - "POST /wake_up HTTP/1.1" 200 OK

WakeUpが完了すると、通常のvLLMとしてOpenAI互換のAPIが利用できます。

NOTE
Level 2のSleep Modeでは /wake_up の後に /collective_rpc (methodとして reload_weights を指定) 、 /reset_prefix_cache のAPIを呼び出す必要があります。

実行結果とパフォーマンス考察

モデル実行時の手順とオーバヘッドは以下の通りでした。

level 1 level 2
Sleep後のGPU使用メモリ 1236MiB 1241MiB
復帰時の手順 1 (wake_up) 3 (wake_up, collective_rpc, reset_prefix_cache)
復帰時のコマンド実行時間 (sec) 0.368304 0.737

Level 1、Level 2のいずれにおいても1秒以内にモデルの復帰が完了するため、vLLMのContainerを再起動するより高速にAPIを使えるようになることがわかります。GPUメモリとしてはLevel 1とLevel 2では削減効果は同等でした。復帰の手順の複雑性を考慮するとLevel 1を標準として利用可能だと考えます。公式ブログではLevel 2の際にCPU側のRAMを節約できると記載があるため、GPUだけでなくSleep中のCPUメモリも削減したい場合はLevel 2を使うと良いでしょう。

実験2. 複数モデルでの切り替え

実験1でvLLM v0.11.0のContainerでSleep Modeを利用可能なことが確認できたため、複数のモデルの切り替えを実行してみます。 vLLMの公式ブログでは同じ環境内の別プロセスとして起動する方法を紹介していましたが、ここではより本番環境に近づけるため1枚のGPUに対して2つのvLLM Containerを起動した状態でモデルの切り替えが実行できるかを試しています。

実験の実行手順としては以下の通りです。

準備

  1. 1つめのモデルのvLLM Containerを起動(8000ポート)
  2. /sleep?level=1 APIで1つめのモデルをSleep
  3. 2つめのモデルのvLLM Containerを起動(8001ポート)
  4. /slee?level=1 APIで2つめのモデルをSleep

実験(各モデルに対して)

  1. /wake_up APIでモデルのリロード
  2. sample.py スクリプトを用いてhealthcheck及び、推論APIの呼び出し

実験用スクリプト

今回は以下のようなBashスクリプトで実験をしています。

#!/bin/bash
set -e

# 切り替えるモデルの定義
MODELS=("meta-llama/Llama-3.2-1B-Instruct 8000" "google/gemma-3-4b-it 8001")

RUN_NAME="vllm_sleep"

# コンテナの作成と起動用関数を定義
create_vllm(){
    c_name=${1#*/}
    echo "Use name ${c_name}"
    docker run -d --name ${RUN_NAME}-${c_name} --rm --runtime nvidia --gpus all \
        -v ~/.cache/huggingface:/root/.cache/huggingface \
        -p $2:8000 \
        --ipc=host \
        --env "VLLM_SERVER_DEV_MODE=1" \
        --env "HUGGING_FACE_HUB_TOKEN=${HF_TOKEN}" \
        vllm/vllm-openai:v0.11.0 \
        --model $1 --enable-sleep-mode
}


# 各モデルを起動し、APIの疎通が確認できたらSleepする
for model in "${MODELS[@]}"; do
    read model_id port <<< "$model"
    c_name="${RUN_NAME}-${model_id#*/}"
    echo "start vllm container"
    create_vllm ${model_id} ${port}
    echo "wait and call an inference"
    uv run sample.py ${port}
    echo "suspend model with level 1"
    curl -X POST "http://localhost:${port}/sleep?level=1"
done

# 各モデルに対してリロードしAPIが実行できれば再度Sleepする
echo "--- start benchmark for model switching ---"
for model in "${MODELS[@]}"; do
    read model_id port <<< "$model"
    c_name=${RUN_NAME}-${model_id#*/}
    echo "Use name ${c_name}"
    curl -X POST "http://localhost:${port}/wake_up"
    echo "wait and call an inference"
    uv run sample.py ${port}
    echo "suspend model with level 1"
    curl -X POST "http://localhost:${port}/sleep?level=1"
done

# 起動しているContainerをCleanupする
docker ps -q | xargs docker kill

Containerの起動状態を確認し、推論APIを呼び出すPythonスクリプト( sample.py )は以下の通りです。

#!/usr/bin/env python3

from openai import OpenAI
import os
import requests
import time
import sys

HOST = "http://localhost"


def wait_to_ready(URL):
    while True:
        try:
            requests.get(f"{URL}/health")
            print("Healthcheck: OK")
            break
        except requests.exceptions.ConnectionError:
            time.sleep(0.1)


def main():
    port = sys.argv[1] if len(sys.argv) > 1 else 8000
    url = f"{HOST}:{port}"
    wait_to_ready(url)
    # Localhost上のvLLMでありOpenAIの実際のキーは不要のためdummyを設定 
    os.environ["OPENAI_API_KEY"] = "dummy"
    client = OpenAI(base_url=f"{url}/v1")
    model_id = client.models.list().data[0].id
    print(f"Model {model_id} found.")
    response = client.responses.create(
        model=model_id,
        input="Write a one-sentence bedtime story about a unicorn."
    )

    print(f"Got response: {response.output_text}")


if __name__ == "__main__":
    main()

実行結果

実験スクリプトを実行すると、 --- start benchmark for model switching --- 以下の出力がおよそ1秒程度で完了することが確認できると思います。このことからvLLMのSleep Modeを利用することで複数のContainer、モデルに対して1つのGPUリソースで効率的なモデルの切り替えを実施できることが確認できました。

NOTE
例示したスクリプトでは meta-llama/Llama-3.2-1B-Instructgoogle/gemma-3-4b-it の順序で起動し、問題なく /sleep , /wake_up が動作することを確認できました。しかし、別の実験として meta-llama/Llama-3.2-3B-Instructgoogle/gemma-3-4b-it で試した場合には、Sleep後に meta-llama/Llama-3.2-3B-Instruct/wake_up させるタイミングでOut of Memory (OOM) Errorが起き、プロセスが停止しました。 meta-llama/Llama-3.2-3B-Instructmeta-llama/Llama-3.2-1B-Instruct より大きなモデルで、必要なGPUメモリが大きいため、Sleep Modeで切り替える場合であっても最低限必要なメモリの量は事前に調査しておく必要があると考えられます。

まとめ

本ブログではモデル切り替えのオーバーヘッドを削減するvLLMのSleep Modeについて紹介しました。Sleep Modeはまだ開発段階の機能ですが、複数のモデルの処理に対してGPUを効率的に利用できる機能ということが確認できました。今後、LLMの需要拡大に向けてこうしたGPUの処理効率を向上させるソフトウェア、仕組みを活用することはますます重要になると考えられます。