​TinyLlama-1.1B：从零开始训练一个精悍的语言模型

©PaperWeekly 原创 · 作者 | StatNLP

单位 | 新加坡科技设计大学

TinyLlama 项目旨在在 3 万亿 tokens 上进行预训练，构建一个拥有 11 亿参数的 Llama 模型。经过精心优化，我们“仅”需 16 块 A100-40G 的 GPU，便可在 90 天内完成这个任务🚀🚀。模型训练已于 2023-09-01 开始，目前还在训练中。

我们采用了与 Llama 2 完全相同的架构和分词器。这意味着 TinyLlama 可以在许多基于 Llama 的开源项目中即插即用。此外，TinyLlama 只有 1.1B 的参数，体积小巧，适用于需要限制计算和内存占用的多种应用。

背景与动机

1.1 从Llama 1/2 中的一些观察开始

以上是从 Llama 2 论文中摘取的训练损失曲线。他们论文中提到即使在对 2 万亿个 token 进行预训练后，模型仍然没有显示出任何饱和迹象。这意味着 Llama 2 的训练可以继续下去，而且还有较大的提升空间。这与之前 Chinchilla Scaling Law 的预期非常不一样（Chinchilla Scaling Law 认为 7B 的模型最优的训练 token 数量是 140B 左右）。

上图是一位 reddit 网友 [1] 分析的 llama2 不同大小模型的训练花费 v.s. Perplexity 的数据，从图中其实也可以观察到，很多情况下，更小的模型达到相同的 Perplexity 所需要的花费反而更小，这更坚定了我们的想法：我们可以在更多的数据上面训练更小的模型。

1.2 为什么要训练更小的Llama模型

• 最近大热的模型推理加速算法 speculative decoding 需要一个小模型来辅助对大模型的加速，Llama2 最小的模型就是 7B 了，但我们认为用一个更小的模型作为 speculator 可以让 speculative decoding 更加实用。（注：此方法要求大小模型的 tokenizer 保持一致，所以 Pythia 等最近相对热门的小模型没法作为 Llama2 的 speculator，而 TinyLlama 可以）

• 更小的模型可以在更多的场景下使用，比如在手机上离线运行，TinyLlama-1.1B 经过 4-bit 量化的模型权重只需要 550MB 的内存。

• 可能可以用在大型单机游戏里面，因为这种场景下必须留足够的显存给游戏本身，所以模型要尽可能的小。

发布时间表和训练细节

我们会根据以下计划逐步发布中间 checkpoint。我们也列了一些基线模型进行比较。

从上面可以看出，TinyLlama 目前的进展非常好🎉🎉。

你也可以在这里 [6] 实时跟踪 TinyLlama 的训练损失。

以下是我们训练设置的一些细节：

速度&显存占用

我们的代码库支持以下特性：

• multi-gpu and multi-node distributed training with FSDP.

• flash attention 2.

• fused layernorm.

• fused swiglu.

• fused cross entropy loss .

• fused rotary positional embedding.

有了这些优化, 我们可以达到 24k tokens/秒/A100 的训练速度，也就是 56% 的 MFU（在 A100-80G 上的 MFU 会更高）。这个速度可以让你可以在 8 个 A100 上用 32 小时训练一个 chinchilla-optimial 的模型（11 亿参数，220 亿token）。

这些优化也大大减少了显存占用, 我们可以把 11 亿参数的模型塞入 40GB 的 GPU 里面还能同时维持 16k tokens 的 per-gpu batch size。只需要把 batch size 改小一点， 你就可以在 RTX 3090/4090 上面训练 TinyLlama。下面是我们的代码库与 Pythia 和 MPT 的训练速度的比较。

▲ Pythia 的数字来自他们的论文。MPT 的数字来自这里 [9]，作者说 MPT-1.3B“was trained on 440 A100-40GBs for about half a day”on 200B tokens。

TinyLlama 是一个相对较小的模型, 同时我们用了 GQA，这意味着它在推理期间也很快。以下是我们测量的一些推理速度：

代码链接

我们的代码可以给初学者做一个入门预训练的简洁参考（3090/4090 能跑，一共就十多个.py文件）。虽然如今预训练比较通用的框架中 Megatron-LM 占据了半壁江山，但是对于小白来说实在上手难度太大了。而我们 TinyLlama 对于小实验室来说也相当友好。不要 999，也不用 99，只需要 16 张 A100 就可以在一个月左右让 LLama 1.1B 的模型完成预训练 （1T token）。

基于这个代码库， 我们公布了 TinyLlama 计划，目标在 90 天里用 16 个 A100 在 3 万亿 token 上训练一个 1.1B 的 TinyLlama，同时我们的代码库提供了一些预训练优化的同时保持了简洁性，3090/4090 可训练，不失为新手入门预训练的好选择。

https://github.com/jzhang38/TinyLlama

参考文献

[1] https://www.reddit.com/r/LocalLLaMA/comments/154cnvf/llama_2_scaling_laws/

[2] https://huggingface.co/stabilityai/stablelm-base-alpha-3b

[3] https://huggingface.co/EleutherAI/pythia-1b/tree/step50000

[4] https://huggingface.co/EleutherAI/pythia-1b

[5] https://huggingface.co/PY007/TinyLlama-1.1B-step-50K-105b

[6] https://links-cdn.wandb.ai/wandb-public-images/links/lance777/pgvhrsny.html

[7] https://huggingface.co/datasets/cerebras/slimpajama-627b

[8] https://huggingface.co/datasets/bigcode/starcoderdata

[9] https://huggingface.co/mosaicml/mpt-1b-redpajama-200b

[10] https://github.com/ggerganov/llama.cpp

[11] https://github.com/vllm-project/vllm

更多阅读

#投 稿 通 道#

 让你的文字被更多人看到 

如何才能让更多的优质内容以更短路径到达读者群体，缩短读者寻找优质内容的成本呢？答案就是：你不认识的人。

总有一些你不认识的人，知道你想知道的东西。PaperWeekly 或许可以成为一座桥梁，促使不同背景、不同方向的学者和学术灵感相互碰撞，迸发出更多的可能性。 

PaperWeekly 鼓励高校实验室或个人，在我们的平台上分享各类优质内容，可以是最新论文解读，也可以是学术热点剖析、科研心得或竞赛经验讲解等。我们的目的只有一个，让知识真正流动起来。

📝 稿件基本要求：

• 文章确系个人原创作品，未曾在公开渠道发表，如为其他平台已发表或待发表的文章，请明确标注 

• 稿件建议以 markdown 格式撰写，文中配图以附件形式发送，要求图片清晰，无版权问题

• PaperWeekly 尊重原作者署名权，并将为每篇被采纳的原创首发稿件，提供业内具有竞争力稿酬，具体依据文章阅读量和文章质量阶梯制结算

📬 投稿通道：

• 投稿邮箱：[e-mail] 

• 来稿请备注即时联系方式（微信），以便我们在稿件选用的第一时间联系作者

• 您也可以直接添加小编微信（pwbot02）快速投稿，备注：姓名-投稿

△长按添加PaperWeekly小编

🔍

现在，在「知乎」也能找到我们了

进入知乎首页搜索「PaperWeekly」

点击「关注」订阅我们的专栏吧

·

·