RFC: Multi-Controller JAX (McJAX) Support for Tunix

[wang2yn84]

Status: Draft for community feedback
Authors: Lance Wang @wang2yn84

---

1. Motivation

Tunix currently leverages Pathways on Cloud for multi-host training, which offers robust distributed runtime and TPU support. However, there are strong reasons to add Multi-Controller JAX (McJAX) as an alternative multi-host backend:

1. TRC Compatibility: Google’s TRC (Tensor Research Cloud) clusters do not support GKE/Pathways setups, but do support McJAX.
2. Broader Familiarity: Many open-source and academic users already use McJAX for distributed runs (via SSH or Slurm launchers).
3. Open Availability: Pathways is not yet fully open-sourced, , limiting reproducibility for non-Google users.

Adding McJAX support enables Tunix to run seamlessly on TPUs, GPUs, and CPUs in multi-host environments without relying on GKE or proprietary infra.

---

2. Goals & Non-Goals

Goals

• Provide multi-host training for Tunix using McJAX, supporting both TPU and GPU platforms.
• Maintain API compatibility with existing Pathways code (resharding API).
• Offer launcher scripts for SSH and Slurm clusters (e.g., TRC).
• Support fault-tolerant checkpointing via Orbax for McJAX topologies.
• Ensure interoperability — Pathways and McJAX can co-exist via environment flag.

Non-Goals

• Re-implement Pathways runtime features such as autoscaling or elastic jobs.
• Guarantee full feature parity across all hardware configurations at day one.

---

3. Background

• Pathways provides a cloud-native single-controller runtime, integrating with GKE and TPU VM topologies. It’s widely used internally but has limited public availability.
• McJAX is a simpler, open alternative for distributed JAX training across multiple controllers, communicating over gRPC or RDMA or NCCL.

---

4 Implementation Phases

• **M0 | Identify the suitable library for distributed system handling
• **M1 | McJAX env prototype on single host | Launch single controller successfully
• **M2 | 2-host TPU v5e-8 | Verify distributed Jax works
• **M3 | End-to-end SFT training
• **M4 | Multi-host GRPO | Verify reward aggregation & gradients
• **M5 | Orbax distributed checkpointing | Reliable recovery
• **M6 | Public Docker + Docs | Release setup guide for TRC and Slurm