OpenLLaMA: LLaMA 的开源复现

TL;DR: 我们发布了 OpenLLaMA 的公开预览版，这是 Meta AI LLaMA 的一个采用宽松许可证的开源复现版本。我们发布了一系列在混合数据上训练的 3B、7B 和 13B 模型。我们的模型权重可以作为现有实现中 LLaMA 的直接替换。

在本仓库中，我们展示了 Meta AI 的 LLaMA 大语言模型的宽松许可证开源复现版本。我们发布了一系列在 1T token 上训练的 3B、7B 和 13B 模型。我们提供了预训练 OpenLLaMA 模型的 PyTorch 和 JAX 权重，以及评估结果和与原始 LLaMA 模型的对比。v2 模型优于在混合数据上训练的旧版 v1 模型。

Hugging Face transformers 的 PyTorch 权重：

• v2 模型
  • OpenLLaMA 3Bv2
  • OpenLLaMA 7Bv2
• v1 模型
  • OpenLLaMA 3B
  • OpenLLaMA 7B
  • OpenLLaMA 13B

EasyLM 的 JAX 权重：

• v2 模型
  • OpenLLaMA 3Bv2 for EasyLM
  • OpenLLaMA 7Bv2 for EasyLM
• v1 模型
  • OpenLLaMA 3B for EasyLM
  • OpenLLaMA 7B for EasyLM
  • OpenLLaMA 13B for EasyLM

更新

2023/07/15

我们发布了 OpenLLaMA 3Bv3 模型，这是一个在 1T token 上训练的 3B 模型，使用的数据集混合与 7Bv2 模型相同。

07/07/2023

我们很高兴发布 OpenLLaMA 7Bv2 模型，该模型基于 Falcon refined-web dataset、starcoder dataset 以及来自 RedPajama 的 wikipedia、arxiv、books 和 stackexchange 的混合数据训练而成。

06/15/2023

我们很高兴发布 OpenLLaMA 13B 的最终 1T token 版本。我们已更新评估结果。对于当前版本的 OpenLLaMA 模型，我们的 tokenizer 在分词前会将多个空字符合并为一个，类似于 T5 tokenizer。因此，我们的 tokenizer 不适用于代码生成任务（例如 HumanEval），因为代码中包含许多空字符。对于代码相关任务，请使用 v2 模型。

06/07/2023

我们很高兴发布 OpenLLaMA 3B 和 7B 的最终 1T token 版本。我们已更新评估结果。我们还很高兴发布 13B 模型的 600B token 预览版，该模型是与 Stability AI 合作训练的。

05/22/2023

我们很高兴发布 OpenLLaMA 7B 模型的 700B token 检查点和 3B 模型的 600B token 检查点。我们还更新了评估结果。我们预计完整的 1T token 训练运行将在本周结束时完成。

05/15/2023

在收到社区反馈后，我们发现之前发布的 checkpoint 的 tokenizer 配置不正确，导致换行符未被保留。为了解决这个问题，我们重新训练了 tokenizer 并重启了模型训练。我们还观察到使用这个新的 tokenizer 时训练损失更低。

权重发布、许可证和使用

我们以两种格式发布权重：一种用于我们的 EasyLM 框架 的 EasyLM 格式，以及一种用于 Hugging Face transformers 库的 PyTorch 格式。我们的训练框架 EasyLM 和 checkpoint 权重均根据 Apache 2.0 许可证以宽松的方式授权。

使用 Hugging Face Transformers 加载权重

预览 checkpoint 可以直接从 Hugging Face Hub 加载。请注意，目前建议避免使用 Hugging Face fast tokenizer，因为我们观察到 自动转换的 fast tokenizer 有时会产生错误的分词。 这可以通过直接使用 LlamaTokenizer 类，或为 AutoTokenizer 类传入 use_fast=False 选项来实现。用法示例如下。

import torch
from transformers import LlamaTokenizer, LlamaForCausalLM

## v2 models
model_path = 'openlm-research/open_llama_3b_v2'
# model_path = 'openlm-research/open_llama_7b_v2'

