目录

1. Mesh Transformer JAX
   1. 更新
2. 预训练模型
   1. GPT-J-6B
      1. 链接
      2. 致谢
      3. 许可证
      4. 模型详情
      5. 零样本评估
3. 架构与使用
   1. 微调
   2. JAX依赖
4. 待办事项

Mesh Transformer JAX

一个基于 JAX 的俳句库，利用 xmap/pjit 操作符实现 Transformer 模型的模型并行。

该并行化方案类似于 原始的 Megatron-LM，由于其高速的二维网格网络，在 TPU 上表现高效。此外，还有一个实验性的模型版本实现了 ZeRo 风格的分片。

此库的设计目标是在 TPUv3 上实现约 400 亿参数规模的可扩展性，超过这一规模则应采用其他并行化策略。对于更大规模的模型，并行化，可以参考其他实现，如 GPT-NeoX 或 DeepSpeed 等。

未来的一个研究方向是将此代码库与 swarm-jax 集成，以通过流水线并行进一步提升可扩展性。

更新

2021年7月12日：新增微调指南

预训练模型

GPT-J-6B

一个拥有 60 亿参数的自回归文本生成模型，基于 The Pile 数据集训练而成。

链接

下载精简权重（仅 bf16 权重，用于推理，9GB)

下载完整权重（包含优化器参数，61GB)

部分训练检查点

Colab 示例

Web 示例

Aran 的博客文章

致谢

本项目离不开 TPU Research Cloud 提供的慷慨算力支持，同时感谢 EleutherAI 的协助。

特别感谢 Google Cloud TPU 团队，他们提供了 Cloud TPU VM 的早期访问权限（现已公开！）。

还要感谢所有以各种方式提供帮助的人士（按字母顺序排列）：

• Aran Komatsuzaki，在实验设计和撰写博客文章方面提供了宝贵建议。
• James Bradbury，在调试 JAX 问题时给予了重要帮助。
• Janko Prester，负责创建 Web 示例的前端界面。
• Laurence Golding，为 Web 示例添加了一些功能。
• Leo Gao，为表格中的基准模型运行了零样本评估。

许可证

GPT-J-6B 的权重采用 Apache License 2.0 版本进行授权。

模型详情

超参数	值
参数量	6,053,381,344
层数	28*
模型维度	4,096
前馈网络维度	16,384
注意力头数	16
每个头的维度	256
上下文长度	2,048
词汇表大小	50,257 （与 GPT-2/3 使用相同的分词器）
位置编码	旋转位置编码（RoPE）
RoPE 维度	64

* 每层由一个前馈块和一个自注意力块组成。

该模型共有 28 层，模型维度为 4096，前馈网络维度为 16384。模型维度被划分为 16 个注意力头，每个头的维度为 256。旋转位置编码（RoPE）应用于每个头的 64 个维度。模型使用 50257 个词汇的 BPE 分词器进行训练，与 GPT-2 和 GPT-3 使用的分词器相同。

零样本评估

模型大致按性能排序，若无性能数据则按训练浮点运算次数（FLOPs）排序。

模型	权重	训练 FLOPs	LAMBADA PPL ↓	LAMBADA Acc ↑	Winogrande ↑	Hellaswag ↑	PIQA ↑	数据集大小 (GB)
随机猜测	✔	0	~很多	~0%	50%	25%	25%	0
GPT-3-Ada‡	✘	-----	9.95	51.6%	52.9%	43.4%	70.5%	-----
GPT-2-1.5B	✔	-----	10.63	51.21%	59.4%	50.9%	70.8%	40
GPTNeo-1.3B‡	✔	3.0e21	7.50	57.2%	55.0%	48.9%	71.1%	825
Megatron-2.5B*	✘	2.4e21	-----	61.7%	-----	-----	-----	174
GPTNeo-2.7B‡	✔	6.8e21	5.63	62.2%	56.5%	55.8%	73.0%	825
GPT-3-1.3B*‡	✘	2.4e21	5.44	63.6%	58.7%	54.7%	75.1%	~800
GPT-3-Babbage‡	✘	-----	5.58	62.4%	59.0%	54.5%	75.5%	-----
Megatron-8.3B*	✘	7.8e21	-----	66.5%	-----	-----	-----	174
GPT-3-2.7B*‡	✘	4.8e21	4.60	67.1%	62.3%	62.8%	75.6%	~800
Megatron-11B†	✔	1.0e22	-----	-----	-----	-----	-----	161
GPT-J-6B‡	✔	1.5e22	3.99	69.7%	65.3%	66.1%	76.5%	825
GPT-3-6.7B*‡	✘	1.2e22	4.00	70.3%	64.5%	67.4%	78.0%	~800
GPT-3-Curie‡	✘	-----	4.00	69.3%	65.6%	68.5%	77.9%	-----
GPT-3-13B*‡	✘	2.3e22	3.56	72.5%	67.9%	70.9%	78.5%	~800
GPT-3-175B*‡	✘	3.1e23	3.00	76.2%	70.2%	78.9%	81.0%	~800
GPT-3-Davinci‡	✘	-----	3.0	75%	72%	78%	80%	-----
Gopher 230B*	✘	6.31E+23	-----	74.50%	70.10%	79.20%	81.80%	1344
MT-NLG 530B*‡	✘	-----	-----	76.6%	73.0%	80.2%	82.0%	-----

