EasyLM

让大型语言模型（LLMs）变得简单，EasyLM 是一个一站式解决方案，用于在 JAX/Flax 中对 LLMs 进行预训练、微调、评估和部署。EasyLM 可以利用 JAX 的 pjit 功能，将 LLMs 训练扩展到数百个 TPU/GPU 加速器。

基于 Hugginface 的 transformers 和 datasets，本仓库提供了一个易于使用和易于定制的代码库，用于训练大型语言模型， 避免了其他许多框架中的复杂性。

EasyLM 基于 JAX/Flax 构建。通过利用 JAX 的 pjit 工具，EasyLM 能够 通过跨多个加速器分片模型权重和训练数据，来训练无法在单个加速器上容纳的大型模型。目前， EasyLM 支持单主机上的多 TPU/GPU 训练，以及 Google Cloud TPU Pods 上的多主机 训练。

目前，支持以下模型：

• LLaMA
• LLaMA 2
• LLaMA 3

Discord Server

我们正在运营一个非官方的 Discord 社区（与 Google 无关），用于讨论在 JAX 中训练 LLMs 的相关内容。点击此链接加入 Discord 服务器。我们为几个基于 JAX 的 LLM 框架设置了专用频道，包括 EasyLM、JaxSeq、Alpa 和 Levanter。

使用 EasyLM 训练的模型

OpenLLaMA

OpenLLaMA 是我们采用宽松许可证重新实现的 LLaMA，可用于 商业用途。请访问项目主页。 OpenLLaMA 可以作为 EasyLM 中 LLaMA 权重的直接替代品。 请参阅LLaMA 文档以获取更多详细信息。

Koala

Koala 是我们基于 LLaMA 微调的新聊天机器人。如果您对 我们的 Koala 聊天机器人感兴趣，可以查看博客文章 以及本地运行文档。

安装

GPU 主机和 Cloud TPU 主机的安装方法有所不同。第一步 是从 GitHub 拉取代码。

git clone https://github.com/young-geng/EasyLM.git
cd EasyLM
export PYTHONPATH="${PWD}:$PYTHONPATH"

在 GPU 主机上安装

GPU 环境可通过 Anaconda 安装。

conda env create -f scripts/gpu_environment.yml
conda activate EasyLM

在 Cloud TPU 主机上安装

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}
}

致谢

• LLaMA 的实现来自 JAX_llama
• JAX/Flax GPT-J 和 RoBERTa 的实现来自 transformers
• 大部分 JAX 工具来自 mlxu
• 代码库深受 JAXSeq 的启发