Levanter: Legible, Scalable, Reproducible Foundation Models with JAX
summary
Levanter is a JAX framework built for legible, scalable, reproducible foundation-model training using Haliax named tensors to eliminate shape-broadcast bugs and separate model logic from parallelism. It adds FSDP in 10 lines (data, parameter, and optimizer sharding) while achieving 50%+ MFU on a v3-256 TPU. Runs are bitwise reproducible across preempted restarts. Implementations include Sophia (a Hessian-approx optimizer scaled to 6.7B and ~2x fewer steps vs Adam), FlashAttention2 to 32k sequences, Hugging Face model compatibility, and TPU/GPU support.
AI-written summary · about
chapters (12)
- 0:00Introduction to Levanter
- 3:14Academic role in foundation models
- 8:11Legibility in tensor code
- 11:25Scaling with named axes
- 13:14Performance benchmarks
- 14:18Reproducibility challenges
- 15:21Bitwise reproducibility in Levanter
- 16:05Sophia optimizer breakthrough
- 18:14Controllable music generation
- 19:19New model features and support
- 20:42Lumia medical AI project
- 21:45Q&A session
related talks
08mDemo: JAX, Flax and Gemma
01h 03mScaling up machine learning without tears (and what do programming languages ha...
14mJAX: Accelerated Machine Learning Research via Composable Function Transformations in Python
11mIntro to JAX: Accelerating Machine Learning research
31mHigh performance machine learning with JAX
01h 00mFedOpt for LLMs
26mTraining Large Deep Learning Models: PyTorch vs JAX
27mFoundation Models: Blueprints for success