The TinyLlama project is an open effort whose goal is to pretrain a 1.1B-parameter Llama model on 3 trillion tokens. The project uses exactly the same architecture and tokenizer as Llama 2, which means TinyLlama can be plugged and played into many open-source projects built on Llama. Its compact 1.1B parameter scale enables it to meet the needs of many applications with strict constraints on compute and memory usage.
Main capabilities and features:
- Pretraining goal: pretrain a 1.1B Llama model on 3 trillion tokens.
- Architecture compatibility: fully identical to the Llama 2 architecture and tokenizer, plug-and-play.
- Training efficiency: through optimization, complete training in about 90 days on 16 A100-40G GPUs.
- Code features: support FSDP multi-GPU and multi-node distributed training, Flash Attention 2, fused LayerNorm, fused cross-entropy loss, fused rotary positional embedding.
- Throughput: reaches 24k tokens per second on A100-40G, with model FLOPs utilization of 56% (without activation checkpointing).
- Memory usage: the 4bit quantized TinyLlama-1.1B weights occupy only 637 MB.
- Intermediate checkpoints: release intermediate checkpoints as planned, and record all intermediate checkpoints.
- Chat model: released the chat model TinyLlama-Chat-V0.1 fine-tuned based on OpenAssistant, and provides a simple fine-tuning script.
- Evaluation results: record evaluation results in EVAL.md.
- Speculative decoding: provides a llama.cpp speculative decoding example.
Usage:
- Pretraining: refer to the instructions in PRETRAIN.md.
- Fine-tuning: the sft directory contains simple full-parameter fine-tuning and inference scripts. The V0.1 chat model was fine-tuned using this script, and the fine-tuning dataset is openassistant-guanaco. For fine-tuning with less than 4GB of memory, refer to the Qlora and bitsandbytes repositories.
- Chat demo: provides a chat demo, allowing direct experience of TinyLlama-Chat-V0.1.
- Speculative decoding: refer to speculative_decoding/README.md.
Applicable scenarios:
- Speculative decoding to assist larger models.
- In memory andDeployment on edge devices with limited computing power, such as real-time machine translation without an internet connection.
- Implement real-time dialogue generation in video games.
- Provide reference code for enthusiasts who want to pretrain language models under 5 billion parameters but do not want to dive into Megatron-LM too early.
Limitations and notes:
- The learning rate of the base model has not yet cooled down, so it is recommended to also use the fine-tuned chat model.
- The project is still under active development, and the team is small, so community feedback and contributions are welcome.
- No extensive hyperparameter tuning has been performed, nor has a more performant fine-tuning dataset been selected.
- Planned work includes: adding pretraining scripts for other datasets, sequence length extrapolation, testing speculative decoding for Llama-2-7B, testing RTX 3090/4090 throughput, adding fine-tuning scripts, properly evaluating the model on downstream tasks, a mobile demo, and exploring retrieval augmentation.
- The project is built on lit-gpt and flash-attention.
- Contributors are from the StatNLP research group at the Singapore University of Technology and Design: Peiyuan Zhang, Guangtao Zeng, Tianduo Wang, and Wei Lu.
FAQ:
- Why pretrain on 3T tokens? Citing the Llama 2 paper, the model still showed no signs of saturation after pretraining on 2T tokens, so it was considered reasonable to pretrain a 1.1B model on 3T tokens. Even if the loss curve eventually stops decreasing, the saturation phenomenon can still be studied.
- The saturation plot in the Pythia paper shows LAMBADA accuracy versus total training tokens (300B); saturation specifically refers to the 70M and 160M models, while the 410M model also did not saturate at 300B tokens.
