Gemma推論の新構成：DataflowとvLLMの連携で効率が23倍超に

Googleは、オープンソースの高速推論エンジンvLLMをGoogle Cloud Dataflowに統合し、オープンモデルGemmaのストリーミング推論を極めて高効率に実行する構成を公開した。わずか数行のコードを追加するだけで、スループット重視の大規模推論パイプラインを構築できる。

LLMを本番のデータパイプラインに組み込む際、最大の障壁となるのが従来のバッチ処理の非効率さだ。GPUは複数入力をまとめて並列処理するのが得意だが、LLMはトークンを1つずつ順番に生成する。短い回答で済むプロンプトと長文を要するプロンプトを同時に処理すると、長い処理が終わるまでGPU全体が占有され、リソースの無駄が生じてしまう。

継続的バッチと専用モデルマネージャー

vLLMはこの課題を「継続的バッチ処理」で解決する。トークンを生成するループごとにバッチを動的に再計算し、完了した処理は即座に返し、新規リクエストをその場でバッチへ割り込ませる仕組みだ。

一方、これを分散処理フレームワークであるDataflowに持ち込むには運用の壁があった。Dataflowは通常CPUコアごとに独立したワーカープロセスを立ち上げるため、全プロセスでGemmaを読み込むとメモリが枯渇してしまう。

そこでDataflowのモデルマネージャー機能が役立つ。パイプライン内にvLLM専用の推論プロセスを1つだけ立ち上げ、各ワーカープロセスが非同期でリクエストをキューへ送る構成を自動で整えてくれる。複雑なマルチプロセス制御を書く必要はなく、RunInference変換を使えば主要なパイプライン処理はわずか5行ほどで記述できる。

23倍超の効率化と検証結果

Googleが公開したベンチマークでは、T4 GPUを1基搭載したワーカーマシンを使用し、P3データセットのプロンプト1万件をgoogle/gemma-2bに通して性能を比較した。

固定サイズバッチを用いる標準パイプラインでは処理に59.137 vCPU時間を要したのに対し、vLLMの継続的バッチを適用したパイプラインはわずか2.481 vCPU時間で完了した。計算リソースの消費量にして23倍以上の改善を達成している。

双方とも手動チューニングを行っていない比較のため、標準パイプライン側もバッチサイズを最適化すれば差は縮まる余地がある。とはいえ、チューニングなしで即座にこれだけのスループットを出せる点こそが実用的だ。別モデルへの変更もハンドラーの文字列を書き換えるだけで対応できる。現在、詳細なコードとノートブックがGoogle Cloudの公式ドキュメントで公開されている。

なぜ重要か

DataflowやApache Beamを使ってテキスト分類、要約、翻訳などのバッチ処理やストリーミングETLを行っている開発者なら、すぐに試す価値がある。インフラのプロセス間通信やGPUメモリ管理に悩まされることなく、使い慣れたパイプラインに最先端の推論最適化を組み込めるからだ。

独自APIの利用を促すだけでなく、定評あるオープンソース推論スタックをマネージド基盤に素直に取り入れたアプローチは合理的だと評価できる。ただし、この構成は総スループットの最大化に特化しているため、1リクエストあたりの応答速度を極限まで縮めたいチャット用途などには向かない点には留意が必要だ。