[Train] JaxTrainer Implementation Tracking Issue

[ryanaoleary]

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

We can update this issue with more steps as milestones become clear.

Use case

Full support for SPMD workloads orchestrated with Ray and RayTrain.

[ryanaoleary]

cc: @Yevet