ローカルLLMのモデル切り替え時間を10秒から2秒未満に短縮

単一のGPUで複数の大規模言語モデルを実行する場合、メモリへのロードとアンロード（スワップ）が必要になることが多く、従来はこのプロセスに大きな遅延が生じていました。llama-sleepという新しいオープンソースツールと特定のLinuxカーネル設定を組み合わせることで、キャッシュされたプロンプトを保持したまま、このスワップ時間を10秒超から2秒未満にまで短縮することに成功しました。

何が起きたか

ローカルAIモデルをセルフホストする開発者は、すべての必要なモデルを同時にGPUメモリに収容できないハードウェア環境において、ボトルネックに直面することが頻繁にあります。llama-swapのようなプロキシを使用する典型的なセットアップでは、GPUメモリを満杯にする2つのモデル間を切り替える際、現在のプロセスを停止し、そのメモリを解放してから、新しいモデルを一から読み込む必要があります。NVIDIA RTX 3060で行ったベースラインテストでは、このコールドスワップには約10.3秒かかりました。この遅延は応答性に悪影響を与えるだけでなく、キー・バリュー（KV）キャッシュを破壊するため、システムはスイッチングのたびに長いシステムプロンプトを再計算しなければなりません。

この非効率性を解決するために、Robert Zaage氏はllama.cppサーバー用の軽量ラッパーであるllama-sleepを開発しました。アイドル期間中にモデルプロセスを終了させる代わりに、llama-sleepはNVIDIAのCUDAチェックポイント機能を使用してモデルの状態をサスペンド（一時停止）します。GPUメモリの内容をホストRAMにコピーし、事実上モデルを「駐車」させます。モデルが必要になった際には、状態をRAMから直接GPUへ復元します。このアプローチにより、ディスクからのモデル重みの再読み込みが不要になり、KVキャッシュも保持されるため、システムはほぼ即座に推論を再開できます。

この駐車メカニズムの初期実装では、スワップ時間は2.1秒に短縮されました。しかし、さらなる最適化により、データ転送速度がオペレーティングシステムのメモリページ処理方法によって制限されていることが判明しました。Linuxカーネルでマルチサイズ透明巨大ページ（mTHP）を有効にすることで、開発者はスワップ時間をさらに1.7秒まで短縮しました。これは元のベースラインに対して6倍の改善であり、音声アシスタントや自動化ワークフローなどのリアルタイムアプリケーションでの高速なモデル切り替えを可能にするものです。

仕組み

llama-sleepのコア技術は、NVIDIAドライババージョン550で導入されたcuda-checkpoint機能に基づいています。この機能により、実行中のCUDAプロセスのプロセス終了を行うことなく、デバイスメモリのスナップショットを取得してホストメモリへ移動させることができます。llama-sleepはリバースプロキシおよびプロセスマネージャーとして動作します。llama-serverをデタッチドデーモンとして実行し、外部からのリクエストを直接受け取らないように保護します。モデルが設定された期間（例えば5秒など）アイドル状態になると、llama-sleepはチェックポイント機能をトリガーします。これによりCUDAコンテキストがサスペンドされ、VRAMデータがシステムRAMにコピーされ、他のタスクのためにGPUリソースが解放されます。

この設計における重要な課題は、モデルが駐車中または復元中の間に到着するリクエストをどう扱うかです。GPUメモリのバックコピー中にリクエストが到着すると、プロセスがクラッシュしたりエラーを返したりする可能性があります。llama-sleepは、プロキシからのヘルスチェックを含むすべての受信リクエストをキューに保持することでこれを解決します。モデルが完全に起動し応答可能になるまで、これらのリクエストはリリースされません。これにより、まだ転送途中のモデルに対してプロキシが「ready」ステータスを認識することを防ぎ、セグメンテーションフォールトを回避し、シームレスな動作を保証します。

