🤖 AI

Fine-Tuning and Serving Gemma 4 31B on Google Cloud TPU: A Technical Comparison with GPU Baselines

本論文は、Google Cloud TPU 上で Gemma 4 31B のファインチューニングとサービングを初めてエンドツーエンドで実証し、GPU ネイティブなワークフローを JAX スタックに移行するために必要なコード適応を詳述するとともに、H100 GPU ベースラインと比較して推論スループットは同等でありながらレイテンシが大幅に低減され、トレーニング速度が 1.61 倍高速でコストが 1.82 倍低く抑えられることを TPU 構成が達成していることを実証する。

原著者： Jatin Kishnani, Mayank Goel, Amit Singh, Pulkit Agrawal, Sairanjan Mishra

公開日 2026-05-26

📖 1 分で読めます☕ さくっと読める

CC BY 4.0

原著者： Jatin Kishnani, Mayank Goel, Amit Singh, Pulkit Agrawal, Sairanjan Mishra

原論文は CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) でライセンスされています。 ✨ これは以下の論文のAI生成解説です。著者が執筆または承認したものではありません。技術的な正確性については原論文を参照してください。免責事項の全文を読む

あなたが、非常に優秀だが非常に高価なロボット（AI モデル）に、電子回路のコンピュータコード（Verilog）を書くという新しいスキルを教えようとしていると想像してください。このトレーニングを行う「学校」として、あなたは以下の 2 つの選択肢を持っています。

GPU スクール: NVIDIA が運営し、有名な H100 チップを使用しています。これは、多くの経験豊富な教師がいる、標準的でよく歩かれた道です。
TPU スクール: Google が運営し、独自のカスタム TPU チップを使用しています。これはより新しく、高度に特化したキャンパスで、より速く安価ですが、入学するには全く新しい言語を学ぶ必要があります。

h2loop.aiからのこの報告書は、Google の TPU スクールを試すことを決めたエンジニアたちによって書かれた「フィールドガイド」です。彼らは 310 億パラメータの AI モデル（Gemma 4）をコード作成に訓練し、その性能を標準的な NVIDIA スクールと比較しました。

以下に、彼らの旅路を分かりやすく解説します。

1. 言語の壁（「ポート」の問題）

最大の障壁はハードウェアそのものではなく、ソフトウェアでした。

GPU スクールは、ほとんどの AI 開発者が既に知っている人気プログラミング言語であるPyTorchを話します。
TPU スクールは、異なる思考様式を必要とする別の言語であるJAXを話します。

著者たちは、PyTorch から JAX へと「授業計画」全体を翻訳する必要がありました。具体的には以下のことを行いました。

複数のチップにまたがるモデルの「脳」の分割方法を再編成しました（これは、本が 2 つの部屋ではなく 4 つの異なる部屋に散らばっている図書館を整理するようなものです）。
TPU スクールでは呼び方が異なるため、モデルの特定の部分を名前変更しました（例えば、「q_proj」を「q_einsum」と呼ぶなど）。
TPU スクールは世界標準（Safetensors）とは異なるファイルシステム（Orbax）を使用するため、作業を保存するためのカスタム「橋」を構築しました。

比喩: これは家を引っ越すようなものです。家具（AI モデル）は同じですが、新しい家（TPU）には異なるドアの幅や間取りがあります。家具を分解し、新しいドアを通って運び込み、再組み立てしないと、入りません。

2. トレーニング競争（ロボットに教える）

モデルのセットアップが完了すると、トレーニング競争が始まりました。

速度: TPU スクールは1.6 倍速く、トレーニングコースを 3.3 時間で完了しました。一方、GPU スクールは 5.4 時間かかりました。
コスト: TPU スクールは時間あたりの料金が安く、かつ完了が早いため、請求総額は2.1 倍安価でした。
- GPU 請求額: 約 119 ドル
- TPU 請求額: 約 56 ドル

なぜ TPU は速かったのでしょうか？
TPU チップを、超高速な内部ハイウェイ（ICI）を通じて互いに極めて素早くバトン（データ）を受け渡す 4 人のランナーのチームだと考えてください。GPU スクールには 2 人のランナーしかいませんでした。個々の GPU ランナーはわずかに速かったものの、TPU チームの協調性とデータを一緒に移動させる能力が、チーム全体としての勝利をもたらしました。

3. 試験（ロボットは学習したか？）

トレーニング後、彼らはロボットに Verilog コーディングの試験を受けさせました。

結果: GPU でトレーニングされたロボットは、最も難しい問題でわずかに高いスコア（約 5% 優位）を記録しました。
注意点: 著者たちは、この違いは一方が「プロンプト」部分を無視し、他方が無視しなかったなど、2 つの学校が試験の採点基準をわずかに異ならせたためであり、どちらかのロボットが本質的に賢かったからではないと指摘しました。この差は統計的に大きなものではありませんでした。

