AI 基礎設施／分散式運算
Ray Pushes TPU Slice Scheduling Down into Serve, Data, and Train, Unifying Multi-Host AI Workflows
Google has unveiled a complete AI library path for running Ray on TPUs, using a topology field to ensure model workers are placed within the same ICI slice. The new release also adds native JAX data batching, JaxTrainer, official containers, and TPU monitoring, though GKE remains the primary validated environment.
Ray 2.55’s TPU support does more than expose accelerators to general-purpose tasks and actors. The second phase of the integration, announced by Google on July 24, brings TPU topology constraints into Ray Serve, Ray Data, and Ray Train. TPU chips are connected in fixed slices through the high-speed ICI interconnect, so multi-host tensor-parallel models must be placed entirely within the same slice. If the scheduler distributes per-chip resources across two slices, the first collective operation may wait indefinitely, leaving the deployment stuck in `DEPLOYING`.
Ray Serve can now specify settings such as `topology: "4x4"` in `accelerator_config`. Each replica creates a slice placement group and applies gang scheduling across the entire topology. Model serving can still use vLLM for load balancing, scaling, and multi-model composition. Ray Data adds `iter_jax_batches()`, which directly produces batches that have been converted to JAX arrays and sharded across devices. This reduces host-side NumPy-to-JAX copying and explicitly supports policies to drop, pad, or raise an error for the final irregular batch.
On the training side, `JaxTrainer` packages the slice shape, workers per host, cross-slice coordination, checkpointing, and fault tolerance into `ScalingConfig`. Google is also providing `rayproject/ray:*-tpu` images containing JAX, Flax, Optax, Orbax, and profiling tools. Ray Dashboard can display TPU utilization and memory usage as well. A public example uses a single TPU v6e `2x4` slice, Qwen3-4B, and LoRA/DPO to connect data preparation, fine-tuning, batch inference, and serving in one workflow.
This reduces the scheduling and environment-assembly costs for existing Ray teams moving to TPUs, while allowing the same Python abstractions to span GPU and TPU workflows. However, simply adding `topology` does not eliminate TPU constraints around static shapes, compilation time, or cross-slice networking costs. The official path also currently depends heavily on GKE, KubeRay, and Google-provided images. Engineering teams should next validate real-world throughput for large models, recovery from Spot VM failures, vLLM’s TPU feature coverage, and portability to Kubernetes environments outside GKE.