JAX と Cloud TPU による本番環境 AI の構築

JAX AI スタックは、Google がサポートする構成可能なライブラリのコレクションで JAX 数値コアを拡張し、大規模な ML 用の堅牢でエンドツーエンドのオープンソース プラットフォームに進化させます。そのため、JAX AI スタックは、ML ライフサイクル全体を網羅する堅牢なエコシステムで構成されています。

  • 産業規模の基盤: JAX AI スタックは大規模なスケール向けに設計されています。数万個のチップにわたるトレーニングのオーケストレートに ML Pathways、復元力とスループットの高い非同期チェックポインティングに Orbax を活用し、最先端のモデルのプロダクション グレードのトレーニングを実現しています。

  • 完全な本番環境対応のツールキット: JAX AI スタックは、柔軟なモデル作成のための Flax、構成可能な最適化戦略のための Optax、再現可能な大規模実行に不可欠な決定論的データ パイプラインの Grain など、開発プロセス全体に対応する包括的なライブラリ セットを提供しています。

  • ピーク時の特殊なパフォーマンス: ハードウェアの利用率を最大化するため、JAX AI スタックは、最先端のカスタム カーネルに Tokamax、トレーニングと推論の速度を向上させる非侵入型の量子化に Qwix、ハードウェア統合型の詳細なパフォーマンス プロファイリングに XProf など、特殊なライブラリを提供しています。

  • 本番環境へのフルパス: JAX AI スタックは、研究からデプロイへのシームレスな移行を実現します。これには、基盤モデルのトレーニングのスケーラブルなリファレンスとしての MaxText、最先端の強化学習（RL）とアライメントのための Tunix、vLLM TPU 統合と JAX サービング ランタイムによる統合推論ソリューションが含まれます。

JAX AI スタックの哲学は、疎結合のコンポーネントのそれぞれを有効に活用することです。JAX はモノリシックな ML フレームワークではなく、スコープが狭く、効率的な配列演算とプログラム変換に重点を置いています。このコア フレームワーク上にエコシステムが構築され、ML モデルのトレーニングと、科学計算などの他のタイプのワークロードの両方に関連する幅広い機能が提供されています。

この疎結合コンポーネントのシステムにより、要件に最適な方法でライブラリを選択して組み合わせることができます。ソフトウェア エンジニアリングの観点から見ると、このアーキテクチャでは、従来はコア フレームワーク コンポーネントと見なされていた機能（データ パイプラインやチェックポイントなど）を、コア フレームワークを不安定にするリスクや、リリース サイクルに巻き込まれるリスクを冒すことなく、反復的に更新できます。ほとんどの機能はモノリシック フレームワークの変更ではなく、ライブラリに実装されているため、コア数値ライブラリの耐久性が高まり、将来のテクノロジー環境の変化にも適応しやすくなっています。

以降のセクションでは、JAX AI スタックの技術的な概要、主な機能、その背後にある設計上の決定事項、最新の ML ワークロード用に耐久性の高いプラットフォームを構築するためにこのスタックがどのように併用されているのかについて説明します。

JAX AI スタックとその他のエコシステム コンポーネント

コンポーネント 機能 / 説明
JAX AI スタックのコアとコンポーネント1
JAX アクセラレータ指向の配列計算とプログラム変換（JIT、grad、vmap、pmap）。
Flax 直感的なモデルの作成と変更が可能な柔軟なニューラル ネットワーク作成ライブラリ。
Optax 構成可能な勾配処理と最適化変換のライブラリ。
Orbax 桁外れのトレーニング復元力を実現する「任意のスケール」の分散チェックポインティング ライブラリ。
Grain スケーラブルかつ決定論的でチェックポインティング可能な入力データ パイプライン ライブラリ。
JAX AI スタック - インフラストラクチャ
XLA TPU、CPU、GPU 用のオープンソースの ML コンパイラ。
Pathways 数万個のチップにわたって計算をオーケストレートするための分散ランタイム。
JAX AI スタック - 高度な開発
Pallas Python で実装された低レベルの高性能カスタム カーネルを作成するための JAX 拡張機能。
Tokamax 最先端の高性能カスタム カーネル（Attention など）のキュレートされたライブラリ。
Qwix 量子化（PTQ、QAT、QLoRA）用の包括的で非侵入型のライブラリ。
JAX AI スタック - アプリケーション
MaxText / MaxDiffusion 基盤モデル（LLM や拡散など）のトレーニング用のスケーラブルなリファレンス フレームワーク。
Tunix 最先端のトレーニング後とアライメント（RLHF、DPO）用のフレームワーク。
vLLM vLLM フレームワークの組み込み統合を使用した高性能 LLM 推論ソリューション。
XProf システム全体のパフォーマンス分析のための、ハードウェア統合型の詳細なプロファイラ。

