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.
You can find the evaluation results of TinyLlama in EVAL.md.
We will be rolling out intermediate checkpoints following the below schedule.
Base models:
We are crafting a note offering possible explaination on why there is a significant improvement from 2T to 2.5T checkpoint (It is related to bos_id issue)
Chat models:
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:
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 GPUsOur codebase supports the following features:
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.
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.5Please 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 repos. We did not undergo extensive hyperparameter tuning nor choose 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:
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 from the StatNLP Research Group of Singapore University of Technology and Design.
If you find our work valuable, please cite:
@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}
}
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.
2. What does "saturation" mean?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.
RetroSearch is an open source project built by @garambo | Open a GitHub Issue
Search and Browse the WWW like it's 1997 | Search results from DuckDuckGo
HTML:
3.2
| Encoding:
UTF-8
| Version:
0.7.4
