EasyLM - LLM Framework

[AaronWard]

EasyLM is a framework for pre-training, fine-tuning, and evaluating and serving large language models in JAX/Flax. EasyLM is designed to be easy to use by hiding the complexities of distributed model/data parallelism but exposing the core training and inference details of large language models, making it easy to customize. EasyLM can scale up LLM training to hundreds of TPU/GPU accelerators without the need to write complicated distributed training code.