【Title】
TinyLlama：开源小语言模型
【Abstract】
我们推出了 TinyLlama，这是一个紧凑的 1.1B 语言模型，在大约 1 万亿个令牌上进行了大约 3 个时期的预训练。 TinyLlama 基于 Llama 2（Touvron 等人，2023b）的架构和标记器构建，利用开源社区贡献的各种进步（例如 FlashAttention（Dao，2023）），实现了更好的计算效率。 尽管规模相对较小，TinyLlama 在一系列下游任务中表现出了出色的性能。 它的性能明显优于规模相当的现有开源语言模型。 我们的模型检查点和代码可在 GitHub 上公开获取：https://github.com/jzhang38/TinyLlama。
【Introduction】
自然语言处理 (NLP) 的最新进展在很大程度上是通过扩大语言模型大小来推动的（Brown 等人，2020；Chowdhery 等人，2022；Touvron 等人，2023a,b）。 在广泛的文本语料库上预训练的大型语言模型 (LLM) 已证明其在各种任务上的有效性（OpenAI，2023；Touvron 等人，2023b）。 一些实证研究证明了法学硕士的涌现能力，这些能力可能只在具有足够多参数的模型中体现出来，例如少样本提示（Brown et al., 2020）和链式推理（Wei et al., 2020）。 ，2022）。 其他研究侧重于对法学硕士的扩展行为进行建模（Kaplan 等人，2020；Hoffmann 等人，2022）。 霍夫曼等人。 （2022）建议，为了训练计算最优模型，模型的大小和训练数据量应该以相同的速度增加。 这提供了在计算预算固定时如何最佳选择模型大小和分配训练数据量的指南。
尽管这些工作显示出对大型模型的明显偏好，但使用较大数据集训练较小模型的潜力仍未得到充分探索。 Touvron 等人没有训练计算最优的语言模型。 (2023a) 强调了推理预算的重要性，而不是仅仅关注训练计算最优语言模型。 推理最优语言模型的目标是在特定推理约束内实现最佳性能，这是通过使用比缩放定律建议的标记更多的标记来训练模型来实现的（Hoffmann 等人，2022）。 图夫龙等人。 (2023a) 表明，较小的模型在使用更多数据进行训练时，可以匹配甚至超越较大的模型。 此外，Thaddée (2023) 指出，现有的缩放法则 (Hoffmann et al., 2022) 在较小模型训练时间较长的情况下可能无法准确预测。
受这些新发现的推动，这项工作的重点是探索较小模型在使用比缩放定律建议的数量多得多的标记进行训练时的行为（Hoffmann 等人，2022）。 具体来说，我们使用大约 3 万亿个令牌训练具有 1.1B 个参数的 Transformer 仅解码器模型（Vaswani 等人，2017）。 据我们所知，这是第一次尝试使用如此大量的数据来训练具有 1B 参数的模型。 遵循与 Llama 2（Touvron 等人，2023b）相同的架构和分词器，我们将模型命名为 TinyLlama。 与现有类似规模的开源语言模型相比，TinyLlama 显示出具有竞争力的性能。 具体来说，TinyLlama 在各种下游任务中都超越了 OPT-1.3B (Zhang et al., 2022) 和 Pythia1.4B (Biderman et al., 2023)。
我们的 TinyLlama 是开源的，旨在提高语言模型研究人员的可访问性。 我们相信其卓越的性能和紧凑的尺寸使其成为语言模型研究的研究人员和从业者的一个有吸引力的平台。
【Pretraining】
本节介绍我们如何预训练 TinyLlama。 首先，我们介绍预训练语料库和数据采样方法的细节。 接下来，我们详细介绍模型架构和预训练期间使用的超参数。
2.1 Pre-training data
我们的主要目标是使预训练过程有效且可重复。 我们采用自然语言数据和代码数据的混合来预训练 TinyLlama，从 SlimPajama (Soboleva et al., 2023) 获取自然语言数据，从 Starcoderdata (Li et al., 2023) 获取代码数据。 我们采用 Llama 的 tokenizer (Touvron et al., 2023a) 来处理数据。
SlimPajama 这是一个基于 RedPajama（Together Computer，2023）为训练语言模型而创建的大型开源语料库。 最初的 RedPajama 语料库是一项开源研究成果，旨在重现 Llama 的预训练数据（Touvron 等人，2023a），其中包含超过 1.2 万亿个代币。 SlimPajama 是通过对原始RedPajama 进行清理和去重而衍生出来的。
Starcoderdata 收集该数据集是为了训练 StarCoder（Li 等人，2023），这是一种强大的开源大型代码语言模型。 它包含 86 种编程语言的约 2500 亿个代币。 除了代码之外，它还包括 GitHub 问题和涉及自然语言的文本代码对。 为了避免数据重复，我们删除了 SlimPajama 的 GitHub 子集，仅从 Starcoderdata 中删除了示例代码数据。
合并这两个语料库后，我们总共有大约 9500 亿个代币用于预训练。 根据 Muennighoff 等人的观察，TinyLlama 在这些令牌上接受了大约三个 epoch 的训练。 （2023），与使用唯一数据相比，对数据进行最多四个时期的重复训练会导致性能下降最小。 在训练过程中，我们对自然语言数据进行采样，以实现自然语言数据和代码数据之间大约 7:3 的比例。
2.2 Architecture
表1：模型架构细节
我们采用与 Llama 2 类似的模型架构（Touvron 等人，2023b）。 我们使用基于 Vaswani 等人的 Transformer 架构。 （2017），详细信息如下：
位置嵌入我们使用 RoPE（旋转位置嵌入）（Su et al., 2021）将位置信息注入到我们的模型中。 RoPE 是最近被许多主流大型语言模型广泛采用的方法，例如 PaLM (Anil et al., 2023)、Llama (Touvron et al., 2023a) 和 Qwen (Bai et al., 2023)。
RMSNorm 在预归一化中，为了获得更稳定的训练，我们在每个变压器子层之前对输入进行归一化。 此外，我们应用 RMSNorm（Zhang 和 Sennrich，2019）作为我们的标准化技术，可以提高训练效率。
SwiGLU 我们没有使用传统的 ReLU 非线性，而是遵循 Llama 2，将 Swish 和门控线性单元结合在一起，称为 SwiGLU (Shazeer, 2020)，作为 TinyLlama 中的激活函数。
分组查询注意力 为了减少内存带宽开销并加快推理速度，我们在模型中使用分组查询注意力（Ainslie et al., 2023）。 我们有 32 个用于查询注意力的头，并使用 4 组键值头。 通过这种技术，模型可以在多个头之间共享键和值表示，而不会牺牲太多性能。
2.3 Speed Optimizations
完全分片数据并行 (FSDP) 在训练期间，我们的代码库集成了 FSDP1，以有效利用多 GPU 和多节点设置。 这种集成对于跨多个计算节点扩展训练过程至关重要，从而显着提高训练速度和效率。
Flash Attention 另一个关键改进是集成了 Flash Attention 2（Dao，2023），这是一种优化的注意力机制。 该存储库还提供融合层范数、融合交叉熵损失和融合旋转位置嵌入，它们在提高计算吞吐量方面发挥着关键作用。
xFormers 我们已将 xFormers（Lefaudeux 等人，2022）存储库中的融合 SwiGLU 模块替换为原始 SwiGLU 模块，进一步提高了我们代码库的效率。 借助这些功能，我们可以减少内存占用，使 1.1B 模型能够适应 40GB 的 GPU RAM。
性能分析以及与其他模型的比较 这些元素的结合使我们的训练吞吐量达到了每 A100-40G GPU 每秒 24,000 个令牌。 与 Pythia-1.0B（Biderman 等人，2023）和 MPT-1.3B 2 等其他模型相比，我们的代码库表现出了卓越的训练速度。 例如，TinyLlama-1.1B 模型对于 300B 代币仅需要 3,456 A100 GPU 小时，而 Pythia 为 4,830 小时，MPT 为 7,920 小时。 这显示了我们优化的有效性以及在大规模模型训练中节省大量时间和资源的潜力。
2.4 Training
我们基于 lit-gpt.3 构建框架。遵循 Llama 2（Touvron 等人，2023b），我们在预训练阶段采用自回归语言建模目标。 与 Llama 2 的设置一致，我们利用 AdamW 优化器（Loshchilov 和 Hutter，2019），将 β1 设置为 0.9，β2 为 0.95。 此外，我们使用余弦学习率计划，最大学习率为 4.0 × 10−4，最小学习率为 4.0 × 10−5。 我们使用 2,000 个预热步骤来促进优化学习。4 我们将批量大小设置为 2M 令牌。 我们将权重衰减指定为0.1，并使用梯度裁剪阈值1.0来调节梯度值。 我们在项目中使用 16 个 A100-40G GPU 预训练 TinyLlama。
图 1：我们的代码库与 Pythia 和 MPT 的训练速度比较。
【Results】
我们在广泛的常识推理和问题解决任务上评估 TinyLlama，并将其与具有相似模型参数的几种现有开源语言模型进行比较。
基线模型 我们主要关注具有纯解码器架构的语言模型，包含大约 10 亿个参数。 具体来说，我们将 TinyLlama 与 OPT-1.3B (Zhang et al., 2022)、Pythia-1.0B 和 Pythia-1.4B (Biderman et al., 2023) 进行比较。
常识推理任务为了了解 TinyLlama 的常识推理能力，我们考虑以下任务：Hellaswag (Zellers et al., 2019)、OpenBookQA (Mihaylov et al., 2018)、WinoGrande (Sakaguchi et al., 2021)、ARC- Easy 和 ARC-Challenge（Clark 等人，2018）、BoolQ（Clark 等人，2019）和 PIQA（Bisk 等人，2020）。 我们采用语言模型评估框架（Gao et al., 2023）来评估模型。 按照之前的实践（Biderman 等人，2023），模型在这些任务的零样本设置中进行评估。 结果如表 2 所示。我们注意到，TinyLlama 在许多任务上都优于基线，并获得了最高的平均分数。
表 2：常识推理任务的零样本性能。
训练期间性能的演变 我们在预训练期间跟踪了 TinyLlama 在常识推理基准上的准确性，如图 2 所示。一般来说， TinyLlama 随着计算资源的增加而改进，在大多数基准测试中超过了 Pythia-1.4B 的准确性。
图 2：预训练期间常识推理基准的性能演变。 Pythia-1.4B的性能也包含在图中进行比较。
问题解决评估我们还使用 InstructEval 基准评估 TinyLlama 的问题解决能力（Chia 等人，2023）。 该基准包括以下任务：
• 大规模多任务语言理解（MMLU）（Hendrycks 等人，2021）：该任务用于衡量模型跨各个学科的世界知识和解决问题的能力。 我们在 5 次拍摄设置中评估模型。
• BIG-Bench Hard (BBH)（Suzgun 等人，2023）：这是 BIG-Bench 基准（Srivastava 等人，2022）中 23 项挑战性任务的子集，旨在衡量语言模型遵循复杂指令的能力 。 这些模型在 3 次设置中进行评估。
• 段落离散推理 (DROP)（Dua 等人，2019）：此阅读理解任务衡量模型的数学推理能力。 我们在 3-shot 设置中评估模型。
• HumanEval（Zheng et al., 2023）：该任务用于衡量模型的编程能力。 这些模型在零样本设置中进行评估。
评估结果如表 3 所示。我们观察到，与现有模型相比，TinyLlama 表现出了更好的问题解决能力。
表 3：InstructEval 基准测试中问题解决任务的表现。
【Conclusion】
在本文中，我们介绍了 TinyLlama，一种开源的小型语言模型。 为了提高开源LLM预训练社区的透明度，我们发布了所有相关信息，包括我们的预训练代码、所有中间模型检查点以及数据处理步骤的详细信息。 凭借其紧凑的架构和良好的性能，TinyLlama 可以在移动设备上启用最终用户应用程序，并作为一个轻量级平台来测试与语言模型相关的各种创新想法。 我们将利用该项目开放、上线阶段积累的丰富经验，旨在开发TinyLlama的改进版本，为其配备多种功能，以增强其在各种任务中的性能和多功能性。 我们将在即将发布的报告中记录进一步的发现和详细结果。
【Acknowledgements】
们对开源社区在我们研究的开放、实时阶段给予的大力支持表示感谢。 特别感谢Qian Liu、Longxu Dou、Hai Leong Chieu 和Larry Law 对我们项目的帮助。 本研究/项目得到新加坡教育部学术研究基金 (AcRF) Tier 2 计划（MOE AcRF Tier 2 奖项编号：MOE-T2EP20122-0011）、新加坡教育部 Tier 3 的支持 计划（奖项编号：MOET320200004）、新加坡国家研究基金会和 DSO 国家实验室的 AI 新加坡计划（AISG 奖项编号：AISG2-RP-2020-016）、AI 新加坡博士奖学金（AISG 奖项编号：AISG2- Doctor-2021-08-007）、SUTD 启动项目 (SKI 2021_03_11) 以及拨款 RS-INSUR-00027-E0901-S00。 本材料中表达的任何意见、调查结果和结论或建议均为作者的观点，并不反映资助机构的观点。