SBTB 2023: David Hall, Levanter: Legible, Scalable, Reproducible Foundation Models with JAX

ai.bythebay.io Nov 2025, Oakland, full-stack AI conference In this talk, I'll describe Levanter, a new JAX framework for training foundation models that we developed at the Stanford Center for Research on Foundation Models (CRFM). We designed Levanter to be legible, scalable, and reproducible. Levanter uses our new named tensor library Haliax to improve code legibility and enable flexible model parallelism: just 10 lines of code add FSDP and/or tensor parallelism without modifying model code at all. This legibility and flexibility does not come at the expense of efficiency: Levanter can achieve in excess of 50%+ model flop utilization on a v3-256 TPU. Thanks to JAX, Levanter also offers bitwise reproducibility. I'll also briefly touch on some of the other features, including Hugging Face compatibility and our online data preprocessing. We'll also briefly describe some of the research that Levanter has enabled or accelerated, including Anticipatory Music Transformers, the new optimizer Sophia, and Backpack architectures. More details available here: https://www.scale.bythebay.io/post/david-hall-levanter-legible-scalable-reproducible-foundation-models-with-jax