Stanford’s Marin foundation model: The first fully open model developed using JAX- Google Developers Blog

This post details the Marin project's use of JAX and the custom Levanter framework to achieve end-to-end reproducibility for foundation models. It highlights specific engineering challenges and solutions: eliminating Python interpreter overhead with `@jax.jit` and fused operations, using `jax.value_and_grad` and gradient checkpointing for efficiency, and employing Pallas Splash Attention. It also covers managing large-scale parallelism with SPMD and Haliax for named tensors, building resilient compute clusters with Google Cloud TPU Multislice and Ray, and ensuring bitwise determinism through JAX's PRNGs and Tensorstore for data loading. The training process itself, including hardware migrations and adaptive data mixtures, is presented as a demonstration of the stack's robustness.