TorchTPU: Expanding TPU Programmabil...
summary
TorchTPU brings PyTorch to Google TPUs with native eager-first and compiled paths, lowering Aten ops to StableHLO via OpenXLA and integrating with torch.compile, DTensor, and Palace/Helium kernels. Out-of-the-box work shows ~70% of thousands of Hugging Face models run without changes, a Llama 3 8B fine-tune runs by changing a few lines, and training scales linearly up to 256-chip linear scaling on Ironwood (TPU v7). Key runtime is TorchTPU with a defer-and-fuse eager mode to fuse ops and reduce the gap with compiled performance.
AI-written summary · about
chapters (16)
- 0:00TorchTPU introduction
- 1:46TPUs and PyTorch need
- 3:11TorchTPU and PyTorch integration
- 4:33Porting PyTorch models to TPUs
- 6:18Demo: Llama 3 8B on TPUs
- 8:26TPU hardware overview
- 10:35TorchTPU eager modes
- 12:18Compilation cache and torch.compile
- 13:43XLA compiler powers TorchTPU
- 14:44Kernel integration: Palace and Helium
- 16:09Distributed training on TPUs
- 17:54Hugging Face model validation
- 19:39Performance benchmarks
- 20:41Optimizing models for TPUs
- 21:43TorchTPU challenges and plans
- 22:05Availability and next steps
related talks
26mEverything Everywhere all at Once: vLLM...
42mCloud TPU Pods: AI Supercomputing for Large Machine Learning Problems
11mUnlock the Future of Generative AI: TorchTitan's Latest...
23mHow Pytorch Powers Training Inference
27mTips and tricks for distributed large model training
28mScaling LLM Workloads on TPUs with Ray: From Pre-Training to Inference
23mHugging Face Accelerate: Making Device-Agnostic ML Training and Inference Easy...
09mTogether Goes Brrr: Threading Research & Production with Torch Compile