EasyLM

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

基于 Hugging Face 的 transformers 和 datasets，这个仓库提供了一个易于使用且易于定制的代码库，用于训练大型语言模型，而无需像其他许多框架那样面对复杂的配置。

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

目前支持以下模型：

• LLaMA
• LLaMA 2
• LLaMA 3

Discord 社区

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

使用 EasyLM 训练的模型

OpenLLaMA

OpenLLaMA 是我们对 LLaMA 的宽松许可复现版本，可用于商业用途。请访问 项目主页 了解更多信息。OpenLLaMA 可以直接替换 EasyLM 中的 LLaMA 权重。更多详情请参阅 LLaMA 文档。

Koala

Koala 是我们在 LLaMA 基础上微调的新款聊天机器人。如果您对我们的 Koala 聊天机器人感兴趣，可以查看 博客文章 和 本地运行指南。

安装

GPU 主机和云 TPU 主机的安装方法有所不同。首先需要从 GitHub 克隆代码库。

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

在 GPU 主机上安装

可以通过 Anaconda 安装 GPU 环境。

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

在云 TPU 主机上安装

TPU 主机虚拟机已预装 Python 和 PIP。只需运行以下脚本即可完成 TPU 主机的设置。

./scripts/tpu_vm_setup.sh

文档

EasyLM 的文档位于 docs 目录中。

参考文献

如果您在研究或应用中使用了 EasyLM，请使用以下 BibTeX 格式引用：

@software{geng2023easylm,
  author = {Geng, Xinyang},
  title = {EasyLM: 一个简单且可扩展的大型语言模型训练框架},
  month = March,
  year = 2023,
  url = {https://github.com/young-geng/EasyLM}
}

致谢

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

EasyLM 快速上手指南

EasyLM 是一个基于 JAX/Flax 的一站式大语言模型（LLM）解决方案，支持预训练、微调、评估和部署。它利用 JAX 的 pjit 功能，可将训练规模扩展至数百个 TPU/GPU 加速器，同时保持代码简洁易用。目前支持 LLaMA、LLaMA 2 和 LLaMA 3 等主流模型。

环境准备

系统要求

• 硬件：支持单主机多卡（GPU/TPU）或多主机集群（Google Cloud TPU Pods）。
• 操作系统：Linux 推荐（Ubuntu/CentOS 等）。
• 软件依赖：
  • Python 3.8+
  • Git
  • Anaconda（仅 GPU 环境需要）
  • Google Cloud TPU VM（仅 TPU 环境需要，已预装 Python 和 Pip）

前置依赖

• 熟悉 JAX/Flax 基础概念（可选，框架已封装大部分复杂度）。
• 拥有 Hugging Face 账号以访问部分模型权重（如 LLaMA 系列）。

安装步骤

1. 克隆项目

首先从 GitHub 拉取代码并设置环境变量：

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

提示：国内用户若克隆速度慢，可尝试使用镜像源或配置 Git 代理。

2. 环境配置

方案 A：GPU 主机安装

推荐使用 Anaconda 管理环境：

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

方案 B：Cloud TPU 主机安装

TPU 虚拟机已预装基础环境，直接运行设置脚本：

./scripts/tpu_vm_setup.sh

基本使用

EasyLM 的核心优势在于简化了大规模模型的训练流程。以下以最基础的模型微调为例展示使用逻辑（具体配置文件需参考 docs/ 目录下的详细文档）。

典型工作流

1. 准备数据：使用 Hugging Face datasets 加载或预处理数据集。
2. 配置任务：修改 YAML 配置文件，指定模型路径（如 LLaMA）、数据集路径及训练参数。
3. 启动训练：使用 JAX 分布式接口启动任务。

示例命令结构

虽然具体参数依任务而定，但启动训练的基本命令格式如下：

# 示例：运行微调脚本（需根据实际配置文件调整参数）
python EasyLM/scripts/train.py \
    --config_path=path/to/your/config.yaml \
    --mesh_dim="1,-1,1" \
    --dtype=bf16

支持的模型

在配置文件中指定以下任一模型即可开始：

• llama (LLaMA)
• llama2 (LLaMA 2)
• llama3 (LLaMA 3)

注意：对于 OpenLLaMA 或 Koala 等衍生模型，可直接替换权重路径使用，无需修改代码架构。详细参数说明请参阅项目 docs/ 目录中的 llama.md 或 koala.md。