1jax-ai-stack Python パッケージに含まれています。

図 1: JAX AI スタックとエコシステムのコンポーネント

JAX AI スタック

アーキテクチャの必須事項: フレームワークを超えたパフォーマンス

モデル アーキテクチャが収束するにつれて（たとえば、マルチモーダル Mixture-of-Experts（MoE）Transformer など）、ピーク パフォーマンスの追求によりメガカーネルが登場しました。メガカーネルは、NVIDIA GPU の CUDA SDK などの下位レベルの API を使用して手動でコーディングされた、特定のモデルのフォワードパス全体（または大部分）になります。このアプローチでは、コンピューティング、メモリ、通信を積極的にオーバーラップさせることで、ハードウェアの使用率を最大化します。コミュニティの最近の研究では、このアプローチにより、GPU での推論スループットが大幅に向上することが示されています。場合によっては 22% を超える向上率が得られています。この傾向は推論に限定されません。大規模なトレーニングで、効率を大幅に向上させるために低レベルのハードウェア制御が使用されていることを示す証拠があります。

この傾向が加速すると、成熟した安定したアーキテクチャのパフォーマンスにとって最終的に重要なものはハードウェアへの低レベルアクセスであるため、現在のすべての高レベル フレームワークの重要性が低下する可能性があります。これは、高レベルのフレームワークの生産性と柔軟性を損なうことなく、エキスパート レベルのハードウェア制御をどのように提供するかという、すべての最新の ML スタックにとって課題を浮き彫りにしています。

TPU がこのレベルのパフォーマンスを実現するための明確なパスを提供するには、エコシステムでハードウェアに近い API レイヤを公開し、高度に特殊化されたカーネルの開発を可能にする必要があります。JAX スタックは、XLA コンパイラの自動化された高レベルの最適化から、Pallas カーネル作成ライブラリのきめ細かい手動制御まで、抽象化の連続体（図 2 を参照）を提供することで、この問題を解決するように設計されています。

図 2: JAX の抽象化の連続体

JAX の抽象化の連続体

コア JAX AI スタック

コア JAX AI スタックは、モデル開発の基盤となる 5 つの主要なライブラリで構成されています。

JAX: 構成可能で高パフォーマンスのプログラム変換の基盤

JAX は、アクセラレータ指向の配列計算とプログラム変換のための Python ライブラリで、高パフォーマンスの数値計算と大規模な ML 用に設計されています。関数型プログラミング モデルと NumPy のような API を備えた JAX は、高レベル ライブラリの確固たる基盤となります。

コンパイラ ファーストの設計により、JAX は XLA（XLA セクションを参照）を活用して、プログラム全体の積極的な分析、最適化、ハードウェア ターゲティングを行うことで、スケーラビリティを促進しています。JAX は関数型プログラミング（純粋関数など）を重視しているため、コアプログラム変換が扱いやすくなっています（重要なことに、構成可能になっています）。

これらのコア変換を組み合わせて、モデルサイズ、クラスタサイズ、ハードウェア タイプ全体でワークロードの高性能とスケーリングを実現できます。

  • jit: Python 関数を最適化された融合 XLA 実行可能ファイルへのジャストインタイム コンパイル。
  • grad: 自動微分。順方向モードと逆方向モード、高階導関数をサポートします。
  • vmap: 関数のロジックを変更することなく、シームレスなバッチ処理とデータ並列処理を可能にする自動ベクトル化。
  • pmap / shard_map: 複数のデバイス（TPU コアなど）にわたる自動並列化。分散トレーニングの基盤を形成します。

XLA の GSPMD（汎用 SPMD）モデルとのシームレスな統合により、JAX はコードの変更を最小限に抑えながら、大規模な TPU Pod 間で計算を自動的に並列化できます。ほとんどの場合、スケーリングには高レベルのシャーディング アノテーションのみが必要です。

Flax: 柔軟なニューラル ネットワークの作成

Flax は、モデル構築に対する直感的でオブジェクト指向のアプローチを提供することで、JAX でのニューラル ネットワークの作成、デバッグ、分析を簡素化します。JAX の関数型 API は強力ですが、PyTorch などのフレームワークに慣れているデベロッパー向けに、パフォーマンスを低下させず、より使い慣れたレイヤベースの抽象化を提供します。

