Scaling up machine learning without tears (and what do programming languages ha...
summary
Concrete tooling that turns JAX-traced array programs into deployable sharding plans and fast estimates: a small DSL plus a simulator that reproduces published trends, an automated partitioner called part IR with a tile registry, and an SPMD lowering that emits distributed types and collectives. Experiments include an 18B Transformer on 256 GPUs / 80→6 GB memory sweep showing different optimal shardings and that pipelining can beat naive strategies, and a simulator that matches external cluster results (within model error).
AI-written summary · about
chapters (10)
- 0:00Introduction
- 2:16What is ML program
- 10:48Parallelizing ML challenges
- 13:20Composing strategies with meshes
- 18:15Modeling to execution with PART
- 23:56PART for automatic scaling
- 32:25Compiler action spaces
- 41:57MLIR and dialect lessons
- 47:21Future directions
- 51:36Conclusion & Q&A
related talks
27mTips and tricks for distributed large model training
21mBeyond the Theory: What Actually Breaks When You Scale Your Disaggregat...
34mPanel Discussion: Training and Inference at Planet Scale
14mLightning Talk: Large-Scale Distributed Training with Dynamo and...
11mEngineering Scalable ML Infrastructure for Modern Apps
13mFrom Silicon to Scale: Training at Scale
27mScaling for Intelligence: Distributed SQL Databases in the Age of AI
26mEfficiently serving LLMs at scale.