EasyLM

Why we included this project

EasyLM takes care of the awkward parts of distributed training so you can work with models that would never fit on one accelerator. It builds on JAX/Flax and uses pjit to shard weights and data across many TPUs or GPUs, handling single-host setups and multi-host runs on Google Cloud TPU Pods. The LLaMA family is supported directly, and the code sits on top of Hugging Face's transformers and datasets libraries. It's also the framework behind OpenLLaMA and the Koala chatbot, so it's a good place to look if you want to see how those models were actually trained. People already comfortable with JAX get an unusually direct route into large-scale LLM work without switching ecosystems.