* 表示由各自作者报告的评估结果，其余数值均由运行 lm-evaluation-harness 工具得出，使用公开权重或 API 接口获取。由于实现细节及零样本任务设定的不同，这些结果可能无法直接比较。更多详情请参阅这篇博客文章。

† Megatron-11B 模型未提供可比指标，且使用其公开权重的多个实现未能复现生成质量和评估结果。（见 1 2 3）因此未进行评估。

‡ 这些模型所用训练数据可能存在测试集污染。OpenAI 的 GPT-3 模型未能针对某些测试集去重训练数据，而 GPT-Neo 模型以及本模型均基于 The Pile 数据集训练，该数据集尚未与任何测试集去重。

架构与使用

本仓库中的大多数脚本设计用于在 TPU 上运行。根据 TPU-VM 架构，TPU-VM 是能够运行任意代码的虚拟机。大多数脚本会启动一个 TPU，通过 SSH 连接到该 TPU 以设置依赖并从本地目录复制代码，随后启动一个 Ray 工作进程，用于接收远程过程调用。

TPU-VM 负责执行模型训练步骤和评估、检查点的保存与加载，而驱动 Python 程序则负责数据加载和整体编排工作（例如何时保存检查点等）。

这意味着大多数脚本（如 train.py、eval_harness.py 等）需要在同一区域的 GCE 虚拟机上运行，以便与 TPU 保持较近的距离，从而降低 RPC 延迟和数据传输成本。其他脚本（通常不带 --tpu 参数的脚本，如 device_sample.py、device_serve.py 或 device_train.py）则需直接在 TPU-VM 上运行。device_* 脚本仅适用于 v3-8 TPU，而不适用于更大规模的 Pod。

此外，还有一个示例脚本 (resharding_example.py) 展示如何将提供的检查点（以 GPT-J-6B 为例，包含 8 个分片）转换为更少的分片数，以便在 GPU 上运行。

微调

要对模型进行微调，请在 TPU VM 上运行 device_train.py。使用 TPU v3-8 时，您可以以约 5000 tokens/秒的速度进行微调，这足以处理中小型数据集。

请阅读逐步指南，以获取详细的微调说明。

JAX 依赖

请注意，此库对 JAX 版本有特定要求。具体而言，要使用 v1 版本的模型（包括 GPT-J 6B），必须安装 jax==0.2.12。这又依赖于 jaxlib==0.1.68。如果不满足此要求，您将遇到难以理解的 xmap 错误。

然而，对于 v2 版本的模型代码（目前尚未公开权重），可以使用最新版本的 JAX。

引用

引用本仓库：

