Cloud TPU と vLLM で LLM 推論を試そう — リソース確保からベンチマークまで
Takashi Sato
AI Infrastructure Specialist, Google Cloud Japan
こんにちは、Google Cloud でインフラ領域を担当している佐藤です。
今回は、Cloud TPU v6e 上で vLLM を使い、大規模言語モデル Qwen3-32B の推論環境を構築する手順をハンズオン形式でお届けします。DWS Flex Start によるリソース確保から、パラメータ チューニング、INT8 量子化、ベンチマークまで一通りカバーしています。
最近、大規模言語モデル（LLM）の推論環境への需要が急速に高まっています。しかし、ハイパフォーマンスなインフラをオンデマンドで調達しようとすると、特定の GPU/TPU Type やリージョンによっては「オンデマンド リソースの即時確保が困難」という課題に直面した—そんな経験はありませんか？
本記事では、Google Cloud の Dynamic Workload Scheduler (DWS) Flex Start モードを活用し、待機キュー経由でリソースの確保を効率的に行いつつ、オープンソースの高スループット LLM 推論エンジンである vLLM を用いて、大規模モデルである Qwen3-32B の推論環境を構築する実践的なハンズオン手順をご紹介します。このガイドを通じて、リソース確保からデプロイ、そしてパフォーマンスのベンチマークまでの流れを解説していきましょう。
TPU 確保方法の比較と DWS Flex Start モードの長所
Google Cloud の GPU/TPU リソースを確保する方法にはいくつかの種類があります。それぞれの長所と短所を理解することで、ワークロードに最適な選択が可能になります。
では、なぜ Flex Start を使うのでしょうか。
オンデマンドや Spot では、「今すぐ使いたいがリソースがない」というエラーが返されることがありますが、DWS Flex Start を使用することで「キュー」にリソース要求が登録されます。バックグラウンドでリソースの空き状況が監視され、確保可能になった瞬間にプロビジョニングが行われるため、張り付いてリソース作成を連打する必要がなく、リソースプロビジョニングの成功率をぐっと上げられるのが大きな長所です。
vLLM と vLLM-TPU の違いとは？
ハンズオンに入る前に、今回利用する vllm-tpu イメージについて補足します。
-
vLLM（通常版）： 主に NVIDIA GPU (CUDA) や AMD GPU (ROCm) 向けに高度に最適化された LLM 推論エンジンです。PagedAttention というメモリ管理技術により、KV キャッシュの断片化を防ぎ、高いスループットを実現します。
-
vLLM-TPU： Google のカスタムシリコンである TPU アーキテクチャ上で、vLLM の PagedAttention や最適化技術を動作させるために特化した拡張実装・環境です。内部的には 2 つのモデルレジストリを確認した上でモデルコードを取得・実行します（下図参照）。そして Torchax によって PyTorch モデルコードが JAX として扱われ、TPU 上での最適な推論実行が可能になります。本ハンズオンでは vllm/vllm-tpu の Docker イメージを利用することで、複雑な依存関係やコンパイラの設定を意識することなく、すぐに TPU のパワーを推論に活用できます。
ハンズオンガイド
1. TPU リソースのリクエスト (Flex Start)
まず、通常の gcloud compute tpus tpu-vm create ではなく、queued-resources create コマンドを使用します。このコマンドによって、1 章で触れた Flex Start によるリソース調達を自動化することができます。
コマンドの意味について
--provisioning-model flex-start を指定することで、キューにリソース要求がエンキューされます。--max-run-duration 24h によって、利用開始から 24 時間後に自動的に終了するように設定しています。消し忘れによる課金防止にも役立ちますね。Spot VM で起動したい場合は、コマンドから alpha を外して --spot を指定することで起動できます。
キューに格納したいリソースを上記コマンドで作成した後はステータスに則って調達されます。ステータスの確認には以下のコマンドを実行してください。
STATE が qr-request-spot のように ACTIVE になれば調達完了です。一方で takashix-qr-request のように WAITING_FOR_RESOURCES の場合、まだ調達できておらずリソース確保を待っている状態です。
※調達できない （STATE が Active にならない）場合は代わりに以下のコマンドを実行してください。
2. インスタンスへの接続と Docker 環境設定
リソースが「ACTIVE」になったら、SSH 接続して環境を準備しましょう。本ガイドでは Qwen3-32B を利用することを想定し、Hugging Face の Token を設定します。
SSH 接続
Docker Image の設定
Hugging Face Token の設定（<your HF token> はご自身の Token に置き換えてください）
3. vLLM サーバーの起動とパラメータ・チューニング
Docker コンテナを --privileged および --net=host で起動します。これは TPU デバイスへの直接アクセスと、ホストの高速なネットワークをコンテナに許可するためです。また --shm-size 100gb を指定して、モデルの重みや共有メモリ領域が不足しないようにしています。
以下のようにイメージのダウンロードが完了したら vLLM サーバーの起動完了です。
上記コマンドにより root@ から始まるプロンプトになっていれば、コンテナ内へのアクセスが成功したことになります。いよいよ vllm serve コマンドでモデルをデプロイしましょう。ここでのパラメータ設定が、推論のパフォーマンス指標であるスループットやレイテンシに大きく影響を与えます。
4. [オプション] INT8 量子化 (W8A8) を用いた Serving
TPU v6e の性能をさらに引き出し、巨大なモデルのメモリ使用量を削減するために、INT8（W8A8）量子化を有効化してモデルをサーブすることが可能です。
vLLM-TPU では内部的に Qwix と呼ばれる JAX 向け量子化ライブラリを使用します。量子化のためのコンフィグファイルはコンテナ内にすでに存在していることがほとんどですが、なかった場合は以下の手順で YAML 形式の設定ファイルを作成し、そのファイルを --additional-config オプションを用いてサーバー起動時に読み込ませてください。
量子化設定ファイル (int8_default.yaml) の作成
INT8 有効化による vLLM サーバーの起動
ファイルの意味と量子化の仕組みについて
-
設定ファイル (int8_default.yaml) の役割: Qwix に対する量子化ルールの定義ファイルです。module_path: '.*' によってモデル内のすべてのレイヤーを対象とし、重み (weight_qtype) と活性化関数 (act_qtype) の双方を int8 フォーマットとして扱うよう指示します。これを W8A8 (Weight 8-bit, Activation 8-bit) 量子化と呼びます。
-
量子化の方法: vLLMがモデルをロードして XLA コンパイルする際、このルールに基づき動的に計算グラフが書き換えられます。ロードされる FP16/BF16 の重みは TPU メモリである HBM 上で INT8 に圧縮・変換され、推論時にも INT8 の行列積として実行されます。これにより、メモリ帯域のボトルネックが緩和されると同時に、TPU v6e に搭載された強力な INT8 演算器の性能が引き出され、スループットの向上とレイテンシの削減が期待できるでしょう。
-
利用するモデル: Qwen3-32B-GPTQ-Int8 などの FP8 や INT8 ですでに Weight が保存されたモデルはフォーマットが vllm-tpu でサポートされていない可能性があるため、基本的には BF16 の重みをロードできるように HF 上の Qwen3-32B などのデフォルトモデルを利用して Post-Quantization を行うことを推奨します。
上記の vllm serve コマンドでモデルの serve に成功すると Application startup complete. というメッセージが表示されます。
5. 動作確認とベンチマーク
Step 5.1. API エンドポイントへのテストリクエスト
別のコンソールを開いて VM およびコンテナへ接続し、OpenAI 互換の API エンドポイント経由で推論テストを行いましょう。新しいコンソールの方で環境変数を設定していない場合は、以下のように再度設定を行ってから SSH コマンドを実施してください。
VM への SSH 接続
コンテナへの接続
推論テストの実施
以下のようにレスポンスが返ってくれば成功です。
Step 5.2. ベンチマークテストの実施
推論サーバーが正しく稼働していることが確認できたら、vLLM に同梱されている公式ベンチマークスクリプトを用いて、本番環境を模した負荷テストを実施しましょう。
ベンチマークパラメータの意味と影響：
ここでは、1000 個のリクエスト（--num-prompts 1000）を並行してサーバーに投げ込んでいます。--random-input-len と --random-output-len を変えることで、たとえば入力を長くすれば RAG のような prefill 負荷が高いケースを、入出力を同程度にすれば翻訳や対話のような decode 負荷が高いケースを再現できます。
入力トークン長（今回は 1800）を長く設定すると、モデルの Prefill（初回計算）フェーズの負荷が高まり、TTFT（Time To First Token: 最初のトークンが出力されるまでの時間）が増大する傾向があります。逆に、このベンチマーク環境で出力スループットを示す Output token throughput (tok/s) が大きく表示されていれば、TPU の並列計算能力をしっかり引き出せている証拠となります。
ベンチマーク結果の例：
結果 1 - 本ガイド記載のパラメータ通りでのベンチマーク (INT8 量子化なし)
結果 2 - 本ガイド記載のパラメータ通りでのベンチマーク (INT8 量子化あり)
結果 3 - 非同期スケジューリング有効でのベンチマーク (INT8 量子化あり)
利用コマンド
非同期スケジューリングの有効化 (--async-scheduling)
CPU 側のリクエストスケジューリングと、TPU 側のモデル実行を非同期で行うことで、ホストとデバイス間の待機時間をなくし、スループットを数 % 〜 10 % 程度押し上げる効果が確認されています。
結果の比較
以下の表にそれぞれの条件でのスループットを比較したところ、INT8 での量子化を行い非同期スケジューリングの有効化も併用すると 123 % のパフォーマンス向上がみられました。
量子化に伴うモデル精度の変化には注意を払う必要がありますが、チューニングを行う際にはぜひ今回利用したオプションもご参照ください。
6. クリーンアップ：リソースの削除
検証が完了したら、余分なコストや Quota の消費を防ぐため、リソースの削除を行います。通常の VM と異なり、キューに格納されたリソースは「SUSPENDED」などの状態に関係なく Quota の割り当てを消費し続けます。今後の別の要求がブロックされるのを防ぐため、明示的に削除コマンドを実行しましょう。
Warning: ハンズオン終了後は、上記の queued-resources delete コマンドを実行してください。tpu-vm delete だけではキューのエントリが残り、Quota 消費の原因となる可能性があります。
まとめ
本記事では、TPU v6e 上で DWS Flex Start を活用した効率的なリソース調達と、vLLM を用いた推論環境の構築・評価までの一連の流れをご紹介しました。
他の Model などの Recipe は以下のリポジトリに公開されていますのでご参照ください。
https://github.com/AI-Hypercomputer/tpu-recipes/tree/main/inference/trillium/vLLM
Flex Start の活用によるリソース枯渇状態からのプロビジョニング成功率向上、vllm-tpu コンテナを用いた容易な最適化環境の構築、そして gpu-memory-utilization などのパラメータチューニングによるスループットとレイテンシのバランス調整は、本番環境における大規模モデル運用において非常に重要なノウハウとなります。
ぜひみなさんの LLM ワークロードでも試してみてください。
