“TinyLlama: An Open-Source Small Language Model”
项目主页：https://prompt-aligned.github.io/
论文地址：https://arxiv.org/pdf/2401.06105.pdf
Github地址：https://github.com/jzhang38/TinyLlama
摘要
TinyLlama是一个1.1B的紧凑语言模型，使用了大约1万亿个标记进行了约3个时期的预训练。它基于Llama 2的架构和分词器，并利用了开源社区的各种进展，如FlashAttention，从而实现了更好的计算效率。尽管尺寸相对较小，TinyLlama在一系列下游任务中表现出色，明显优于现有的具有相似尺寸的开源语言模型。模型的检查点和代码可以在GitHub上公开获取。
简介
自然语言处理（NLP）的最新进展主要是通过增加语言模型的规模来推动的。大型语言模型（LLMs）在广泛的任务上展示了其有效性。一些实证研究表明，只有具有足够多参数的模型才能展现出LLMs的新能力，例如少样本提示和思维链推理。其他研究关注LLMs的规模行为建模。然而，训练更大的模型和增加训练数据量应该以相同的速度增长，以实现计算最优模型。然而，训练更小的模型并使用更大的数据集的潜力仍未得到充分探索。相比于训练计算最优的语言模型，注重推理预算的语言模型更重要。通过训练更小的模型并使用更多的数据，可以达到甚至超过更大模型的性能。现有的规模定律可能无法准确预测训练更小模型更长时间的情况。
这项工作探索了在训练过程中使用比缩放定律建议的更多标记的较小模型的行为。研究人员使用了大约3万亿个标记来训练一个具有11亿参数的Transformer解码器模型。该模型名为TinyLlama，与现有的类似规模的开源语言模型相比表现出竞争力，并在各种下游任务中超过了OPT-1.3B和Pythia1.4B。TinyLlama是开源的，旨在提高语言模型研究者的可访问性。
预训练
预训练数据
采用自然语言数据和代码数据混合的方式进行预训练，其中自然语言数据来自SlimPajama，代码数据来自Starcoderdata。预训练过程中，使用Llama的分词器，总共预训练了约9500亿个标记，训练时间为三个时期。在训练过程中，自然语言数据和代码数据的比例为7:3。
架构
采用了类似于Llama 2的Transformer架构，使用RoPE注入位置信息，使用RMSNorm进行预归一化，使用SwiGLU作为激活函数，使用grouped-query attention来减少内存带宽开销和加速推理。
RoPE、RMSNorm、SwiGLU和grouped-query attention都是近期主流大语言模型中广泛采用的方法。
优化
本文介绍了一种名为FSDP的多GPU和多节点训练方法，以及优化的注意力机制Flash Attention、融合的层归一化、交叉熵损失和旋转位置嵌入等功能，这些功能共同提高了计算吞吐量。同时，本文还介绍了使用原始SwiGLU模块替换xFormers中的融合SwiGLU模块的方法，从而进一步提高了代码效率。实验结果表明，TinyLlama-1.1B模型的训练速度优于Pythia-1.0B和MPT-1.3B模型，且能够在较短时间内完成大规模模型训练，具有较高的实用价值。
训练
我们基于lit-gpt构建了我们的框架。在预训练阶段，我们采用了自回归语言建模目标，遵循Llama 2的设置。我们使用AdamW优化器，将β1设置为0.9，β2设置为0.95。此外，我们使用余弦学习率调度，最大学习率为4.0×10-4，最小学习率为4.0×10-5。我们使用了2000个预热步骤来促进优化学习。我们将批量大小设置为2M个标记。我们将权重衰减设置为0.1，并使用梯度剪裁阈值1.0来调节梯度值。我们在项目中使用16个A100-40G GPU对TinyLlama进行预训练。
结果
在常识推理任务中，TinyLlama在许多任务上表现优于基线模型，并获得了最高的平均分数。此外，我们还追踪了TinyLlama在常识推理基准测试中的准确性，并使用InstructEval基准测试评估了其问题解决能力。
TinyLlama在多项基准测试中表现出更好的问题解决能力。在预处理数据时，TinyLlama意外地插入了过多的EOS标记，这可能会对模型产生负面影响。但在修正后，TinyLlama的性能显著提高。评估任务包括MMLU、BBH、DROP和HumanEval。其中，MMLU用于衡量模型的世界知识和问题解决能力，BBH用于衡量模型在复杂指令遵循方面的能力，DROP用于衡量模型的数学推理能力，而HumanEval用于衡量模型的编程能力。
总结
TinyLlama是一个开源的、小规模的语言模型，具有紧凑的架构和良好的性能。它可以在移动设备上实现终端用户应用，并作为测试与语言模型相关的创新想法的轻量级平台。该项目将不断积累经验，开发改进版本的TinyLlama，增强其性能和多功能性。未来的报告将详细记录进一步的发现和结果。
评论0