Introduction to EasyLM
EasyLM simplifies the handling of large language models (LLMs) by providing a comprehensive solution for pre-training, fine-tuning, evaluation, and deployment. Designed for use with JAX/Flax, EasyLM efficiently scales the training of LLMs across hundreds of TPU/GPU accelerators using JAX's pjit functionality.
Built on Strong Foundations
This project builds upon the capabilities of Hugging Face's transformers and datasets. It provides an easy-to-use codebase that is fully customizable for training large language models, removing the complexity found in many other frameworks.
Scalable and Flexible Training
EasyLM leverages JAX/Flax for its implementations. By utilizing JAX's pjit utility, it allows users to train large models that exceed the capacity of a single accelerator by distributing the model weights and data across several accelerators. It supports both multi-TPU/GPU training on a single host and multi-host training on Google Cloud TPU Pods.
Supported Models
EasyLM currently supports several significant models:
Community Interaction
For those interested in discussions about training LLMs with JAX, an unofficial Discord community is available. It covers various JAX-based LLM frameworks including EasyLM, JaxSeq, Alpa, and Levanter. Join the Discord server to connect with others.
Models Trained with EasyLM
OpenLLaMA
OpenLLaMA is a reproduction of LLaMA with a permissive license allowing commercial use. It can be used as a direct replacement for LLaMA weights in EasyLM. More details can be found on the project main page.
Koala
Koala is a newly developed chatbot that builds upon LLaMA. For those interested, a comprehensive blogpost and documentation for running it locally are available.
Installation Guide
The installation procedure varies slightly depending on whether EasyLM is being installed on a GPU host or a Cloud TPU host.
Installation on a GPU Host
For a GPU host, users can employ Anaconda for setting up the environment:
git clone https://github.com/young-geng/EasyLM.git
cd EasyLM
export PYTHONPATH="${PWD}:$PYTHONPATH"
conda env create -f scripts/gpu_environment.yml
conda activate EasyLM
Installation on a Cloud TPU Host
For Cloud TPU hosts, the VM already has Python and PIP pre-installed. Users can run a setup script to configure the host:
./scripts/tpu_vm_setup.sh
Comprehensive Documentation
Detailed documentation for EasyLM can be found in the docs directory, providing users with the necessary resources to maximize the utility of the framework.
References and Credits
Researchers and developers who find EasyLM useful are encouraged to cite it using the provided BibTeX. The project's infrastructure draws from several influential JAX-based resources, including JAX_llama for LLaMA implementations, transformers for GTP-J and RoBERTa models, and mlxu for JAX utilities. The EasyLM codebase is also inspired by JAXSeq.