2025 年第 1 四半期の AI Hypercomputer デベロッパー エクスペリエンスの強化: 迅速に構築し、大規模にスケール

※この投稿は米国時間 2025 年 5 月 17 日に、Google Cloud blog に投稿されたものの抄訳です。

ノートブックでイテレーションを行う場合でも、大規模なクラスタをオーケストレートする場合でも、最先端の AI モデルの構築はワクワクするものです。しかし、トレーニングのスケールアップには、複雑なインフラストラクチャへの対応、多数のインスタンス間でのソフトウェアや依存関係の構成、パフォーマンスのボトルネックの特定など、大きな課題を伴うことがあります。

Google Cloud は、規模に関係なく AI トレーニングを容易にすることに注力しています。Google は、TPU や GPU などの強力なハードウェアだけでなく、デベロッパーの皆様の生産性を向上させるために設計された一連のツールと機能によって、AI Hypercomputer システムを継続的に進化させています。インタラクティブな開発から、トレーニングの最適化、デプロイの簡素化まで、ワークフローの合理化に役立つ最新の機能強化について、いくつか詳しく見ていきましょう。

Pathways on Cloud でノートブックからスケール

Jupyter ノートブックが提供する迅速なイテレーションは気に入っていても、数千のアクセラレータにスケールするとなると、使い慣れた環境から離れることになります。同時に、ワークロードを大規模に実行するためには、さまざまなツールを学ばなければならず、現実的とは言えません。また、短時間で終わる可能性のある反復テストのために、大規模なアクセラレータ クラスタを何週間も拘束することも然りです。

使いやすさと大規模なスケールを天秤にかけるべきではありません。JAX を使用すると、1 つのアクセラレータ用のコードを記述し、それを数千のアクセラレータ向けに簡単にスケールアップできます。大規模でマルチタスクのスパース活性化 ML システムを作成するためのオーケストレーション システム、Pathways on Cloud は、このコンセプトをさらに進化させ、インタラクティブなスーパーコンピューティングを実現します。Pathways は、アクセラレータのプールを動的に管理し、数千台ものデバイスにわたってデータの移動と計算をオーケストレートします。その結果、Jupyter ノートブックから直接 1 つのみのアクセラレータでテストを開始し、それを調整して、同じインタラクティブなセッション内で数千のアクセラレータにスケールできます。このようにして、規模を犠牲にすることなく、研究開発において迅速なイテレーションが可能になりました。

Pathways on Cloud を使用すると、規模に応じてコードを書き直す必要がなくなります。数時間しかかからないテストのために、何週間もハードウェアをオーバープロビジョニングするのはやめましょう。すでに活用している技術や知識を基に、イテレーションを高速化し、スーパーコンピューティングの能力をオンデマンドで活用しましょう。こちらの動画をご覧になり、Pathways on Cloud で真のインタラクティブなスケーリングを実現する方法をご確認ください。ただ Google Kubernetes Engine（GKE）クラスタで JupyterHub を実行する以上のことを行えます。

Xprofiler で迅速にデバッグし、スマートに最適化

ジョブをスケールアップする際には、アクセラレータが使用されていることを知っているだけでは不十分です。アクセラレータがどのように使用されているかと、速度低下やクラッシュの原因も理解しておく必要があります。そうでなければ、実行全体を停止させる厄介なメモリ不足エラーを検出するにはどうすればよいでしょうか。

Google Cloud アクセラレータの詳細なパフォーマンス分析用ツール、Xprofiler ライブラリをご紹介しましょう。これはコード実行のプロファイリングとトレースを行い、特に、XLA コンパイラによって生成された高レベル演算（HLO）に関する重要な分析情報を提供します。Xprofiler を使用すると、実用的な分析情報を簡単に取得できます。コマンドラインから Xprofiler インスタンスを起動するだけで、実行中の詳細なプロファイルとトレースログをキャプチャできます。その後、TensorBoard を使用してこのデータをすばやく分析します。パフォーマンスのボトルネックを可視化し、ルーフライン分析でハードウェアの制限を把握し（ワークロードが計算依存型かメモリ依存型か）、エラーの根本原因をすばやく特定できます。Xprofiler は、最高のパフォーマンスを発揮するようにコードを最適化し、AI インフラストラクチャを最大限に活用できるようにします。