この設計により、トレーニング済みモデル コンポーネントの変更や結合が簡素化されます。LoRA や量子化などの手法では、操作可能なモデル定義が必要です。Flax の NNX API は、Pythonic インターフェースを介してこれを提供します。NNX はモデルの状態をカプセル化し、ユーザーの認知負荷を軽減します。また、モデル階層のプログラムによるトラバーサルと変更を可能にします。

主な強み

  • 直感的なオブジェクト指向 API: モデルの構築を簡素化し、サブモジュールの置換や部分的な初期化などの高度なユースケースを可能にします。
  • Core JAX との一貫性: Flax は、JAX の関数型パラダイムと完全に互換性のあるリフト変換を提供し、JAX のパフォーマンスを最大限に引き出しながら、デベロッパーの使いやすさを向上させます。

Optax: 構成可能な勾配処理と最適化戦略

Optax は、JAX 用の勾配処理と最適化のライブラリです。これは、モデルビルダーに、ディープ ラーニング モデルなどのアプリケーションをトレーニングするために柔軟に再結合可能なビルディング ブロックを提供するように設計されています。これは、コア JAX ライブラリの機能を基盤として、ML モデルのトレーニングに使用できる損失関数、オプティマイザー関数、高性能ライブラリ、それに関連する手法を提供します。

目的

損失の計算と最小化は、ML モデルのトレーニングを可能にする中核となるものです。自動微分をサポートするコア JAX ライブラリは、モデルをトレーニングするための数値機能を提供しますが、一般的なオプティマイザー（RMSProp や Adam など）や損失（CrossEntropy や MSE など）の標準実装は提供しません。これらの関数を実装することは可能ですが（一部の上級デベロッパーはそうするでしょう）、オプティマイザーの実装にバグがあると、モデルの品質に関する問題の診断が難しくなります。Optax は、正確性とパフォーマンスがテストされたこれらのアルゴリズムの実装を提供します。ユーザーがこのような重要な部分を実装する必要はありません。

最適化理論の分野はいまだに研究領域ではありますが、トレーニングにおける中心的な役割を担っているため、本番環境の ML モデルのトレーニングに不可欠な要素となっています。この役割を果たすライブラリは、研究成果を迅速に反映できる柔軟性と、本番環境のモデル トレーニングで信頼できる堅牢性とパフォーマンスを備えている必要があります。また、標準方程式に一致する最先端のアルゴリズムのテスト済みの実装も提供する必要があります。Optax ライブラリは、モジュラー構成が可能なアーキテクチャと、正しい読み取り可能なコードを重視することで、これを実現するように設計されています。

デザイン

Optax は、読みやすく、十分にテストされ、効率的なコア アルゴリズムの実装を提供することで、研究速度と研究から本番環境への移行の両方を強化するように設計されています。Optax はディープ ラーニングのコンテキスト以外でも使用されますが、このコンテキストでは、JAX の哲学に沿って純粋関数型で実装された、よく知られた損失関数、最適化アルゴリズム、勾配変換のコレクションと見なすことができます。よく知られている損失関数とオプティマイザーのコレクションにより、ユーザーは簡単かつ確実に作業を始めることができます。

Optax のモジュール型アプローチにより、複数のオプティマイザーの連結を行い、その後に他の一般的な変換（勾配クリッピングなど）を適用し、MultiStep や Lookahead などの一般的な手法を使用してそれらをラップすることで、数行のコードで強力な最適化戦略を実現できます。柔軟なインターフェースにより、新しい最適化アルゴリズムを研究し、shampoo や muon などの強力な 2 次最適化手法を使用できます。

# Optax implementation of a RMSProp optimizer with a custom learning rate
#  schedule, gradient clipping and gradient accumulation.
optimizer = optax.chain(
  optax.clip_by_global_norm(GRADIENT_CLIP_VALUE),
  optax.rmsprop(learning_rate=optax.cosine_decay_schedule(init_value=lr,decay_steps=decay)),
  optax.apply_every(k=ACCUMULATION_STEPS)
)

# The same thing, in PyTorch
optimizer = optim.RMSprop(model_params, lr=LEARNING_RATE)
scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=TOTAL_STEPS)
for i, (inputs, targets) in enumerate(data_loader):
    # ... Training loop body ...
    if (i + 1) % ACCUMULATION_STEPS == 0:
        torch.nn.utils.clip_grad_norm_(model.parameters(), GRADIENT_CLIP_VALUE)
        optimizer.step()
        scheduler.step()
 optimizer.zero_grad()

