Scaling PyTorch Training with Ray Data on Uber's Michelangelo AI Platform
summary
Uber’s Michelangelo migrates feature transforms and dataloaders to migrated from Spark to native PyTorch on Ray Data and replaced Petastorm with Ray Data, running a 4.5 billion‑row, 100+ TB job on 100 workers in production. Key wins: fused PyTorch transforms, zero‑copy Arrow I/O, a producer thread for background batch prep and an epoch‑dependent seed for reproducible shuffling. Optimizations yield ~15× throughput, 19× less CPU time and 3× less memory vs. the Spark baseline; some pipelines initially slow become up to 55% faster after tuning.
AI-written summary · about
chapters (15)
- 0:00Introduction and overview
- 1:24Training workflow evolution
- 3:17Native transformation migration
- 6:32Why Ray Data for preprocessing
- 7:57Native transformation workflow
- 10:29Performance and reliability
- 13:57Scheduling loop optimization
- 15:07Data loading migration
- 16:57Architecture of data loaders
- 17:51Migration strategy
- 19:57Shuffling and reproducibility
- 22:50Removing CPU bottlenecks
- 25:04Optimized performance results
- 26:01GPU utilization nuances
- 26:59Lessons and future work
related talks
33mUber's GenAI Leap: Batch Predictions Using Ray and vLLM
32mHow Uber Optimize Marketplaces with Ray
19mMichelangelo ML Platform at Uber: Past, Present and Future
21mHow Ray and Anyscale Make it Easy to do Massive-scale ML on Aerial Imagery
33mOvercoming Distributed ML Challenges with Ray Train
22mLarge Scale Data Loading and Data Preprocessing with Ray
44mPanel: ML Platform on Ray
28mScaling Multimodal Data and Reinforcement Learning with Ray