4. 就職面接（モデルの提供）

トレーニング後、モデルはリアルタイムでユーザーの質問に答えるために実用化されなければなりません。これを「推論」と呼びます。彼らは、異なる長さの質問にモデルが答える速度をテストしました。

短い質問（単純なタスク）:
- GPU スクールは、非常に短い質問への回答でわずかに速かったです。
- これらの迅速なタスクに対する回答ごとのコストもわずかに安価でした。
- 比喩: コーヒー 1 杯を買うだけなら、地元の店（GPU）の方がわずかに効率的です。
長い質問（複雑なタスク）:
- 質問が長くなると（4,000 語以上）、TPU スクールは圧倒的な差をつけました。
- 速度: TPU は、長い質問への回答開始（Time-to-First-Token）において23 倍速くでした。
- 容量: TPU は、速度を落とすことなく、同時に長い質問をする人を4 倍多く処理できました。
- 比喩: 小説 1 冊を書く必要がある場合、TPU スクールには巨大な図書館と、シームレスに協力して働ける作家チームがあります。一方、GPU スクールは圧倒され、本を落とし始めます。

5. 最終判断

著者たちは、彼らの特定のニーズ（大規模モデルのトレーニングとユーザーへの提供）において、Google の TPU が勝者であると結論付けました。

トレーニング: TPU は明確なチャンピオンです（速く、はるかに安価）。
推論: TPU は複雑なタスクや長いタスクにおいてチャンピオンです。非常に単純で短いタスクについては、GPU がわずかに優れています。
総コスト: トレーニングコストと 1 日分のユーザー提供コストを合計すると、TPU 構成は全体として1.8 倍安価でした。

結論:
TPU スクールへの切り替えには、新しい言語を学び、家具を修理し、橋を架けるなど、多くの事前の努力が必要でした。しかし、一度落ち着けば、その学校はより速く回り、コストが安く、従来の GPU スクールよりもはるかに大きな複雑なタスクを処理できました。重い AI 作業を行う企業にとって、切り替えのための追加努力は、節約とパフォーマンスに見合う価値がありました。

技術サマリー：Gemma 4 31B LoRA SFT における TPU と GPU の比較

問題定義

大規模言語モデル（LLM）のファインチューニングとサービングは、データの可用性ではなく、計算コストと反復速度によって次第に制約されるようになっています。Google Cloud TPU は高い理論的性能を提供しますが、成熟した PyTorch ベースの GPU エコシステムと比較して、多くのエンジニアリング努力を要すると見なされることがよくあります。本レポートは、標準的な GPU ネイティブなトレーニングレシピ（PyTorch、HuggingFace TRL、FSDP）を JAX ベースの TPU スタック（Tunix、Qwix、Orbax）へポートするエンドツーエンドのプロセスを文書化することで、オープンツールにおけるギャップに対処します。本研究は、Verilog コード生成のためのLoRA 教師ありファインチューニング（SFT）およびその後の推論サービングを特に対象とし、同一のハイパーパラメータ下での TPU と GPU ハードウェアの実証的な比較を提供します。

手法

著者は、Gemma 4 31B のファインチューニングとサービングを 2 つの異なるハードウェア構成で並行して比較しました。

トレーニング: TPU v5p-8（4 チップ）対 2× NVIDIA H100 80GB GPU。
推論: TPU v6e-8（Trillium、8 チップ）対 2× NVIDIA H100 80GB GPU。
データセット: トレーニングには CodeV-R1（1 万件の自然言語から Verilog へのペア）、評価には NVlabs/verilog-eval（156 件の仕様から RTL への問題）を使用しました。

エンジニアリング適応（GPU から TPU へ）

手法の核心は、トレーニングおよび推論パイプラインを PyTorch から JAX へ包括的にポートする作業でした。

