Skip to content

Latest commit

 

History

History
188 lines (135 loc) · 12.6 KB

README_zh-CN.md

File metadata and controls

188 lines (135 loc) · 12.6 KB
Raw

TinyLlama-1.1B

English | 中文

Chat Demo

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

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

新闻

  • 2023-12-18：
    • 添加两个文档 1, 2 说明训练曲线、项目时间表和错误修复的变化。
  • 2023-10-03:
    • 在speculative decoding中添加llama.cpp的代码示例。具体请查看 speculative_decoding/README.md。
    • 2023-10-02: 1. 1T-token检查点刚发布。2. 我们在huggingface上记录了所有中间检查点。
    • 2023-09-28: 启用Discord服务器。
  • 2023-09-18:
    • 发布了一个 chat demo，欢迎点击链接来尝试我们的模型。
  • 2023-09-16:

发布时间表

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

基座模型:

Date 模型权重 Tokens Step Commonsense Avg
2023-09-01 Pythia-1.0B 300B 143k 48.30
2023-09-04 TinyLlama-1.1B-intermediate-step-50k-105b (ModelScope) 105B 50k 46.11
2023-09-16 TinyLlama-1.1B-intermediate-step-240k-503b (ModelScope) 503B 240K 48.28
2023-10-01 TinyLlama-1.1B-intermediate-step-480k-1T 1T 480k 50.22
2023-11-04 TinyLlama-1.1B-intermediate-step-715k-1.5T 1.5T 715k 51.28
2023-11-20 TinyLlama-1.1B-intermediate-step-955k-2T 2T 955k 51.64
2023-12-11 TinyLlama-1.1B-intermediate-step-1195k-2.5T 2.5T 1195k 53.86
2023-12-28 TinyLlama-1.1B-intermediate-step-1431k-3T 3T 1431k 52.99

对话模型:

Date 模型权重 Tokens Step Commonsense Avg
2023-09-16 TinyLlama-1.1B-Chat-V0.1 (ModelScope) 503B 240K 49.57
2023-10-1 TinyLlama-1.1B-Chat-V0.3 1T 480K 51.36
2023-11-04 TinyLlama-1.1B-Chat-V0.4 1.5T 715K 52.30

需要注意的是，由于我们的现在模型还处于训练初期，学习率并没有完全稳定下来，为了更好的体验我们的模型，您可以下载我们 聊天模型 或者通过 chat demo 来尝试我们的模型。

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

潜在场景

小型但强大的语言模型对许多应用都很有用。以下是一些潜在的场景：

  • 帮助对大型模型进行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

速度极快

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

  • 使用FSDP进行多GPU和多节点分布式训练
  • flash attention 2
  • 融合层归一化 (fused layernorm)
  • 融合swiglu (fused swiglu)
  • 融合交叉熵损失 (fused cross entropy loss)
  • 融合旋转位置嵌入 (fused rotary positional embedding)

致谢：flash attention 2、融合层归一化、融合交叉熵损失和融合旋转位置嵌入来自于FlashAttention仓库；融合swiglu来自于xformers。

有了这些优化, 我们可以达到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

开始预训练

请参考PRETRAIN.md。

微调

  • 我们在 sft 中添加了我们进行微调和推理的代码。并且基于这个代码我们在openassistant-guanaco 数据集上进行了微调，得到了我们的第一版聊天模型。
  • 如果您希望在 RAM 小于 4GB 的 GPU 上对用我们的模型进行微调，可以参考并使用 Qlora 和 bitsandbytes 项目。
  • 目前微调的时候我们并没有广泛对超参进行搜索，也没有选择潜在更优的 instruction 数据集。我们希望促进 NLP 社区对于我们的TinyLlama模型的开放研究，并开源更好的微调聊天模型。我们也会把这些模型放在这个项目中。

TODO

该项目仍在积极开发中。我们团队很小，非常欢迎社区的反馈和贡献。以下是我们计划进行的一些工作：

  • Add scripts for pretraining on other datasets.
  • Sequence length extrapolation.
  • Test out speculative decoding for Llama-2-7B.
  • Test the throughput on RTX 3090/4090.
  • Add fine-tuning scripts.
  • Properly evaluate the model on downstream tasks.
  • A demo running on mobile phones.
  • Explore retrieval-augmentation.

Star History

Star History Chart

Acknowledgements

这个仓库基于出色的开源项目lit-gpt和flash-attention构建.

@online{lit-gpt,
  author    = {Lightning AI},
  title     = {Lit-GPT},
  url       = {https://github.com/Lightning-AI/lit-gpt},
  year      = {2023},
}
@article{dao2023flashattention2,
  title     ={Flash{A}ttention-2: Faster Attention with Better Parallelism and Work Partitioning},
  author    ={Dao, Tri},
  year      ={2023}
}

Citation

此项目目前由Peiyuan Zhang，Guangtao Zeng，Tianduo Wang和Wei Lu贡献。

如果您觉得我们的工作有价值， 可以引用:

@misc{zhang2024tinyllama,
      title={TinyLlama: An Open-Source Small Language Model}, 
      author={Peiyuan Zhang and Guangtao Zeng and Tianduo Wang and Wei Lu},
      year={2024},
      eprint={2401.02385},
      archivePrefix={arXiv},
      primaryClass={cs.CL}
}