コンテナ イメージで設定の手間を省く

多数の優れた AI フレームワークやライブラリから選択できるとはいえ、適切なドライバや依存関係を正しく構成するのは複雑で時間がかかる可能性があります。特に数百または数千のインスタンスにスケールする際には、構成を誤ると、費用のかかるエラーや遅延につながりかねません。こうした難題を回避するために、Google は一般的な AI 開発ニーズに合わせて設計された、ビルド済みの最適化されたコンテナ イメージを提供しています。

GPU 上の PyTorch では、GPU で高速化されたインスタンス コンテナ イメージにより、すぐに実行できる環境を提供します。Google は NVIDIA と緊密に連携して、NVIDIA CUDA ツールキット、NCCL、NVIDIA NeMo などのフレームワークといった重要なソフトウェアのテスト済みバージョンを組み込んでいます。Canonical のおかげで、これらは最適化された Ubuntu LTS で実行されます。最適なパフォーマンスが得られるように調整された安定した環境ですぐに作業を開始できるため、互換性の課題を回避し、設定時間を大幅に節約できます。

また、（TPU または GPU で）JAX を使用している場合は、Google Cloud 上の JAX for AI 向けにキュレートされたコンテナ イメージとレシピを使用して、効率的に作業を開始できます。テスト済みのすぐに使える JAX 環境を使用して、手動での依存関係の追跡や構成の手間を省きましょう。

実績のあるレシピで GPU トレーニングの効率を向上

設定時間の短縮だけでなく、特に大規模な場合には、トレーニング中の有用なコンピューティング時間（「ML グッドプット」）を最大化することが重要です。ジョブの失敗によって無駄なサイクルが生じると、費用が大幅に膨らみ、結果を得るのが遅れる可能性があります。そのため、Google は、このような課題に対処するための手法とすぐに使えるレシピを提供しています。

非同期や複数ティアのチェックポイントなどの手法は、トレーニングの速度を落とさずにチェックポイントの頻度を高め、保存 / 復元オペレーションを高速化します。AI Hypercomputer は、リセット、ホットスワップ、スケーリングのいずれかのアクションをインテリジェントに選択して、中断を自動的に処理できます。NVIDIA とのパートナーシップにおいて作成された ML グッドプット レシピは、これらの手法をバンドルして、NVIDIA NeMo や NVIDIA Resiliency Extension（NVRx）と統合し、Google Cloud での PyTorch トレーニングの効率と信頼性を高める包括的なソリューションを提供します。

また、Google Cloud Storage や Parallelstore などのさまざまなストレージ オプションのトレーニング パフォーマンスをベンチマーク評価できる、最適化されたレシピ（チェックポイント機能を含む）も追加されました。そして最後に、NVIDIA で高速化されたインスタンス（NVIDIA Blackwell 上に構築）、A4 向けのレシピが追加されました。トレーニング レシピには、PyTorch と JAX を使用した、最大 512 個の Blackwell GPU をトレーニングするスパースおよび高密度のモデルが含まれます。

MaxText による最先端の JAX LLM 開発

Google Cloud で LLM に JAX を使用するデベロッパー向けに、MaxText は TPU と GPU の両方での高度なトレーニング、チューニング、サービングを提供します。また最近、教師ありファインチューニング（SFT）や Direct Preference Optimization（DPO）などの主要なファインチューニング手法に加え、一時停止と再開や弾力性のあるトレーニングなどの復元力のあるトレーニング機能のサポートが追加されました。MaxText は、Google が NVIDIA と共同で開発した JAX 最適化とパイプライン並列処理技術を活用して、数万台の NVIDIA GPU にわたってトレーニング効率を向上させます。また、最新のオープンモデルである Gemma 3、Llama 4 のトレーニングと推論（Scout と Maverick）、DeepSeek v3 のトレーニングと推論のサポートとレシピも追加されました。

