[论文解读] TinyLlama: An Open-Source Small Language Model

Peiyuan Zhang, Guangtao Zeng+2 more|arXiv (Cornell University)|Jan 4, 2024

Natural Language Processing Techniques被引用 57

一句话总结

TinyLlama 是一个紧凑的 1.1B 参数的解码器仅语言模型，使用 Llama 2 架构和开源加速，在 ~3 trillion tokens 上进行预训练（≈3 epochs），在同等规模的开源模型中取得强劲表现。

ABSTRACT

We present TinyLlama, a compact 1.1B language model pretrained on around 1 trillion tokens for approximately 3 epochs. Building on the architecture and tokenizer of Llama 2, TinyLlama leverages various advances contributed by the open-source community (e.g., FlashAttention and Lit-GPT), achieving better computational efficiency. Despite its relatively small size, TinyLlama demonstrates remarkable performance in a series of downstream tasks. It significantly outperforms existing open-source language models with comparable sizes. Our model checkpoints and code are publicly available on GitHub at https://github.com/jzhang38/TinyLlama.

研究动机与目标

• 评估在非常大规模的预训练数据下，1.1B 的小型模型是否能达到强劲的性能。
• 利用开源高效性来提升训练速度和内存使用效率。
• 在常识和推理任务上，将 TinyLlama 与同等规模的现有开源模型进行比较。
• 通过发布数据、代码和检查点来展示开放性和可重复性。

提出的方法

• 采用 Llama 2 风格的解码器架构和分词器。
• 在 SlimPajama 自然语言数据和 Starcoderdata 代码数据的混合数据集上进行预训练（约 ~950B tokens）。
• 总共训练约 3 epochs，总标记数约为 ~3 trillion tokens。
• 应用速度/效率优化：Fully Sharded Data Parallel (FSDP)、FlashAttention、xFormers 调整，以及分组查询注意力。
• 使用 RoPE 位置嵌入和带有 SwiGLU 激活的 RMSNorm。
• 预训练遵循自回归语言模型目标，使用 AdamW、余弦学习率调度、预热，以及 2,000 warmup steps。

实验结果

研究问题

• RQ1 在异常大规模数据集（~3T tokens）上训练时，1.1B 参数的模型是否能取得有竞争力的性能？
• RQ2 开源效率改进是否能够在不牺牲性能的情况下实现更快的训练和更低的内存占用？
• RQ3 在常识推理和问题解决基准上，TinyLlama 与其他 1B 规模的开源模型相比如何？

主要发现

• TinyLlama 在多项下游任务的零-shot 评估中显著优于 OPT-1.3B 和 Pythia-1.4B。
• 在常识推理基准上，与同等规模的开源语言模型相比，表现具有竞争力（例如 HellaSwag、OpenBookQA、WinoGrande、ARC、BoolQ、PIQA）。
• 使用 ~3T tokens 的训练和效率优化可实现高吞吐量（≈24,000 tokens/s 每个 A100-40G），并且比可比模型所需的 GPU 小时更少。
• 在 InstructEval 任务（MMLU、BBH、HumanEval、DROP）上，TinyLlama 展现出比基线更强的解决问题能力。
• 该模型保持开源，发布了预训练代码、中间检查点以及数据处理细节。