©PaperWeekly Original · Author | StatNLP Institution | Singapore University of Technology and Design
The TinyLlama project aims to pre-train on 3 trillion tokens, building a Llama model with 1.1 billion parameters. With careful optimization, we only need 16 A100-40G GPUs to complete this task in 90 days🚀🚀. The model training started on 2023-09-01 and is currently ongoing.
We adopted the same architecture and tokenizer as Llama 2. This means TinyLlama can be plug-and-play in many Llama-based open-source projects. Additionally, TinyLlama has only 1.1B parameters, making it compact and suitable for various applications that require limited computation and memory usage.
Background and Motivation
1.1 Observations from Llama 1/2
The above is an excerpt of the training loss curve from the Llama 2 paper. They mentioned that even after pre-training on 2 trillion tokens, the model showed no signs of saturation. This means that the training of Llama 2 can continue, and there is still significant room for improvement. This is very different from the expectations of the previous Chinchilla Scaling Law (which suggests that the optimal training token count for a 7B model is around 140B).
A Reddit user [1] analyzed the training costs vs. Perplexity of different sizes of Llama 2 models. From the graph, it can be observed that in many cases, smaller models require less cost to achieve the same Perplexity, which further strengthens our idea:We can train smaller models on more data.
1.2 Why Train a Smaller Llama Model
-
The recently popular model inference acceleration algorithm, speculative decoding, requires a smaller model to assist in accelerating the larger model. The smallest model of Llama 2 is 7B, but we believe that using a smaller model as a speculator can make speculative decoding more practical. (Note: This method requires the tokenizers of the small and large models to be consistent, so relatively popular small models like Pythia cannot serve as speculators for Llama 2, while TinyLlama can.)
-
Smaller models can be used in more scenarios, such as running offline on mobile devices. The TinyLlama-1.1B model weights, after 4-bit quantization, only require 550MB of memory.
-
It may be used in large single-player games, as in such scenarios, sufficient memory must be reserved for the game itself, so the model needs to be as small as possible.
Release Schedule and Training Details
We will gradually release intermediate checkpoints according to the following schedule. We have also listed some baseline models for comparison.
From the above, it can be seen that TinyLlama is currently making excellent progress🎉🎉.You can also track TinyLlama’s training loss in real-time here [6].
Here are some details of our training setup:
Speed & Memory Usage
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.
With these optimizations, we can achieve 24k tokens/second/A100 training speed, which is 56% MFU (the MFU on A100-80G will be higher). This speed allows you to train a Chinchilla-optimal model with 11 billion parameters and 220 billion tokens in 32 hours using 8 A100s. These optimizations also significantly reduce memory usage, allowing us to fit an 11 billion parameter model into a 40GB GPU while maintaining a per-GPU batch size of 16k tokens. By simply reducing the batch size a bit, you can train TinyLlama on RTX 3090/4090. Below is a comparison of our codebase’s training speed with Pythia and MPT.
▲ The numbers for Pythia come from their paper. The numbers for MPT come from here [9], the author stated that MPT-1.3B “was trained on 440 A100-40GBs for about half a day” on 200B tokens.TinyLlama is a relatively small model, and we used GQA, which means it is also fast during inference. Here are some measured inference speeds:
Code Link
Our code can serve as aconcise reference for beginners to pre-training (can run on 3090/4090, with just over ten .py files). Although Megatron-LM currently occupies a significant portion of the commonly used pre-training frameworks, it is too difficult for beginners to get started. Our TinyLlama is also quite friendly for small labs. You don’t need 999, or even 99, just 16 A100s can complete the pre-training of the LLama 1.1B model in about a month (1T token).
Based on this codebase, we announced the TinyLlama plan, aiming to train a 1.1B TinyLlama on 3 trillion tokens in 90 days using 16 A100s, while our codebase provides some pre-training optimizations while maintaining simplicity, making it a good choice for beginners to get started with pre-training.
https://github.com/jzhang38/TinyLlama
References
[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
Further Reading
#Submission Channel#
Let Your Words Be Seen by More People
How can more quality content reach readers in a shorter path, reducing the cost for readers to find quality content? The answer is: people you don’t know.
There are always some people you don’t know who know what you want to know. PaperWeekly may serve as a bridge to facilitate the collision of scholars and academic inspiration from different backgrounds and directions, sparking more possibilities.
PaperWeekly encourages university laboratories or individuals to share various quality content on our platform, which can be latest paper interpretations, analysis of academic hotspots, research insights, or competition experience explanations, etc. Our only goal is to let knowledge flow.
📝 Basic Requirements for Submissions:
• The article must be an original work by the individual, not previously published in public channels. If it has been published or is pending publication on other platforms, please clearly indicate.
• Submissions are recommended to be written in markdown format, with images sent as attachments, requiring clear images without copyright issues.
• PaperWeekly respects the author’s right to attribution and will provide competitive remuneration for each accepted original submission, specifically based on the article’s readership and quality tier.
📬 Submission Channel:
• Submission Email:[email protected]
• Please include your immediate contact information (WeChat) in your submission so we can contact the author as soon as the manuscript is selected.
• You can also directly add the editor’s WeChat (pwbot02) for quick submission, with a note: name – submission.
△ Long press to add PaperWeekly editor
🔍
Now, you can also find us on “Zhihu”
Search for “PaperWeekly” on the Zhihu homepage
Click “Follow” to subscribe to our column
··