前のコード スニペットは、カスタム学習率、勾配クリッピング、勾配累積を使用してオプティマイザーを設定する方法を示しています。

主な強み

  • 堅牢なライブラリ: 正確性と読みやすさに重点を置いた、損失、オプティマイザー、アルゴリズムの包括的なライブラリを提供します。
  • モジュール式のチェーン可能な変換: この柔軟な API を使用すると、トレーニング ループを変更することなく、強力で複雑な最適化戦略を宣言的に作成できます。
  • 機能的でスケーラブル: 純粋な関数型実装は、JAX の並列化メカニズム（pmap など）とシームレスに統合されるため、同じコードを使用して単一のホストから大規模なクラスタにスケーリングできます。

Orbax / TensorStore - 大規模分散チェックポイント処理

Orbax は、単一デバイスから大規模な分散トレーニングまで、あらゆる規模に対応するように設計された JAX 用のチェックポインティング ライブラリです。断片化されたチェックポインティングの実装を統合し、非同期チェックポインティングや多層チェックポインティングなどの重要なパフォーマンス機能をより多くのユーザーに提供することを目的としています。Orbax は、大規模なトレーニング ジョブに必要な復元力を実現し、チェックポイントを公開するための柔軟な形式を提供します。

システム状態全体をスナップショットする一般的なチェックポイントと復元システムとは異なり、Orbax を使用した ML チェックポインティングでは、トレーニング モデルの重み、オプティマイザーの状態、データローダの状態の再開に必要な情報のみが選択的に保持されます。この目標を絞り込んだアプローチにより、アクセラレータのダウンタイムを最小限に抑えることができます。Orbax は、I/O オペレーションと計算をオーバーラップさせることでこれを実現します。これは、大規模なワークロードにとって重要な機能です。アクセラレータのアイドル時間がデバイスからホストへのデータ転送の時間にまで短縮されます。次のトレーニング ステップと重複する場合もあるため、チェックポインティングはパフォーマンスの観点からほぼ無料になります。

Orbax は、配列データの効率的な並列読み取りと書き込みに TensorStore を使用します。Orbax API はこの複雑さを抽象化し、JAX のモデルの標準表現である PyTree を処理するためのユーザー フレンドリーなインターフェースを提供します。

主な強み

  • 広範囲で採用: 毎月数百万件のダウンロードがある Orbax は、ML アーティファクトを共有するための一般的なメディアとして機能します。
  • 複雑さを簡素化: Orbax は、非同期保存、アトミック性、ファイル システムの詳細など、分散チェックポインティングの複雑さを抽象化します。
  • 柔軟性: Orbax では、一般的なユースケース向けの API が用意されているだけでなく、特殊な要件を処理するようにワークフローをカスタマイズできます。
  • パフォーマンスとスケーラビリティ: 非同期チェックポインティング、効率的なストレージ形式（OCDBT）、インテリジェントなデータ読み込み戦略などの機能により、Orbax は数万のノードを含むトレーニング実行にスケーリングできます。

Grain: 決定論的でスケーラブルな入力データ パイプライン

Grain は、JAX モデルのトレーニングと評価用のデータを読み取って処理するための Python ライブラリです。柔軟性、高速性、決定論的であり、大規模なワークロードのトレーニングを成功させるために不可欠なチェックポインティングなど、高度な機能をサポートしています。一般的なデータ形式とストレージ バックエンドをサポートしており、ネイティブでサポートされていないユーザー固有の形式とバックエンドへのサポートを拡張するための柔軟な API も提供しています。Grain は主に JAX で動作するように設計されていますが、フレームワークに依存せず、JAX を実行する必要はありません。他のフレームワークでも使用できます。

目的

データ パイプラインはトレーニング インフラストラクチャの重要な部分を形成します。一般的な変換を効率的に表現できる柔軟性と、アクセラレータを常にビジー状態に保つことができる十分なパフォーマンスが必要です。また、複数のストレージ形式とバックエンドに対応できる必要があります。ステップ時間が長いため、大規模なモデルを大規模にトレーニングするには、通常のトレーニング ワークロードで必要とされる要件に加えて、主に決定論と再現性に関するデータ パイプラインの要件が必要になります2。Grain ライブラリは、これらのニーズに対応する柔軟なアーキテクチャで設計されています。

