Ray、TPU slice スケジューリングを Serve、Data、Train に組み込み、マルチホスト AI ワークフローを統合

Ray 2.55 の TPU サポートは、通常の task や actor からアクセラレーターを認識できるようにするだけではない。Google が7月24日に発表した第2段階の統合では、TPU のトポロジー制約が Ray Serve、Ray Data、Ray Train に組み込まれた。TPU チップは固定された slice 内で ICI 高速インターコネクトを介して接続されているため、マルチホストのテンソル並列モデルは、必ず単一の slice 全体に配置されなければならない。スケジューラーがチップ単位のリソースを2つの slice に分散させると、最初の collective が永遠に完了せず、デプロイが `DEPLOYING` のまま停止する可能性がある。

Ray Serve では、`accelerator_config` に `topology: "4x4"` などを指定できるようになった。replica が slice placement group を作成し、トポロジー全体に対して gang scheduling を実行する。モデルサービングでは引き続き、vLLM のロードバランシング、スケーリング、複数モデルの構成機能を利用できる。Ray Data には `iter_jax_batches()` が追加され、JAX array への変換とデバイスシャーディングを完了したバッチを直接生成できる。これにより、NumPy から JAX へのホスト側コピーが削減されるほか、最後の不規則なバッチについて、破棄、パディング、エラーのいずれで処理するかを明示的に指定できる。

トレーニング向けの `JaxTrainer` は、slice の形状、ホストごとのワーカー数、slice 間の協調、checkpoint、フォールトトレランスを `ScalingConfig` に集約する。Google は同時に、JAX、Flax、Optax、Orbax、分析ツールを含む `rayproject/ray:*-tpu` イメージを提供する。Ray Dashboard でも TPU 使用率とメモリ使用量を表示できる。公開されたサンプルでは、単一の TPU v6e `2x4` slice、Qwen3-4B、LoRA/DPO を使用し、データ準備、ファインチューニング、バッチ推論、サービングまでを一連のワークフローとして構成している。

これにより、既存の Ray チームが TPU を採用する際のスケジューリングや環境構築のコストが下がり、同じ Python 抽象化を GPU と TPU のワークフローで共通利用できるようになる。ただし、「`topology` を追加するだけ」で TPU の静的形状、コンパイル時間、slice 間のネットワークコストが解消されるわけではない。また、公式の利用経路は現時点で GKE、KubeRay、Google 提供イメージに大きく依存している。エンジニアリングチームは今後、大規模モデルの実効スループット、Spot 障害からの復旧、vLLM の TPU 機能カバレッジ、GKE 以外の Kubernetes 環境におけるポータビリティを検証する必要がある。