Scaling LLM Workloads on TPUs with Ray: From Pre-Training to Inference
summary
Ray on GKE implements topology-aware scheduling that atomically reserves contiguous TPU slices (examples: 4-node 4x4 or 2x2x2 topologies) to avoid XLA deadlocks, adds dynamic sub-slicing to carve a 32-chip slice into smaller isolated partitions at runtime, and integrates TPU images, placement-group APIs, and preemption-aware checkpoints so JAX and PyTorch jobs (one process per host vs per device) and Ray Serve LLM inference scale across multi-slice TPU fleets with autoscaling and observability.
AI-written summary · about
chapters (14)
- 0:00Introduction to TPU scaling
- 2:55TPU slice fragmentation problem
- 4:20Native topology aware scheduling
- 7:06Ray core TPU integration
- 8:54Slice placement group API
- 10:58Dynamic sub-slicing in Ray
- 11:43Ray Serve LLM on TPUs
- 13:28Production deployment with Ray Service
- 14:54Ray Train integration with TPUs
- 17:21Multi-slice training setup
- 18:25Elastic training and preemption
- 20:52Data pipeline optimization
- 22:38Observability and diagnostics
- 23:00Future roadmap and next steps
related talks
26mIntroduction to Distributed ML Workloads with Ray on Kubernetes
23mAn Open Source Post-Training Stack: Kubernetes + Ray + PyTorch + vLLM
31mSpotify's Approach to Distributed LLM Training with Ray on GKE
31mRay at Scale: Apple's Approach to Elastic GPU Management
28mRay on Public-Cloud k8s - experiments, lessons & best practices
28mScaling AI & Machine Learning Workloads With Ray on AWS, Kubernetes, & BERT
22mRay, a Unified Distributed Framework for the Modern AI Stack
30mTopology-Aware Scheduling for AI Training & Inference With Kueue