最終的なパフォーマンス向上は、メモリコピー自体の最適化によってもたらされます。GPUとホストRAMの間でのギガバイト単位のデータ転送は、ページフォールトの影響を強く受けます。標準的な4KBメモリページでは、コピー中にオペレーティングシステムが数百万個の個別ページエントリを管理する必要があり、大きなオーバーヘッドが発生します。Linuxカーネルバージョン6.8以降でマルチサイズ透明巨大ページを有効にすることで、システムは16KBから1MBまでのより大きなメモリブロックを使用できるようになります。これによりページフォールトの数が劇的に減少し、CPUによるメモリ管理がボトルネックとなるのではなく、PCIeバスの理論上の帯域幅限界に近いデータ転送が可能になります。

主要な詳細

• パフォーマンス向上: パラメータ数26Bの2つのモデル間のスワップ時間が、10.3秒から1.7秒に短縮されました。
• キャッシュの保持: スワップ中もKVキャッシュはそのまま維持されるため、数千トークンに及ぶプロンプトの再計算が不要になります。
• ハードウェア要件: ドライババージョン570以降のNVIDIA GPUとLinuxホストシステムが必要です。
• ソフトウェア依存関係: llama-swap v263、llama.cpp b11552、およびGo言語で書かれたスタンドアロンのllama-sleepバイナリを使用します。
• カーネル設定: 最適なパフォーマンスを得るには、Linuxカーネルで匿名メモリに対するマルチサイズ透明巨大ページ（mTHP）を有効にする必要があります。
• メモリのトレードオフ: 駐車されたモデルは、GPUメモリ使用量と同量のホストRAMを消費するため、十分なシステムメモリが必要です。

なぜ重要なのか

ローカルAIインフラストラクチャを構築するエンジニアにとって、この進展はコンシューマー向けハードウェアでのマルチモデルアーキテクチャの実現可能性を一変させます。以前は、モデルのスワップに伴う高い遅延のため、音声コマンド用には小さなモデル、複雑な推論用にはより大きなモデルといった、異なるタスク用に特化されたモデルを使用することは現実的ではありませんでした。ユーザーは、単純なタスクでリソースを無駄遣いしながら単一の大型モデルをロードし続けるか、コンテキスト切り替え時の長い遅延を受け入れるかのどちらかを選択せざるを得ませんでした。スワップ時間が2秒未満であれば、システムは知覚できるラグなしに需要に応じてリソースを動的に割り当てることができます。

KVキャッシュの保持は、ユーザー体験と計算効率の両方において同様に重要です。長いシステムプロンプトや会話履歴の再計算には時間とエネルギーの両方が消費されます。この状態をホストRAMに保持することで、llama-sleepはスワップ後の最初のトークン生成がほぼ瞬時に行われることを保証します。これは、コンテキストの連続性が不可欠なホームオートメーションやインタラクティブエージェントなどのアプリケーションにおいて特に価値があります。モデルをバックグラウンドで駐車できる能力により、システムは常に応答可能な状態を保ち、クラウドAPIの柔軟性とローカル推論のプライバシーとの間のギャップを埋めます。

実施できること

• ドライバの更新: CUDAチェックポイント機能をサポートするために、NVIDIA GPUドライバをバージョン570以降に更新してください。

• llama-sleepのインストール: llama-sleepのスタンドアロンGoバイナリをダウンロードし、既存のllama-swap設定に統合してください。

• アイドルタイムアウトの設定: モデルがアクティブに使用中でない場合のみ駐車されるよう、llama-sleepの--idleパラメータをあなたの使用パターンに合わせて設定してください。

• 巨大ページの有効化: Linuxカーネル6.8以降を実行している場合、16KBから1MBのサイズに対応する関連sysfsエントリに"always"を書き込むことで、マルチサイズ透明巨大ページを有効にしてください。

• メモリ使用量の監視: ホストシステムに、駐車されたモデルと他の実行中のプロセスを保持するのに十分な空きRAMがあることを確認してください。

• スワップ遅延のテスト: あなたの特定のハードウェア環境でのパフォーマンス向上を検証するために、スワップ後のリクエストから最初のトークン出力までの時間を測定してください。