TinyLlama：1.1B迷你型羊驼，冲、冲、冲！ - 文章 - 开发者社区


          
论文题目：TinyLlama: An Open-Source Small Language Model
          
论文链接：https://arxiv.org/pdf/2401.02385.pdf
          
Github：https://github.com/jzhang38/TinyLlama

概要

研究背景：

自然语言处理（NLP）的进步很大程度上依赖于语言模型规模的扩大。
大型语言模型（LLMs）在广泛的任务中表现出色，且在参数数量足够大的模型中可能展现出新的能力，如少样本提示（few-shot prompting）和链式推理（chain-of-thought reasoning）。
尽管大型模型受到青睐，但训练更小的模型并使用更大的数据集的潜力尚未充分探索。

模型设计：

提出了TinyLlama，一个基于Llama 2架构和tokenizer的紧凑1.1B参数语言模型。
使用了大约1万亿token进行预训练，大约3个epoch。
利用了开源社区的贡献，如FlashAttention，提高了计算效率。
采用了Transformer架构，包括RoPE、RMSNorm、SwiGLU和分组查询注意力等技术。
实现了全分片数据并行（FSDP）、Flash Attention等速度优化措施。

picture.image

实验结论：

TinyLlama在一系列下游任务中表现出色，显著优于现有类似规模的开源语言模型。
在常识推理和问题解决任务上，TinyLlama超越了OPT-1.3B和Pythia-1.4B。
TinyLlama的训练吞吐量达到了每A100-40G GPU每秒24,000个tokens。
与Pythia-1.0B和MPT-1.3B等模型相比，TinyLlama的训练速度更快，例如，TinyLlama-1.1B模型在300B tokens的训练中所需的A100 GPU小时数远少于Pythia和MPT模型。

讨论

1、 TinyLlama在预训练数据的选择上有何独特之处，这对模型性能有何影响？

答案： TinyLlama采用了自然语言数据和代码数据的混合数据集进行预训练，这种数据混合策略可能有助于模型更好地理解和生成自然语言，同时也能处理代码相关的任务。这种数据选择可能使得模型在处理自然语言任务时更加灵活和准确，同时也能应对编程相关的语言理解任务。

2、TinyLlama在常识推理和问题解决任务上的优异表现可能归因于哪些因素？

答案： TinyLlama在常识推理和问题解决任务上的优异表现可能归因于其预训练数据的多样性和规模，以及模型架构的先进性。混合数据集可能使得模型在处理不同类型任务时更加稳健。此外，模型架构中采用的RoPE、RMSNorm、SwiGLU等技术可能有助于模型更好地捕捉长距离依赖和上下文信息，这对于常识推理和问题解决任务至关重要。

3、TinyLlama具体训练过程？

答案：

TinyLlama的训练过程涉及以下几个关键步骤：

预训练数据准备：
使用了两个主要的数据集：SlimPajama和Starcoderdata。
SlimPajama是一个基于RedPajama的大型开源语料库，经过清理和去重，包含了大约950亿个token。
Starcoderdata是一个用于训练StarCoder的大型代码语言模型的数据集，包含约250亿个token，覆盖86种编程语言，以及GitHub问题和文本-代码对。
为了避免数据重复，移除了SlimPajama中的GitHub子集，仅从Starcoderdata中采样代码数据。
TinyLlama在这些token上进行了大约三个epoch的训练。
模型架构：
采用了与Llama 2相似的Transformer架构，基于Vaswani等人（2017）的工作。
模型具有特定的隐藏层大小、中间隐藏层大小、上下文长度、头数、层数和词汇表大小。
使用了RoPE（Rotary Positional Embedding）来注入位置信息。
在预归一化中应用了RMSNorm（Root Mean Square Layer Normalization）以提高训练稳定性。
使用了SwiGLU作为激活函数，它是Swish和Gated Linear Unit的结合。
为了减少内存带宽开销和加速推理，使用了分组查询注意力。
速度优化：
在训练过程中集成了Fully Sharded Data Parallel (FSDP)来有效地利用多GPU和多节点设置。
集成了Flash Attention，一个优化的注意力机制，以及融合的层归一化、融合的交叉熵损失和融合的旋转位置嵌入，以提高计算吞吐量。
从xFormers库中替换了融合的SwiGLU模块，以进一步增强代码库的效率。
训练过程：
基于lit-gpt框架构建训练框架。
使用AdamW优化器，设置了特定的β1和β2值，以及学习率调度策略。
使用了2000个warmup步骤来优化学习过程。
设置了2M tokens的批次大小，权重衰减为0.1，并使用1.0的梯度裁剪阈值来调节梯度值。
使用16个A100-40G GPU进行预训练。