Trillium TPU で最高のパフォーマンスを得られるよう、行列乗算、集団コンピューティング、高帯域幅メモリ（HBM）テストなど、数百のアクセラレータで複数のスライスにスケールアップするマイクロベンチマーク レシピが追加されました。これらの指標は、パフォーマンスの最適化に特に役立ちます。GKE の本番環境ワークロードについては、自動アプリケーション モニタリングをぜひご検討ください。

PyTorch / XLA 2.7 と torchprime を使用して TPU で PyTorch を活用

Google は、TPU で PyTorch を使用するユーザーに、高パフォーマンスの統合エクスペリエンスを提供できるよう尽力しています。そのため、最近リリースされた PyTorch / XLA 2.7 には、パフォーマンスの大幅な改善が含まれており、これは特に TPU 上で vLLM を推論に使用しているユーザーにとってメリットとなります。このバージョンでは、柔軟性と相互運用性を高める重要な新機能も追加されており、PyTorch / XLA コード内から JAX 関数を直接呼び出せるようになりました。

さらに、TPU 上で PyTorch / XLA の機能を活用できるように、torchprime が導入されました。これは、TPU 上で PyTorch モデルをトレーニングするためのリファレンス実装です。Torchprime は、大規模で高パフォーマンスなモデル トレーニングのためのベスト プラクティスを示すように設計されているため、PyTorch / XLA 開発の第一歩として最適です。

RecML を使用して最先端のレコメンダーを構築

生成 AI が注目を集めることが多いとはいえ、非常に効果的なレコメンダー システムが多くのアプリケーションの基盤であることに変わりはありません。こうしたシステムの大規模なトレーニングにおいては、TPU 固有のメリットがあります。ディープ ラーニング レコメンダー モデルは、大規模なエンベディング テーブルに依存してユーザー、アイテム、その特徴を表現することが多いため、これらのエンベディングを効率的に処理することが不可欠です。この点において、TPU が強みを発揮します。特に、特殊な統合データフロー プロセッサである SparseCore と組み合わせて使用する場合には顕著です。SparseCore は、レコメンダーではよくある、膨大なスパース エンベディングのルックアップと処理を高速化する目的に特化したもので、他の方法と比べてトレーニングを大幅に高速化します。

この機能を活用できるように、Google は RecML を提供しています。これは、TPU 向けに最適化された、使いやすく高パフォーマンスの、大規模なディープ ラーニング レコメンダー システム ライブラリであり、BERT4Rec、Mamba4Rec、SASRec、HSTU などの最先端のレコメンダー モデルをトレーニングするためのリファレンス実装を提供します。RecML は SparseCore を使用してパフォーマンスを最大化するため、TPU ハードウェアを効率的に活用して、レコメンダー モデルのトレーニングとスケーリングを簡単に高速化できます。

一緒に構築しましょう

AI デベロッパー エクスペリエンスの向上は、Google Cloud の継続的なミッションです。Pathways を使用したインタラクティブなテストのスケーリングから、Xprof を使用したボトルネックの特定や、最適化されたコンテナとフレームワークのレシピを使用した迅速な開始まで、こうした AI Hypercomputer の改善点によってプロセスが円滑化されるため、より迅速にイノベーションを達成できます。これらは Google Cloud Next 25 で発表した AI Hypercomputer の他のイノベーションを基盤としています。

これらの新機能の詳細を調べ、コンテナ イメージをスピンアップし、JAX と PyTorch のレシピを試し、MaxText、torchprime、RecML などのオープンソース プロジェクトに貢献しましょう。皆様のフィードバックが、Google Cloud における AI 開発の未来を形作ります。ぜひ一緒に未来を築き上げましょう。

-AI Infrastructure 担当グループ プロダクト マネージャー、Alex Spiridonov

-AI Infrastructure 担当シニア プロダクト マネージャー、Niranjan Hira