Description

This issue will serve as an implementation tracker for a JaxTrainer to support jax and SPMD workloads. The initial support for this framework will target SPMD with multi-host TPUs on Kubernetes.

Milestone 1: MVP of JaxTrainer with single-slice multi-host TPUs

• Add default TPU info to Ray node labels
  • [Core] Add default TPU Ray node labels #55064
• Add API change to ScalingConfig to specify topology and accelerator arguments
• Add JaxTrainer wrapper of DataParallelTrainer to RayTrain with SPMD scheduling support
  • [Train] Implement a JaxTrainer to support SPMD with TPUs #55207
• Extensively test Jax training workload with multi-host TPUs and Anyscale and KubeRay operators
• Add documentation and guides to RayTrain docs

Milestone 2: Full support for TPU multi-slice

• Implemented in: [Train] Add TPU multi-slice support to JaxTrainer #58629

Milestone 3: Elastic training with multi-slice TPUs in Ray Train

• Detect Alive TPU slices: [Core] Add TPU util to determine number of ready multi-host slices #61300
• Update elastic policy to scale for TPU slices: [Train] Update elastic policy to handle multi-host TPUs with JaxTrainer #61299

Use case

Full support for SPMD workloads orchestrated with Ray and RayTrain.