@misc{mesh-transformer-jax,
  author = {Wang, Ben},
  title = {{Mesh-Transformer-JAX: Model-Parallel Implementation of Transformer Language Model with JAX}},
  howpublished = {\url{https://github.com/kingoflolz/mesh-transformer-jax}},
  year = 2021,
  month = May
}

引用 GPT-J-6B 的权重：

@misc{gpt-j,
  author = {Wang, Ben 和 Komatsuzaki, Aran},
  title = {{GPT-J-6B: A 6 Billion Parameter Autoregressive Language Model}},
  howpublished = {\url{https://github.com/kingoflolz/mesh-transformer-jax}},
  year = 2021,
  month = May
}

如果您使用本仓库或其中的预训练权重做出了有趣的事情，我们非常乐意了解！欢迎在 GitHub 上提交问题或通过个人资料中的邮箱与我们联系。

Mesh Transformer JAX 快速上手指南

Mesh Transformer JAX 是一个基于 JAX xmap/pjit 算子实现的 Haiku 库，专为在 TPU 上进行 Transformer 模型并行训练而设计。它支持类似 Megatron-LM 的并行策略，适用于扩展至约 400 亿参数的模型（如 GPT-J-6B）。

环境准备

系统要求

• 硬件: 推荐使用 Google Cloud TPU VM (特别是 v3-8 配置)。大部分脚本专为 TPU 架构优化。
  • 注意: device_* 系列脚本仅支持 v3-8 TPU，不支持更大的 Pod 配置。
• 操作系统: Linux (TPU VM 默认环境)。
• 网络: 若在本地运行驱动脚本连接远程 TPU，需确保与 TPU 位于同一 Google Cloud 区域以最小化延迟。

前置依赖

本库对 JAX 版本有严格要求，版本不匹配会导致 xmap 报错。

• 对于 v1 模型 (包括 GPT-J-6B):
  • jax==0.2.12
  • jaxlib==0.1.68
• 对于 v2 模型代码 (无公开权重):
  • 可使用最新版本的 JAX。

其他依赖包括 haiku, ray, numpy 等，通常包含在项目的 requirements.txt 中。

安装步骤

1. 克隆仓库

   git clone https://github.com/kingoflolz/mesh-transformer-jax.git
   cd mesh-transformer-jax
   

2. 安装 Python 依赖 由于国内访问 PyPI 可能较慢，建议配置国内镜像源（如清华源）进行安装。

   # 针对 GPT-J-6B (v1 模型) 的严格版本安装
   pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple
   
   # 手动锁定关键版本以防自动升级导致错误
   pip install jax==0.2.12 jaxlib==0.1.68 -i https://pypi.tuna.tsinghua.edu.cn/simple
   

3. 下载预训练权重 (可选) 如需使用 GPT-J-6B 进行推理或微调，可下载精简版权重（仅 bf16，用于推理，约 9GB）：

   wget https://the-eye.eu/public/AI/GPT-J-6B/step_383500_slim.tar.zstd
   # 解压文件 (需要 zstd 工具)
   tar --use-compress-program=unzstd -xf step_383500_slim.tar.zstd
   

基本使用

大多数脚本设计为在 TPU VM 上直接运行，或通过本地驱动程序远程调用。以下是在 TPU VM 上进行微调 (Fine-tuning) 的最简示例。

1. 准备数据与配置

确保你的数据集已预处理为模型可接受的格式，并修改配置文件（通常在 configs/ 目录下）以指向你的数据路径和检查点路径。

2. 运行微调脚本

在 TPU VM 终端中执行 device_train.py。以下命令假设你已配置好环境并拥有 v3-8 TPU：

python device_train.py \
    --tpu <your-tpu-name> \
    --config configs/gpt_j_6b_finetune.json \
    --initial_checkpoint_path /path/to/downloaded/checkpoint \
    --data_path /path/to/your/dataset

• --tpu: 指定 TPU 名称（如果在 TPU VM 内部运行且为单卡/单节点测试，部分脚本可能不需要此参数，具体视脚本逻辑而定，但标准流程通常涉及 Ray 集群设置）。
• 性能参考: 在 TPU v3-8 上，微调速度约为 5000 tokens/秒。

3. 推理示例 (Device Sample)

若要进行简单的文本生成测试，可使用 device_sample.py：

python device_sample.py \
    --tpu <your-tpu-name> \
    --config configs/gpt_j_6b_inference.json \
    --checkpoint_path /path/to/checkpoint \
    --prompt "Once upon a time" \
    --tokens_to_generate 100

提示: 详细的微调步骤、参数说明及高级用法（如显存分片转换以适应 GPU），请参阅项目根目录下的 howto_finetune.md 和 resharding_example.py。