#JAX

GoogleはRayのTPU対応について、Ray Serve、Ray Data、Ray Trainで推論・データ処理・分散学習を行う方法を公開しました。TPUスライスの配置をtopology指定で扱い、vLLMによるLLM提供とJAX学習を実行できます。

Ray Serveではaccelerator_config.topologyを指定し、複数ホストにまたがるモデルを同じTPUスライスへ配置してvLLMで提供できます。指定しない場合、ワーカーが別スライスに分散してデプロイが停止するため、実運用では重要な設定です。 Ray Dataのiter_jax_batches()は、デバイス分割済みのJAX配列を渡し、ホスト側コピーによる入力ボトルネックを抑えます。 Ray TrainのJaxTrainerは、TPUトポロジーを指定して分散学習、チェックポイント、障害回復、複数スライスへの拡張を扱えます。JAXは各ワーカー関数内でimportする必要があります。 Ray 2.55以降のTPU対応を前提とする開発者向け解説で、モデル/APIの新規提供や利用料金変更ではありません。 Google Cloud TPUで大規模推論・後学習・バッチ推論を行うチームは、手作業の配置制御を減らしつつ、RayService、公式TPUイメージ、ダッシュボード監視を組み合わせられます。

Googleは、397BパラメータのQwen 3.5 MoEをIronwood TPUで高速化するシステム最適化手法を公開しました。 JAX/Pallas、データ並列とエキスパート並列の組み合わせにより、prefill中心ワークロードで最大4.7倍の推論高速化を報告しています。

対象はGoogle自身のGemini APIやGoogle AI Studioで提供されるモデルではなく、Qwen 3.5-397B MoEをIronwood TPU7x上で提供・推論するための開発者向け技術解説です。 Data ParallelismとExpert Parallelismを組み合わせ、トークンルーティング向け通信融合やカスタム低レベルカーネルを実装したと説明されています。 Batched Ragged Page AttentionやGated DeltaNetブロックの最適化により、メモリ帯域とTensorCore MXUの利用率を高め、prefill中心の処理で最大4.7倍の高速化を達成したとしています。 モデルは397Bパラメータであり、ローカル実行、量子化、Ollama/LM Studio/llama.cpp/MLX、商用利用条件はこの記事では扱っていません。 APIの移行期限・破壊的変更・料金変更は記載されていません。大規模MoE推論基盤を設計するTPU/JAX開発者に関係する内容です。

Google Developers Blogが、MaxText、JAX、Pathways、Cloud TPUを使ったelastic trainingの解説を公開しました。分散学習中にTPUワーカーを意図的に停止しても、メインプロセスを再起動せず約2分以内に復旧する例を示しています。

Googleは、JAX AIスタックのMaxText、Pathways、Orbax、Cloud TPU、GKEを使い、分散LLM学習の途中でワーカー障害が起きても学習プロセスを生かしたまま復旧するelastic trainingを紹介しました。対象は大規模モデルをTPU上で分散学習するMLエンジニアや研究開発チームです。 通常の分散学習では、1台のマシンが消えるとall-reduceが停止し、全ワーカーが終了してジョブ全体を再起動する必要があります。記事の例では、Pathwaysが障害をPython例外として扱い、壊れたワーカーだけを置き換え、OrbaxがCloud Storage上のチェックポイントから復旧します。 実験では、TPUワーカーを意図的に停止した後、次の学習ステップまでの総ダウンタイムは2分未満で、その大半はKubernetesが代替Podをスケジュールする時間だったと説明されています。コントローラのPythonプロセスは同じPIDのまま再起動されません。 提供条件として、記事は再現方法を含む開発者向け解説です。対象API/ランタイムはMaxText、JAX、Pathways、Orbax、Cloud TPU、GKEで、モデル提供やGemini/Gemmaの新リリースではありません。破壊的変更や移行期限は示されていませんが、実運用ではチェックポイント、Pod置換、TPU/GKE構成の設計が重要になります。