T最新テックトレンド

Google I/O 2024: AIモデル開発の未来を拓くGoogle Cloud TPUと包括的ソフトウェアスタック

Google for Developers
0:00--:--

近年、人工知能（AI）の進化は目覚ましく、その中でも特に大規模言語モデル（LLM）を中心とした生成AI技術は、私たちの生活やビジネスに変革をもたらし続けています。Googleは、AI分野におけるイノベーションの最前線を走り続けており、その原動力となっているのが、同社が独自に開発したAIアクセラレータであるTPU（Tensor Processing Unit）と、それを最大限に活用するための包括的なソフトウェアスタックです。

Google I/O 2024では、GoogleがAIの急速な進化をどのように推進してきたか、そして開発者が最先端のAIモデルを構築、学習、デプロイ、運用するための最新のハードウェアとソフトウェアツールが発表されました。本記事では、これらの発表内容を深く掘り下げ、その重要性、具体的な機能、ビジネスへの影響、そして将来性について詳細に解説します。

AIモデル開発の4つの段階とスケーリングの新次元

大規模言語モデル（LLM）のような複雑なAIモデルを構築し、運用するプロセスは、大きく分けて4つの段階に分類されます。

  1. モデル構築（Model Building）: ニューラルネットワークのアーキテクチャを設計し、基本的なフレームワーク上で実装する段階です。
  2. 事前学習（Pre-training）: 大規模なデータセットを用いてモデルに基本的な知識とパターンを学習させる段階で、主に次のトークンを予測する能力を習得させます。
  3. 事後学習（Post-training）: 事前学習済みモデルを特定のタスクやユーザーの指示に従うように微調整し、より有用なモデルに仕上げる段階です。これには、教師ありファインチューニングや強化学習が含まれます。
  4. 推論（Inference）: 学習済みのモデルを実際にユーザーに提供し、新しいデータに対して予測や応答を生成させる段階です。

GoogleのAIイノベーションのタイムラインを見ると、Gemini、Gemma、Vio、Imagen、そして最新のOmniといった多様なモデルやモダリティが、驚くべきペースでリリースされてきたことがわかります。これらのモデルは、マルチモーダルな理解、高度な推論能力、そして複雑なタスクを自動実行するエージェントワークフローといった画期的な機能を次々と解放してきました。

特に注目すべきは、AIの「知能（Intelligence）」が向上するにつれて、モデルを学習させるための「計算資源（Compute）」の使われ方が変化している点です。従来の事前学習や事後学習だけでなく、推論フェーズ、特に「Thinking Models」と呼ばれる、推論時に多数のトークンを消費して複雑な問題を熟考するタイプのモデルにおいて、膨大な計算資源が求められるようになってきました。これにより、AIモデルを支えるハードウェアとソフトウェアスタックは、学習と推論のそれぞれのニーズに特化した最適化が不可欠となっています。

Google Cloud TPUのハードウェア革新：学習と推論の専門化

Googleは、このようなAIワークロードの変化に対応するため、TPUハードウェアのアーキテクチャを学習用と推論用に専門化させるという革新的なアプローチを採用しています。

TPU 8t: 大規模学習ワークロードに最適化された設計

TPU 8tは、大規模なAIモデルの学習ワークロードに特化して設計された、Googleの最新世代のTPUです。その主な特徴は以下の通りです。

  • 圧倒的な計算能力: TPU 8tは、1ポッドあたり121エクサフロップスのFP4演算能力を誇ります。これは、大規模なデータセットと複雑なモデルの学習に不可欠な、驚異的な処理能力です。
  • 優れたワットあたりの性能: 従来のTPUと比較して、ワットあたりの性能が2倍向上しています。TPU 8tは、環境負荷を低減しつつ、より効率的な学習を可能にします。
  • スループットとスケーリング効率の最大化: 大規模な学習ジョブでは、モデルを複数のチップやノードに分散させて学習させる「並列処理」が不可欠です。TPU 8tは、このスケーリング効率を最大限に高めることに重点を置いており、モデルの規模が大きくなっても線形に近いスケーリング性能を維持できるよう設計されています。

TPU 8tは、何日、何週間にもわたる膨大な計算を必要とする基盤モデルの事前学習や、特定のドメインデータを使った大規模な教師ありファインチューニングにおいて、比類ない性能を発揮します。

