TinyLlama-1.1B小型但强大的语言模型
TinyLlama 项目旨在在 3 万亿 tokens 上进行预训练，构建一个拥有11亿参数的Llama模型。经过精心优化，"仅"需16块A100-40G的GPU，便可在90天内完成这个任务。训练已于2023-09-01开始。
该项目采用了与Llama 2完全相同的架构和分词器。这意味着TinyLlama可以在许多基于Llama的开源项目中即插即用。此外，TinyLlama只有1.1B的参数，体积小巧，适用于需要限制计算和内存占用的多种应用。
发布时间表
项目团队会根据以下计划逐步发布中间checkpoint。同时也列了一些基线模型进行比较。
| Date | HF Checkpoint | Tokens | Step | HellaSwag Acc_norm |
|---|---|---|---|---|
| Baseline | StableLM-Alpha-3B | 800B | -- | 38.31 |
| Baseline | Pythia-1B-intermediate-step-50k-105b | 105B | 50k | 42.04 |
| Baseline | Pythia-1B | 300B | 143k | 47.16 |
| 2023-09-04 | TinyLlama-1.1B-intermediate-step-50k-105b | 105B | 50k | 43.50 |
| 2023-09-16 | -- | 500B | -- | -- |
| 2023-10-01 | -- | 1T | -- | -- |
| 2023-10-16 | -- | 1.5T | -- | -- |
| 2023-10-31 | -- | 2T | -- | -- |
| 2023-11-15 | -- | 2.5T | -- | -- |
| 2023-12-01 | -- | 3T | -- | -- |
潜在场景
小型但强大的语言模型对许多应用都很有用。以下是一些潜在的场景：
- 帮助对大型模型进行speculative decoding。
- 在边缘装置上运行，比如离线的实时机器翻译 (TinyLlama的4比特量化版本的模型权重只需要550MB的内存)。
- 在游戏中实现实时对话生成(因为还得给游戏本身留显存所以模型要小)。
此外，项目代码可以给初学者做一个入门预训练的简洁参考。如果你要训练50亿以下参数的语言模型, 你其实不需要Megatron-LM。
训练细节
以下是训练设置的一些细节：
| Setting | Description |
|---|---|
| Parameters | 1.1B |
| Attention Variant | Grouped Query Attention |
| Model Size | Layers: 22, Heads: 32, Query Groups: 4, Embedding Size: 2048, Intermediate Size (Swiglu): 5632 |
| Sequence Length | 2048 |
| Batch Size | 2 million tokens (2048 * 1024) |
| Learning Rate | 4e-4 |
| Learning Rate Schedule | Cosine with 2000 warmup steps |
| Training Data | Slimpajama & Starcoderdata |
| Data Preprocessing | Excluded GitHub subset of Slimpajama; Sampled all code from Starcoderdata |
| Combined Dataset Size | Around 950B tokens |
| Total Tokens During Training | 3 trillion (slightly more than 3 epochs/143k steps) |
| Natural Language to Code Ratio | 7:3 |
| Hardware | 16 A100-40G GPUs |
速度极快
代码库支持以下特性：
- 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的训练速度的比较。
| Model | A100 GPU hours taken on 300B tokens |
|---|---|
| TinyLlama-1.1B | 3456 |
| Pythia-1.0B | 4830 |
| MPT-1.3B | 7920 |
Pythia的数字来自他们的论文。MPT的数字来自这里，作者说MPT-1.3B"was trained on 440 A100-40GBs for about half a day" on 200B tokens。
TinyLlama是一个相对较小的模型，使用了GQA，这意味着它在推理期间也很快。以下是测量的一些推理速度：
| Framework | Device | Settings | Throughput (tokens/sec) |
|---|---|---|---|
| Llama.cpp | Mac M2 16GB RAM | batch_size=1; 4-bit inference | 71.8 |
| vLLM | A40 GPU | batch_size=100, n=10 | 7094.5 |