## v1 models
# model_path = 'openlm-research/open_llama_3b'
# model_path = 'openlm-research/open_llama_7b'
# model_path = 'openlm-research/open_llama_13b'

tokenizer = LlamaTokenizer.from_pretrained(model_path)
model = LlamaForCausalLM.from_pretrained(
    model_path, torch_dtype=torch.float16, device_map='auto',
)

prompt = 'Q: What is the largest animal?\nA:'
input_ids = tokenizer(prompt, return_tensors="pt").input_ids

generation_output = model.generate(
    input_ids=input_ids, max_new_tokens=32
)
print(tokenizer.decode(generation_output[0]))

有关更高级的用法，请遵循 transformers LLaMA 文档。

使用 LM-Eval-Harness 进行评估

可以使用 lm-eval-harness 对该模型进行评估。然而，由于上述分词器问题，我们需要避免使用快速分词器以获得正确的结果。这可以通过向 lm-eval-harness 的这一部分 传入 use_fast=False 来实现，如下面的示例所示：

tokenizer = self.AUTO_TOKENIZER_CLASS.from_pretrained(
    pretrained if tokenizer is None else tokenizer,
    revision=revision + ("/" + subfolder if subfolder is not None else ""),
    use_fast=False
)

使用 EasyLM 加载权重

若要在我们的 EasyLM 框架中使用这些权重，请参阅 EasyLM 的 LLaMA 文档。请注意，与原始 LLaMA 模型不同，我们的 OpenLLaMA 分词器和权重是完全从头开始训练的，因此不再需要获取原始的 LLaMA 分词器和权重。

数据集与训练

v1 模型是在 RedPajama 数据集上训练的。v2 模型是在 Falcon refined-web 数据集、StarCoder 数据集以及 RedPajama 数据集中的 wikipedia、arxiv、book 和 stackexchange 部分的混合数据上训练的。我们遵循与原始 LLaMA 论文完全相同的预处理步骤和训练超参数，包括模型架构、上下文长度、训练步数、学习率调度和优化器。我们的设置与原始设置之间的唯一区别在于所使用的数据集：OpenLLaMA 采用开放数据集，而非原始 LLaMA 所使用的数据集。

我们在云 TPU-v4s 上使用 EasyLM训练模型，这是我们开发的一个基于 JAX 的训练流水线，用于训练和微调大型语言模型。我们采用常规数据并行和完全分片数据并行 （也称为 ZeRO stage 3）的组合，以平衡训练吞吐量和内存使用。总体而言，我们的 7B 模型达到了每个 TPU-v4 芯片每秒超过 2200 个 token 的吞吐量。训练损失如下图所示。

评估

我们使用 lm-evaluation-harness 在广泛的任务上对 OpenLLaMA 进行了评估。LLaMA 的结果是通过在相同的评估指标上运行原始 LLaMA 模型生成的。我们注意到，我们针对 LLaMA 模型的结果与原始 LLaMA 论文略有不同，我们认为这是由于不同的评估协议所致。在 lm-evaluation-harness 的这个 issue 中也报告了类似的差异。此外，我们展示了 GPT-J 的结果，这是一个由 EleutherAI 在 Pile 数据集上训练的 6B 参数模型。

原始 LLaMA 模型训练了 1 万亿个 token，而 GPT-J 训练了 5000 亿个 token。我们在下表中展示了结果。OpenLLaMA 在大多数任务上表现出与原始 LLaMA 和 GPT-J 相当的性能，并在某些任务上优于它们。