TPU 8i: 低レイテンシ推論ワークロードに最適化された設計

一方で、TPU 8iは、学習済みのモデルを使った推論ワークロードに特化して最適化されています。推論のニーズは学習とは異なり、スループットだけでなく、ユーザーからのリクエストに対する応答速度（レイテンシ）とコスト効率が非常に重要になります。

  • 低レイテンシの実現: TPU 8iは、従来の3D Torusトポロジーと比較して、レイテンシが2倍低いという特徴を持っています。これにより、リアルタイム性が求められるインタラクティブなAIアプリケーションにおいて、迅速な応答を可能にします。
  • コスト効率の向上: 推論ワークロードは、学習ワークロードと比較して、より多くのユーザーによって頻繁に利用される可能性があります。TPU 8iは、このような大量の推論リクエストを低コストで処理できるよう設計されており、AIサービスの運用コスト削減に貢献します。

TPU 8tとTPU 8iの専門化されたアーキテクチャは、GoogleがAIの進化をさらに加速させるための戦略的な一歩であり、開発者がそれぞれのワークロードの特性に最適なハードウェアを選択できるようになります。これにより、リソースの最適化とAIモデルの性能最大化が同時に達成されます。

TPUの能力を最大限に引き出すソフトウェアスタック

優れたハードウェアだけでは、AIの真の可能性を引き出すことはできません。Googleは、TPUハードウェアの能力を最大限に引き出すため、革新的なソフトウェアスタックと開発ツールを提供しています。

vLLM TPU inference: 高効率な推論サービス

vLLM TPU inferenceは、JAXおよびPyTorchフレームワークで構築されたLLMに対し、低レイテンシで高スループットな推論サービングを提供する基盤です。特に大規模なコンテキスト長と高い並列度が求められる現代のLLM推論において、以下の重要な課題を解決します。

  • 効率的なメモリ最適化:
    • PagedAttention: LLMの推論において、Attentionメカニズムで使用されるKey-Value（KV）キャッシュは、コンテキスト長の増加とともに膨大なメモリを消費し、メモリ断片化の問題を引き起こします。PagedAttentionは、KVキャッシュを固定サイズのブロックに仮想化することで、メモリの利用効率を劇的に向上させ、メモリ断片化を排除します。
    • Prefix caching: エージェントワークフローや対話型アプリケーションでは、複数のリクエストが共通のプレフィックスを持つことがよくあります。vLLMは、以前に計算されたKVキャッシュを再利用することで、冗長な計算を省き、メモリ効率とレイテンシを改善します。
    • Chunked prefill: 長いプロンプトの処理を効率化します。
  • インテリジェントなスケジューリングシステム:
    • Continuous batching: 推論リクエストはプロンプト長、生成時間、到着パターンが不規則です。従来の静的なバッチ処理ではなく、vLLMはトークンレベルで動的にワークロードをスケジューリングすることで、ハードウェア利用率を向上させ、低レイテンシを実現します。
    • Sequence parallelism: モデルを複数のデバイスに効率的に分散し、並列処理を最適化します。
    • TPU aware scheduling: TPUのハードウェア特性を考慮したスケジューリングにより、TPUの利用率を最大化します。
  • ハードウェア効率:
    • Speculative decoding: より小さな補助モデルが先行してトークンを予測し、それを大規模モデルが検証することで、推論速度を向上させます。
    • Prefill/decode disaggregation: プロンプト処理とトークン生成のフェーズを分離し、それぞれのフェーズに最適化された処理を行います。
    • TPU runtime optimizations: TPUのメモリレイアウトとハードウェア認識実行戦略を最大限に活用し、スループットを最大化します。
  • 統一された推論バックエンド: vLLM TPU inferenceは、PyTorchとJAXの両方のフレームワークをサポートし、GPUとTPUの双方で動作するポータビリティを提供します。これにより、開発者は基盤となるハードウェアやソフトウェアスタックを気にすることなく、既存のアプリケーション層を変更せずにモデルをデプロイ、運用できます。