フレームワークおよびメッシュ構成: PyTorch FSDP から明示的なデバイスメッシュ構成（v5p-8 の場合 fsdp=1, tp=4）を備えた JAX へ移行しました。著者は、Gemma 4 のアテンションヘッド数が 8 で割り切れないため、tp=8 は失敗したと指摘しています。
LoRA 注入および命名: HuggingFace（PyTorch）と Tunix（JAX）のモジュール名の不一致に対処しました。特に、Tunix は K と V の投影を単一の kv_einsum テンソルに融合させるため、チェックポイントの結合時にカスタム処理が必要でした。
シャーディングおよび最適化: Qwix によって注入されたパラメータがベース重みから不一致の PartitionSpec ランクを継承していた LoRA シャーディング注釈の修正を実装しました。オプティマイザは Optax AdamW に切り替えられ、カスタム状態シャーディングロジックが必要となりました。
データパイプライン: HuggingFace の DataLoader を Google の Grain ライブラリに置き換え、損失マスキング（システムプロンプトや推論チェーンの除去）およびシーケンス処理のためのカスタムロジックを実装しました。
チェックポインティングおよび変換: Orbax（JAX 形式）から LoRA アダプタを結合し、推論用に Safetensors（PyTorch/HF 形式）へ戻すカスタムパイプラインを開発しました。これには、テンソル形状の転置や融合された kv_einsum の分割への対応が含まれます。
推論セットアップ: 特定の Docker イメージ（vllm/vllm-tpu:gemma4）を使用して v6e-8 上で vLLM-TPU をデプロイしました。起動時の XLA コンパイル、HBM 割り当て、Gemma 4 の不均一なアテンション次元によるチャンク化されたマルチモーダル入力の無効化などの制約を克服しました。

主要な結果

トレーニング性能

速度: TPU v5p-8 は、2×H100 ベースライン（5.39 時間）と比較して 1.61 倍高速（3.34 時間）にトレーニングを完了しました。
スループット: TPU は 763 トークン/秒 を達成し、GPU の 486 トークン/秒を上回りました。
コスト: 時間単価の低さ（ $16.80 対$ 22.12）と高速な実行により、TPU によるトレーニングは 2.12 倍安価（ $56.11 対$ 119.23）でした。
損失収束: 両モデルは同様の収束軌跡を示しました。絶対損失の差（TPU が 0.072、GPU が 2.258）は、損失計算範囲の違い（TPU はアシスタントのみの損失を使用、GPU は完全シーケンス損失を使用）に起因しました。

推論性能

短コンテキスト（≤2048 トークン）: GPU はピーク出力スループットでわずかな優位性（+6%）を示し、P99 尾部レイテンシもわずかに良好でした。これは、チップあたりの HBM バンド幅が高く、XLA バケットパディングのオーバーヘッドがないことに起因します。
長コンテキスト（≥4096 トークン）: TPU v6e-8 は 優位性 を示しました。
- スループット: 4096 トークンにおいて出力トークン/秒が 66% 高い 値を示しました。
- TTFT（Time-to-First-Token）: 4096 トークンで 23.6 倍高速（61 ms 対 1,443 ms）、8192 トークンで 7.1 倍高速 でした。
- スケーラビリティ: TPU は fp8 KV キャッシュ（実効容量の倍増）と高い集約計算能力により、中程度のコンテキストで最大 32 の同時リクエスト（QPS）を維持しましたが、GPU は約 8 で飽和しました。
コスト: 長コンテキストのワークロードにおいて、TPU は出力トークンあたり 41% 安価 でした。

評価品質

Verilog 生成ベンチマークにおいて、GPU でトレーニングされたモデルはわずかに高い pass@1（0.697 対 0.641） および pass@5（0.814 対 0.795） を達成しました。著者は、このわずかな差異をハードウェアの能力ではなく、トレーニング目的の違い（損失マスキング）および確率的サンプリングのばらつきに起因するとしています。

意義と主張

本論文は、TPU ハードウェア上での Gemma 4 31B のファインチューニングおよびサービングの 最初のエンドツーエンド実証 を提供し、オープンツールエコシステムにおける重要なギャップを実質的に解消したと主張しています。

生産環境への適合性: 本作業は、特定の Docker 構成、メッシュ設定、チェックポイント変換スクリプトを含む、TPU インフラストラクチャ上で Gemma 4 をデプロイするための再現可能かつ生産環境対応のレシピを提供します。
総所有コスト（TCO）: 本研究は、代表的なトレーニングおよびサービングのワークロードにおいて、TPU 構成が GPU ベースラインよりも 1.82 倍安価 であると結論付けています。この優位性は主にトレーニングコストの節約によって駆動されており、推論コストは長コンテキストアプリケーションにおいて大幅に累積します。
ハードウェアの適合性: 著者は、TPU を 長コンテキストワークロード およびコストに敏感な生産パイプラインのための優先インフラストラクチャとして位置付けていますが、短コンテキストスループット および エコシステムの成熟度（より単純なプロトタイピング）における GPU の優位性も認めています。

本レポートは、JAX/Tunix へのポートに伴うエンジニアリングオーバーヘッド（メッシュ構成、シャーディング修正、パイプライン再構築に約 1 週間と推定）は無視できないものであると強調しつつも、これは多数のトレーニング実行にわたって償却される 一回性のコスト であり、大規模 LLM 適応のための TPU は実行可能かつ優れた代替手段であると結論付けています。