EasyLM 使大型语言模型 (LLM) 变得简单,是在 JAX/Flax 中预训练、微调、评估和服务 LLM 的一站式解决方案。 EasyLM 可以利用 JAX 的 pjit 功能将 LLM 训练扩展到数百个 TPU/GPU 加速器。
该存储库建立在 Hugginface 的转换器和数据集之上,提供了一个易于使用且易于自定义的代码库,用于训练大型语言模型,而无需许多其他框架的复杂性。
EasyLM 使用 JAX/Flax 构建。通过利用 JAX 的 pjit 实用程序,EasyLM 能够通过在多个加速器上分割模型权重和训练数据来训练不适合单个加速器的大型模型。目前,EasyLM 支持在单个主机中进行多个 TPU/GPU 训练,以及在 Google Cloud TPU Pod 上进行多主机训练。
目前支持以下型号:
我们正在运行一个非官方的 Discord 社区(与 Google 无关),用于讨论与 JAX 培训法学硕士相关的问题。点击此链接加入 Discord 服务器。我们为多个基于 JAX 的 LLM 框架提供专用频道,包括 EasyLM、JaxSeq、Alpa 和 Levanter。
OpenLLaMA 是我们许可的 LLaMA 复制品,可用于商业目的。在这里查看项目主页。 OpenLLaMA 可以替代 EasyLM 中的 LLaMA 权重。请参阅 LLaMA 文档了解更多详细信息。
Koala 是我们在 LLaMA 之上进行微调的新聊天机器人。如果您对我们的 Koala 聊天机器人感兴趣,您可以查看有关在本地运行它的博客文章和文档。
GPU 主机和 Cloud TPU 主机的安装方法有所不同。第一步是从 GitHub 拉取。
git clone https://github.com/young-geng/EasyLM.git
cd EasyLM
export PYTHONPATH= " ${PWD} : $PYTHONPATH "
GPU环境可以通过Anaconda安装。
conda env create -f scripts/gpu_environment.yml
conda activate EasyLM
TPU 主机虚拟机预装了 Python 和 PIP。只需运行以下脚本即可设置 TPU 主机。
./scripts/tpu_vm_setup.sh
EasyLM 文档可以在 docs 目录中找到。
如果您发现 EasyLM 在您的研究或应用中有用,请使用以下 BibTeX 进行引用:
@software{geng2023easylm,
author = {Geng, Xinyang},
title = {EasyLM: A Simple And Scalable Training Framework for Large Language Models},
month = March,
year = 2023,
url = {https://github.com/young-geng/EasyLM}
}