デモ: Food Tracking Assistant 動画では、Gemma 4の31BパラメータモデルをvLLM TPU inferenceでサービングするデモが紹介されました。ユーザーが料理の画像とテキストを入力すると、AIエージェントが以下のステップで栄養情報をログに記録します。

  1. コンテンツの特定: 画像とテキストから料理の材料（鶏肉、米、野菜など）を分析・識別します。
  2. 栄養情報の検索: 外部の「栄養データベース」をツールコールで呼び出し、特定された材料の栄養価を調べます。
  3. 要約の生成: カロリー、タンパク質、炭水化物、脂質などの合計を計算し、最終的な出力として提示します。 このデモは、vLLMがマルチモーダル入力、エージェントによるツール利用、および複雑な推論ワークフローをいかに効率的に処理できるかを示す具体的な例です。

Tunix: 強化学習と事後学習の加速

Tunix（Tune-in-JAX）は、Googleが開発したJAXベースの軽量なフレームワークで、AIモデルの事後学習（Post-training）の全段階をサポートします。特に、モデルを人間の価値観や特定のタスクにアライメントさせるための強化学習において、その真価を発揮します。

  • 多岐にわたる学習テクニックのサポート: Tunixは、教師ありファインチューニング（SFT）、強化学習（PPO、GRPO、GSPO、RPOなど）、知識蒸留（Knowledge Distillation）といった最先端の事後学習技術をサポートします。これにより、開発者はモデルのニーズに応じて最適な学習アプローチを選択できます。
  • エージェント的ワークフローへの対応: Tunixは、複雑なエージェント的ワークフローにおけるモデルの訓練をサポートします。これにより、モデルが現実世界の課題を解決するために、外部ツールと連携したり、複数の思考ステップを踏んだりする能力を向上させることができます。
  • Kaggleでの活用: Tunixは、Kaggleハッカソンでも活用され、数千人の参加者がTunixを使って古いGemmaモデルに数学の問題を解かせ、推論の跡を示すように学習させました。これは、Tunixが研究コミュニティで広く利用され、実用的な成果を生み出していることを示しています。
  • 知識蒸留の革新: Tunixは知識蒸留をサポートしており、大規模モデルの知識をより小さく、より効率的なモデルに転移させることを可能にします。これは、大規模モデルの計算コストが高い推論フェーズにおいて、特に有用な技術です。大規模モデルが持つ豊富な「確率分布」情報を、より小さなモデルの学習シグナルとして活用することで、高い性能を維持しながらも、より高速で安価な推論を実現できます。

デモ: GRPOによる強化学習 動画では、TunixとGRPO（Group Relative Policy Optimization）を使って、Gemma 3の4Bパラメータモデルをファインチューニングするデモが紹介されました。このデモの目的は、少ない計算資源でモデルの推論能力を向上させることです。

  1. Actor & Reference モデルの定義: 学習対象となるActorモデルと、学習の安定性を保つためのReferenceモデルを設定します。
  2. データセットのセットアップ（Grainを使用）: 大規模なラベル付きデータセットが不要な強化学習の特性を活かし、GoogleのデータパイプラインライブラリであるGrainを用いて、モデルがエージェントとして振る舞う際のシナリオをシミュレーションするためのデータセットを効率的に準備します。
  3. 報酬のスコアリング: モデルの出力が望ましい振る舞いをしているか（例えば、ツールを正しく呼び出し、推論の思考プロセスを示し、適切な形式で出力しているか）を評価するためのルールベースの報酬システムを定義します。
  4. モデルの学習: 定義された報酬システムに基づき、TPU上でモデルを学習させます。Tunixは、チェックポイントの保存、学習進捗のログ記録、TPU全体でのワークロード分散といった重い処理を自動で処理します。
  5. TensorBoardによる評価: 学習の成果はTensorBoardでリアルタイムに可視化され、報酬カーブやさまざまなメトリクスを通じて、モデルの学習状況とベースラインモデルとの比較を行います。

MaxText: フロンティアモデルの学習を再現可能に

