The TinyLlama project aims to pretrain a 1.1B Llama model on 3 trillion tokens. With some proper optimization, we can achieve this within a span of "just" 90 days using 16 A100-40G GPUs 🚀🚀. The training has started on 2023-09-01.
We adopted exactly the same architecture and tokenizer as Llama 2. This means TinyLlama can be plugged and played in many open-source projects built upon Llama. Besides, TinyLlama is compact with only 1.1B parameters. This compactness allows it to cater to a multitude of applications demanding a restricted computation and memory footprint.
- 2023-09-28: Add a discord server.
- 2023-09-18: 1. We added a chat demo so that you can play with TinyLlama-Chat-V0.1 right away.
- 2023-09-16: 1. We released the intermediate checkpoint trained on 503B tokens. 2. We released a chat model finetuned on OpenAssisant and simple finetuning scripts is added. 3. More eval benchmarks are added and documented in EVAL.md.
You can find the evaluation results of TinyLlama in EVAL.md.
We will be rolling out intermediate checkpoints following the below schedule.
| Date | HF Checkpoint | Tokens | Step | Commonsense Avg |
|---|---|---|---|---|
| 2023-09-01 | Pythia-1.0B | 300B | 143k | 48.30 |
| 2023-09-04 | TinyLlama-1.1B-intermediate-step-50k-105b | 105B | 50k | 46.11 |
| 2023-09-16 | TinyLlama-1.1B-intermediate-step-240k-503b | 503B | 240K | 48.28 |
| 2023-09-16 | TinyLlama-1.1B-Chat-V0.1 | 503B | 240K | 49.57 |
| 2023-10-01 | -- | 1T | -- | -- |
| 2023-10-16 | -- | 1.5T | -- | -- |
| 2023-10-31 | -- | 2T | -- | -- |
| 2023-11-15 | -- | 2.5T | -- | -- |
| 2023-12-01 | -- | 3T | -- | -- |
Note that the learning rate of the base model has not cooled down yet so we recommend you to also use the finetuned chat model.
Meanwhile, you can track the live cross entropy loss here.
Tiny but strong language models are useful for many applications. Here are some potential usecases:
- Assisting speculative decoding of larger models. (See this tutorial by Andrej Karpathy)
- Deployment on edge devices with restricted memory and computational capacities, for functionalities like real-time machine translation without an internet connection (the 4bit-quantized TinyLlama-1.1B's weight only takes up 550MB RAM).
- Enabling real-time dialogue generation in video games.
Moreover, our code can be a reference for enthusiasts keen on pretraining language models under 5 billion parameters without diving too early into Megatron-LM.
Below are some details of our training setup:
| 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. See Issue 27 for a minor bug |
| 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/1430k steps) |
| Natural Language to Code Ratio | 7:3 |
| Hardware | 16 A100-40G GPUs |
Our codebase supports the following features:
- multi-gpu and multi-node distributed training with FSDP.
- flash attention 2.
- fused layernorm.
- fused swiglu.
- fused cross entropy loss .
- fused rotary positional embedding.
Credit: flash attention 2, fused layernorm, fused cross entropy loss, and fused rotary positional embedding are from the FlashAttention repo. Fused swiglu is from xformers.
Thanks to those optimizations, we achieve a throughput of 24k tokens per second per A100-40G GPU, which translates to 56% model flops utilization without activation checkpointing (We expect the MFU to be even higher on A100-80G). It means you can train a chinchilla-optimal TinyLlama (1.1B param, 22B tokens) in 32 hours with 8 A100. Those optimizations also greatly reduce the memory footprint, allowing us to stuff our 1.1B model into 40GB GPU RAM and train with a per-gpu batch size of 16k tokens. You can also pretrain TinyLlama on 3090/4090 GPUs with a smaller per-gpu batch size. Below is a comparison of the training speed of our codebase with that of Pythia and MPT.
| Model | A100 GPU hours taken on 300B tokens |
|---|---|
| TinyLlama-1.1B | 3456 |
| Pythia-1.0B | 4830 |
| MPT-1.3B | 7920 |
The Pythia number comes from their paper. The MPT number comes from here, in which they say MPT-1.3B " was trained on 440 A100-40GBs for about half a day" on 200B tokens.
The fact that TinyLlama is a relatively small model with grouped query attention means it is also fast during inference. Below are some throughputs that we measure:
| 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 |
Please refer to PRETRAIN.md for instructions on how to pretrain TinyLlama.
We include a simple full-parameter finetuning & inference script in sft. Our V0.1 chat model is finetuned using this script. The FT dataset we use is openassistant-guanaco. For finetuning with less than 4GB RAM, we refer you to the Qlora and bitsandbytes repo. We did not undergone extensive hyperparameter tuning nor choosing more performant FT datasets. We hope the community can explore on finetuning TinyLlama and come up with better chat models. I will include community-finetuned models in this repo.
This project is still under active development. We are a really small team. Community feedback and contributions are highly appreciated. Here are some things we plan to work on:
- 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.
This repository is built upon lit-gpt and flash-attention. Be sure to explore this fantastic open-source project if it's new to you!
@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}
}
This project is currently contributed by Peiyuan Zhang, Guangtao Zeng, Tianduo Wang and Wei Lu.
If you find our work valuable, please cite:
@online{tinyllama,
author = {Peiyuan Zhang, Guangtao Zeng, Tianduo Wang, Wei Lu},
title = {TinyLlama},
url = {https://github.com/jzhang38/TinyLlama},
year = {2023},
month = {Sep},
}
1. Why would pretraining a 1.1B model for so long make sense? Doesn't it contradict the Chinchilla Scaling Law?
Above is the training loss curve taken from the Llama 2 paper. Here I quote from that paper: "We observe that after pretraining on 2T Tokens, the models still did not show any sign of saturation". That is why we believe pretraining a 1.1B model for 3T tokens is a reasonable thing to do. Even if the loss curve does not go down eventually, we can still study the phenomenon of saturation and learn something from it.
The figure from the Pythia paper displays the LAMBADA accuracy plotted against the total training tokens (300B). The term "saturation" pertains specifically to the 70M and 160M models. Notably, even the 410M model does not saturate with 300B tokens, as it continues to show an increasing trend, similar to the trend of larger models.