Microsoftの研究者は2024年1月4日、公式ブログ「Splitwise improves GPU usage by splitting LLM inference phases」において、大規模言語モデル(LLM)を使用した推論においてGPUの使用効率を向上させる「Splitwise」と名付けられたシステムについて解説した。

Splitwiseでは、LLM推論を2つのフェーズした上で別のマシンに割り当て、それぞれの特性に応じて利用するハードウェアが最適になるように設計することで、ハードウェアのコストを削減しながらスループットを最大化することに成功したという。

LLM推論でハードウェアの利用効率を制限する要因

LLM推論には、プロンプトフェーズとトークン生成フェーズという、異なる特性を持つ2つのフェーズがある。プロンプトフェーズではすべてのユーザのプロンプトを並行して処理できるのに対して、トークン生成フェーズでは各出力トークンを順番に生成する必要がある。

そのため、プロンプトフェーズは計算負荷が高くなり、トークンフェーズはメモリ負荷が高くなるという特性の不一致が発生する。このことが、LLM推論におけるハードウェアの利用効率を制限する主な要因だと研究者は指摘する。

  • LLM推論プロセスにおけるプロンプトフェーズとトークン生成フェーズの例 出典:Microsoft Research Blog

    LLM推論プロセスにおけるプロンプトフェーズとトークン生成フェーズの例 出典:Microsoft Research Blog

Splitwiseの特徴

Splitwiseでは、プロンプトフェーズとトークン生成フェーズを異なるマシンに分割して配置し、利用可能なハードウェアが最適になるように設計することで、両方のフェーズでのハードウェアの使用効率を高めている。加えて、プロンプトフェーズとトークンフェーズにわたる混合バッチを処理するための混合プールを設けている。この混合プールは実行時の計算需要に応じて動的にサイズ設定される。プロンプトマシンからトークンマシンに対しては、低遅延のInfiniBand経由で状態コンテキストを転送する。

  • Splitwiseアーキテクチャの概要 出典:Microsoft Research Blog

    Splitwiseアーキテクチャの概要 出典:Microsoft Research Blog

Splitwiseの評価結果

Microsoftでは、Splitwiseを使用してNVIDIAのDGX-A100やDGX-H100などさまざまなGPUを備えたクラスタを設計して評価を行い、Splitwizeクラスタが通常のクラスタと比較して同じコストでスループットを最大化できることを確認したという。具体的には、20%低いコストで1.4倍高いスループットを達成できることや、同じコストと電力で2.35倍のスループットを達成できることが示されている。

Splitwiseに関する研究のより具体的な成果は、Microsoftが2023年11月に公開した以下の論文にまとめられている。