MaxTextは、JAXで実装された高性能な大規模言語モデルのリファレンス実装であり、フロンティアスケールのAIモデル学習のためのGoogleの知見と経験を凝縮したものです。

  • フロンティアスケールモデルの学習を簡素化: MaxTextは、Gemma、Qwen、DeepSeek、Llama、Mistralといった主要なオープンソースモデル向けの最適化されたレシピを提供します。これにより、開発者はゼロから学習パイプラインを構築する手間を省き、これらのモデルを容易に学習、微調整できます。
  • スケーラブルな学習アーキテクチャ: シングルホストマシンから、数千のチップを擁する大規模なマルチノードクラスターまで、シームレスにスケーリングできる学習アーキテクチャを提供します。これにより、研究者は小規模な実験から大規模なモデル学習まで、同じコードベースで対応できます。
  • JAXエコシステムとの統合: MaxTextは、Flax（ニューラルネットワークライブラリ）、Orbax（チェックポイント管理）、Grain（データパイプライン）、Chex（JAXコードのテスト）、Qwik（JAXのユーティリティ）、Tunix（事後学習）といったJAXエコシステムのツールと緊密に統合されています。これにより、学習プロセス全体を通じて一貫性のある開発体験を提供し、研究者がインフラ管理ではなくモデル開発に集中できるようになります。
  • Googleのハードラーニングをレシピに: 大規模モデルの学習には、分散コンピューティング、メモリ管理、並列処理など、数多くの複雑な課題が伴います。MaxTextは、Googleがこれらの課題を解決してきた経験から得た「ハードラーニング」を、すぐに使えるトレーニング設定や構成として提供することで、開発者がこれらの複雑さを抽象化し、再現性の高い研究開発を行えるようにします。

JAXとTorchTPU: 柔軟なモデル構築フレームワーク

AIモデル構築の根幹を支えるのが、フレームワークの選択です。Googleは、JAXとTorchTPUという2つの強力な選択肢を提供し、開発者に柔軟性と高性能なコンピューティングへのアクセスを提供しています。

  • JAX: 表現力豊かな数値計算ライブラリ

    • JAXは、Pythonで数値計算を行うためのライブラリであり、NumpyライクなAPIを持ちながら、自動微分、JITコンパイル、Vmap（ベクトル化）、Pmap（並列化）といった「Composable Transformations」を提供します。
    • 自動微分（Grad）: シンプルなPython関数で複雑な勾配計算を自動で行えます。
    • JITコンパイル（Jit）: Python関数をXLA（Accelerated Linear Algebra）コンパイラによって最適化されたアクセラレータプログラムにコンパイルし、実行速度を劇的に向上させます。
    • Vmap/Pmap: バッチ処理やデバイス間の並列処理を簡潔に表現し、スケーラブルなモデル学習を可能にします。
    • JAXは、そのミニマルな設計と強力な変換機能により、大規模なMLシステムを構築するための理想的な基盤となっています。開発者は低レベルの最適化を意識することなく、ハードウェアの能力を最大限に引き出すことができます。
    • JAXエコシステム全体が、Flax, Optax, Orbax, Grain, Chex, Qwik, Tunixといったライブラリ群で構成され、モジュール管理、パラメーター管理、データ処理、テスト、推論など、AIモデル開発のあらゆる側面をサポートします。
    • JAXが提供する一貫した開発体験は、CPU、GPU、TPUといった異なるハードウェアバックエンドにおいても、同じコードベースで高性能なAIモデルを構築することを可能にします。

  • TorchTPU: PyTorchとTPUのシームレスな統合

    • TorchTPUは、MetaのPyTorchチームとの緊密な協力のもと、PyTorchユーザーがTPUを容易に利用できるようにするためにゼロから再構築された新しいアプローチです。
    • 完全にネイティブ: 新しいAPIを学ぶ必要はありません。PyTorchの既存のAPIと機能を使って、TPU上でモデルを実行できます。
    • 高いポータビリティ: コードの変更を最小限に抑え、あるいはゼロコード変更で、既存のPyTorchモデルをGPUからTPUに移行できます。
    • 優れたパフォーマンス: Torch Compile、カーネルサポート、およびプロファイリングツールを通じて、TPU上でのモデル実行性能を最適化します。
    • エコシステムの活性化: PyTorchコミュニティ全体がTPU上でモデルを構築、学習できるようになり、TPUエコシステムの成長とイノベーションを促進します。
    • 使いやすさ: TorchTPUは、TPUへの移行を容易にするためにゼロから設計されました。これにより、開発者はインフラ管理の複雑さに煩わされることなく、モデル開発に集中できます。
    • Kinetic: Kerasチームが開発したKineticは、Python関数をTPUまたはGPUでリモート実行するためのシンプルなライブラリです。デコレータ1つで分散コンピューティングを抽象化し、インフラ管理の手間を大幅に削減します。これは、TorchTPUが目指す使いやすさを象徴するプロジェクトの一つです。