2PaLM 論文のセクション 5.1 で、グラデーション クリッピングを有効にしているにもかかわらず、損失の急増が非常に大きかったことが指摘されています。解決策として、問題のあるデータバッチを削除し、損失が急増する前のチェックポイントからトレーニングを再開することが記述されています。これは、完全に決定論的で再現可能なトレーニング ステップでのみ可能になります。

デザイン

最上位レベルで入力パイプラインを構成する方法は 2 つあります。1 つはデータワーカーの個別のクラスタとして構成する方法、もう 1 つはアクセラレータを駆動するホストにデータワーカーを配置する方法です。Grain はさまざまな理由から後者を選択しています。

アクセラレータは、通常はトレーニング ステップ中にアイドル状態になる強力なホストと組み合わされるため、入力データ パイプラインの実行に適しています。この実装には、入力とコンピューティング全体で一貫したシャーディングのビューを提供することで、データ シャーディングのビューを簡素化するという利点もあります。データワーカーをアクセラレータ ホストに配置すると、ホスト CPU が飽和するリスクがあるという見方もありますが、RPC を使用してコンピューティング負荷の高い変換を別のクラスタにオフロードできないというわけではありません3。

API 側では、複数のプロセスと柔軟な API をサポートする純粋な Python 実装により、Grain では、よく理解されている変換パラダイムに基づいてパイプライン ステージを構成することで、任意の複雑なデータ変換を実装できます。

Grain は、Parquet や TFDS などの他の一般的なデータ形式とともに、ArrayRecord や Bagz などの効率的なランダム アクセス データ形式もサポートしています。Grain には、ローカル ファイル システムからの読み取りと Cloud Storage からの読み取りのサポートがデフォルトで含まれています。一般的なストレージ形式とバックエンドのサポートに加えて、ストレージ レイヤへのクリーンな抽象化により、既存のデータソースのサポートを追加することも、Grain ライブラリと互換性を持たせるために既存のデータソースをラップすることも可能です。

3マルチモーダル データ パイプラインは、このように動作する必要があります。たとえば、画像と音声のトークナイザーは、独自のアクセラレータ上の独自のクラスタで実行されるモデル自体であり、入力パイプラインは RPC 呼び出しを行ってデータのサンプルをトークンのストリームに変換します。

主な強み

  • 決定論的なデータフィード: データワーカーをアクセラレータと同じ場所に配置し、安定したグローバル シャッフルとチェックポイント可能なイテレータと組み合わせることで、Orbax を使用してモデルの状態とデータ パイプラインの状態を整合性のあるスナップショットで一緒にチェックポインティングできます。これにより、トレーニング プロセスの決定論性が向上します。
  • 強力なデータ変換を可能にする柔軟な API: 柔軟で純粋な Python 変換 API を使用すると、入力処理パイプライン内で広範なデータ変換を実行できます。
  • 複数の形式とバックエンドの拡張可能なサポート: 拡張可能なデータソース API は、一般的なストレージ形式とバックエンドをサポートしており、新しい形式とバックエンドのサポートを追加できます。
  • 強力なデバッグ インターフェース: データ パイプラインの可視化ツールとデバッグモードを使用すると、データ パイプラインのパフォーマンスを内省、デバッグ、最適化できます。

拡張された JAX AI スタック

コアスタックに加えて、専門ライブラリの豊富なエコシステムが、エンドツーエンドの ML 開発に必要なインフラストラクチャ、高度なツール、アプリケーション レイヤ ソリューションを提供します。

基盤となるインフラストラクチャ: コンパイラとランタイム

XLA: ハードウェアに依存しないコンパイラ中心のエンジン

目的

XLA（Accelerated Linear Algebra）は Google のドメイン固有のコンパイラで、JAX に統合されており、TPU、CPU、GPU のハードウェア デバイスをサポートしています。XLA は、TPU、GPU、CPU をターゲットとするハードウェアに依存しないコード ジェネレータとして設計されました。

XLA コンパイラのコンパイラ ファーストの設計は、急速に進化する研究環境において持続的な優位性を生み出すための基本的なアーキテクチャ上の選択です。一方、他のエコシステムでの一般的なカーネル中心のアプローチでは、パフォーマンスのために手動で最適化されたライブラリに依存しています。これは、安定し、確立されたモデル アーキテクチャには非常に効果的ですが、イノベーションにはボトルネックとなります。新しい研究で新しいアーキテクチャが導入されると、エコシステムは、新しいカーネルが作成され、最適化されるまで待たなければなりません。しかし、コンパイラ中心の設計では新しいパターンに一般化できることが多く、最先端の研究に最初から高性能なパスを提供できます。