任务/指标	GPT-J 6B	LLaMA 7B	LLaMA 13B	OpenLLaMA 3Bv2	OpenLLaMA 7Bv2	OpenLLaMA 3B	OpenLLaMA 7B	OpenLLaMA 13B
anli_r1/acc	0.32	0.35	0.35	0.33	0.34	0.33	0.33	0.33
anli_r2/acc	0.34	0.34	0.36	0.36	0.35	0.32	0.36	0.33
anli_r3/acc	0.35	0.37	0.39	0.38	0.39	0.35	0.38	0.40
arc_challenge/acc	0.34	0.39	0.44	0.34	0.39	0.34	0.37	0.41
arc_challenge/acc_norm	0.37	0.41	0.44	0.36	0.41	0.37	0.38	0.44
arc_easy/acc	0.67	0.68	0.75	0.68	0.73	0.69	0.72	0.75
arc_easy/acc_norm	0.62	0.52	0.59	0.63	0.70	0.65	0.68	0.70
boolq/acc	0.66	0.75	0.71	0.66	0.72	0.68	0.71	0.75
hellaswag/acc	0.50	0.56	0.59	0.52	0.56	0.49	0.53	0.56
hellaswag/acc_norm	0.66	0.73	0.76	0.70	0.75	0.67	0.72	0.76
openbookqa/acc	0.29	0.29	0.31	0.26	0.30	0.27	0.30	0.31
openbookqa/acc_norm	0.38	0.41	0.42	0.38	0.41	0.40	0.40	0.43
piqa/acc	0.75	0.78	0.79	0.77	0.79	0.75	0.76	0.77
piqa/acc_norm	0.76	0.78	0.79	0.78	0.80	0.76	0.77	0.79
record/em	0.88	0.91	0.92	0.87	0.89	0.88	0.89	0.91
record/f1	0.89	0.91	0.92	0.88	0.89	0.89	0.90	0.91
rte/acc	0.54	0.56	0.69	0.55	0.57	0.58	0.60	0.64
truthfulqa_mc/mc1	0.20	0.21	0.25	0.22	0.23	0.22	0.23	0.25
truthfulqa_mc/mc2	0.36	0.34	0.40	0.35	0.35	0.35	0.35	0.38
wic/acc	0.50	0.50	0.50	0.50	0.50	0.48	0.51	0.47
winogrande/acc	0.64	0.68	0.70	0.63	0.66	0.62	0.67	0.70
平均	0.52	0.55	0.57	0.53	0.56	0.53	0.55	0.57

我们已从基准测试中移除了任务 CB 和 WSC，因为我们的模型在这两项任务上的表现异常高。我们假设训练集中可能存在基准数据污染。

联系方式

我们非常希望收到来自社区的反馈。如果您有任何问题，请提交 issue 或联系我们。

OpenLLaMA 由以下人员开发： Xinyang Geng* 和 Hao Liu*，来自 Berkeley AI Research。 *同等贡献

致谢

我们感谢 Google TPU Research Cloud 项目提供了部分计算资源。我们特别感谢 TPU Research Cloud 的 Jonathan Caton 帮助我们组织计算资源，Google Cloud 团队的 Rafi Witten 和 Google JAX 团队的 James Bradbury 帮助我们优化训练吞吐量。我们还要感谢 Charlie Snell、Gautier Izacard、Eric Wallace、Lianmin Zheng 以及我们的用户社区提供的讨论和反馈。

OpenLLaMA 13B v1 模型是与 Stability AI 合作训练的，我们感谢 Stability AI 提供的计算资源。我们特别感谢 David Ha 和 Shivanshu Purohit 在协调后勤和提供工程支持方面的工作。

引用

如果您在研究或应用中发现 OpenLLaMA 有用，请使用以下 BibTeX 进行引用：

@software{openlm2023openllama,
  author = {Geng, Xinyang and Liu, Hao},
  title = {OpenLLaMA: An Open Reproduction of LLaMA},
  month = May,
  year = 2023,
  url = {https://github.com/openlm-research/open_llama}
}

@software{together2023redpajama,
  author = {Together Computer},
  title = {RedPajama-Data: An Open Source Recipe to Reproduce LLaMA training dataset},
  month = April,
  year = 2023,
  url = {https://github.com/togethercomputer/RedPajama-Data}
}

@article{touvron2023llama,
  title={Llama: Open and efficient foundation language models},
  author={Touvron, Hugo and Lavril, Thibaut and Izacard, Gautier and Martinet, Xavier and Lachaux, Marie-Anne and Lacroix, Timoth{\'e}e and Rozi{\`e}re, Baptiste and Goyal, Naman and Hambro, Eric and Azhar, Faisal and others},
  journal={arXiv preprint arXiv:2302.13971},
  year={2023}
}