OpenXLA: ハードウェアとフレームワークを繋ぐユニバーサルコンパイラ

JAXとTorchTPUといった上位レベルのフレームワークの基盤となっているのが、OpenXLAというオープンソースコンパイラです。

  • 汎用オープンソースコンパイラ: OpenXLAは、JAX、PyTorch、TensorFlowといった多様なMLフレームワークと、TPU、GPU、CPUなどの様々なハードウェアアクセラレータを橋渡しする役割を担います。
  • 分散コンピューティングの自動化: TPUトポロジービジュアライザーで示されるように、TPUは数百、数千のチップが高速なインターコネクトで接続された複雑なネットワークを形成しています。OpenXLAは、この複雑な分散コンピューティング環境を自動化し、開発者がハードウェアの低レベルな詳細を意識することなく、大規模な並列計算を効率的に実行できるようにします。
  • 「単一の巨大なコンピュータ」としてコードを記述: OpenXLAの目標は、開発者がまるで単一の巨大なコンピュータ上でコードを書いているかのように感じられる環境を提供することです。これにより、開発者は、モデルのアーキテクチャやアルゴリズムの革新に集中できます。

ビジネスへの影響と将来性

GoogleのTPUと包括的なソフトウェアスタックの進化は、AI開発の未来に多大な影響を与えます。

  1. AI開発の民主化と加速: MaxTextが提供するバトルテスト済みのレシピや、Tunixによる強化学習の容易化、TorchTPUによるPyTorchからのシームレスな移行は、AIモデル開発の参入障壁を大幅に下げます。これにより、より多くの研究者や開発者が、最先端のAI技術を活用し、自身のアイデアを迅速に具現化できるようになります。
  2. コスト効率の向上: TPU 8iのような推論特化型ハードウェアと、vLLM TPU inferenceのような効率的な推論サービスは、AIモデルの運用コストを削減し、より多くの企業がAIサービスを大規模にデプロイできるようになります。また、知識蒸留のような技術は、小規模で効率的なモデルの利用を促進し、さらなるコスト削減に貢献します。
  3. 革新的なAIアプリケーションの創出: マルチモーダル理解やエージェント的ワークフローといった機能は、Food Logging Assistantのような新たなタイプのAIアプリケーションを可能にします。これらのアプリケーションは、従来のAIでは実現不可能だった、より複雑で人間らしいインタラクションや自動化されたタスク実行を提供し、さまざまな産業に新たなビジネス価値をもたらすでしょう。
  4. 研究と開発の生産性向上: 開発者が低レベルのインフラ管理や最適化の複雑さから解放され、モデルの設計、アルゴリズムの改善、新しい機能の探求といった本質的な研究開発に集中できるようになることで、AI分野全体のイノベーションが加速します。
  5. オープンソースコミュニティへの貢献: Googleは、MaxText、Tunix、OpenXLAといった多くのプロジェクトをオープンソースとして公開し、Meta-PyTorchなどのコミュニティと協力することで、AI技術の透明性とアクセス性を高めています。これは、AIエコシステム全体の発展にとって極めて重要です。

Googleは、「AI Systems DevLabs 2026」のようなイベントを通じて、これらの技術に関する深い知見を共有し、開発者コミュニティとの連携を強化しています。これは、AIの未来を共に築き、その恩恵を広く社会にもたらそうとするGoogleのコミットメントを示しています。

まとめ

Google I/O 2024で発表されたGoogle Cloud TPUとそれに付随するソフトウェアスタックは、AIモデル開発のあらゆる段階において、開発者体験を根本から変革する可能性を秘めています。学習と推論に特化したハードウェア、効率的なメモリ管理やスケジューリングを実現する推論サービス、強化学習やフロンティアモデル学習を簡素化するフレームワーク、そしてそれら全てを支える汎用コンパイラOpenXLA。これらすべてが連携することで、AIの進化はさらに加速し、より高度で、よりアクセスしやすく、より効率的なAIの未来が切り開かれることでしょう。Googleは、ハードウェアとソフトウェアの両面からAI開発を支援し、AIの可能性を最大限に引き出すための道を舗装し続けています。

← 一覧に戻る