デザイン

XLA は、JAX がトレース プロセス中に生成する計算グラフをジャストインタイム（JIT）でコンパイルすることで機能します（たとえば、関数に @jax.jit デコレーターが付いている場合など）。

このコンパイルは、マルチステージ パイプラインに従って行われます。

  1. JAX 計算グラフ
  2. High-Level Optimizer（HLO）
  3. Low-Level Optimizer（LLO）
  4. ハードウェア コード
  • JAX グラフから HLO へ: JAX 計算グラフが XLA の HLO 表現に変換されます。この高レベルの最適化では、オペレーター フュージョンや効率的なメモリ管理など、強力なハードウェアに依存しない最適化が適用されます。StableHLO 言語は、このステージのバージョニングされた永続的インターフェースとして機能します。
  • HLO から LLO へ: 高レベルの最適化の後、ハードウェア固有のバックエンドが引き継ぎ、HLO 表現をマシン指向の LLO に変換します。
  • LLO からハードウェア コードへ: 最終的に LLO は効率性の高いマシンコードにコンパイルされます。TPU の場合、このコードはハードウェアに直接送信される Very Long Instruction Word（VLIW）パケットとしてバンドルされます。

スケーリングの場合、XLA の設計は並列処理を中心に構築されています。アルゴリズムを使用して、チップ上の行列乗算ユニット（MXU）を最大限に活用します。チップ間では、XLA は SPMD（単一プログラム、複数データ）を使用します。これは、すべてのデバイスで単一のプログラムを使用するコンパイラ ベースの並列化技術です。この強力なモデルは JAX API を介して公開され、高レベルのシャーディング アノテーションを使用してデータ、モデル、パイプラインの並列処理を管理できます。

より複雑な並列処理パターンでは、MPMD（複数プログラム、複数データ）も可能です。PartIR:MPMD などのライブラリを使用すると MPMD アノテーションも提供できます。

主な強み
  • コンパイル: 計算グラフのジャストインタイム コンパイルにより、メモリ レイアウト、バッファ割り当て、メモリ管理の最適化が可能になります。カーネルベースの手法などの場合、この負担はデベロッパーに課せられます。ほとんどの場合、XLA はデベロッパーの速度を損なうことなく優れたパフォーマンスを実現できます。
  • 並列処理: XLA は SPMD を使用して複数の形式の並列処理を実装しており、これは JAX レベルで公開されています。これにより、シャーディング戦略を表現し、数千個のチップにわたるモデルのテストとスケーラビリティを実現できます。

Pathways: 大規模な分散コンピューティング用の統合ランタイム

Pathways は、分散トレーニングと推論の抽象化を提供し、フォールト トレランスと復元が組み込まれています。ML 研究者は単一の強力なマシンを使用しているかのようにコーディングできます。

目的

大規模なモデルをトレーニングしてデプロイするには、数百から数千個のチップが必要です。これらのチップは、多数のラックとホストマシンに分散されています。トレーニング ジョブは、これらのチップとそれぞれのホストが、並列化（シャーディング）された XLA コンピューティングで連携して動作する必要がある大規模な同期プログラムです。大規模言語モデルの場合、数万個以上のチップが必要になる可能性があるため、このサービスは、Pod 内でチップ間相互接続（ICI）とオンチップ相互接続（OCI）ファブリックを使用するだけでなく、データセンター ファブリック内の複数の Pod にまたがって使用できるようにする必要があります。

デザイン

ML Pathways は、ホストと TPU チップ間で分散コンピューティングを調整するために使用するシステムです。このサービスは、数十万のアクセラレータにわたるスケーラビリティと効率性を実現するように設計されています。大規模なトレーニングでは、複数の Pod ジョブ用の単一の Python クライアント、Megascale XLA の統合、コンパイル サービス、リモート Python が提供されます。また、クロススライス並列処理とプリエンプション許容度もサポートしており、リソースのプリエンプションからの自動復旧が可能です。

Pathways には、XLA 計算グラフを単一の TPU Pod を超えて拡張できる最適化されたクロスホスト コレクティブが組み込まれています。XLA のデータ、モデル、パイプラインの並列処理のサポートを拡張し、データセンター ネットワーク（DCN）を使用して TPU スライス境界を越えて動作するようにします。これは、DCN 通信を XLA 通信プリミティブで管理する分散ランタイムを統合することで実現されます。

主な強み