Source-linked AI summary
TinyLlama: An Open-Source Small Language Model
Peiyuan Zhang, Guangtao Zeng, Tianduo Wang, Wei Lu
TL;DR
The potential of training smaller language models on substantially larger datasets remains under-explored. TinyLlama investigates this question with a compact 1.1B-parameter Transformer trained on blended natural-language and code data using Llama 2’s architecture and tokenizer. It surpasses similarly sized open-source baselines across various downstream tasks and provides an open-source platform for language-model research and applications.
Problem
The potential of training smaller language models with substantially larger datasets remains under-explored.
Method
TinyLlama trains a 1.1B-parameter decoder-only Transformer with Llama 2’s architecture and tokenizer on blended SlimPajama and StarCoder data.
Results
TinyLlama surpasses OPT-1.3B and Pythia-1.4B across various downstream tasks and outperforms baselines on many commonsense reasoning tasks.
Takeaways & Limitations
TinyLlama’s compact architecture and promising performance support its use for end-user applications and lightweight language-model research.
Abstract
from arXiv · showhide
We present TinyLlama, a compact 1.1B language model pretrained on around 1 trillion tokens for approximately 3 epochs. Building on the architecture and tokenizer of Llama 2, TinyLlama leverages various advances contributed by the open-source community (e.g., FlashAttention and Lit-GPT), achieving better computational efficiency. Despite its relatively small size, TinyLlama demonstrates remarkable performance in a series of downstream tasks. It significantly outperforms existing open-source language models with comparable sizes. Our model checkpoints and code are publicly available on GitHub at https://github.com/jzhang38/TinyLlama.
1 Introduction
Large language models have advanced NLP through scale, but the potential of training smaller models on substantially larger datasets remains under-explored. TinyLlama addresses this gap with a 1.1B-parameter model trained on up to 3 trillion tokens and reports competitive downstream performance against similarly sized open-source models.
- Scaling language model size has driven recent NLP progress and broad task effectiveness.
- Smaller models trained with larger datasets remain less studied than large-model approaches.Inference-optimal training emphasizes more tokens than standard scaling laws recommend, and smaller models can match or outperform larger counterparts.
- TinyLlama is a 1.1B-parameter decoder-only Transformer trained with up to 3 trillion tokens using Llama 2’s architecture and tokenizer.
- TinyLlama surpasses OPT-1.3B and Pythia-1.4B across various downstream tasks.
2 Pre-training
TinyLlama’s pre-training combines curated natural-language and code data with a decoder-only Transformer architecture and efficiency-oriented training optimizations. Version 1.1 further uses multi-stage specialization, reduces total training tokens to 2 trillion, and provides general, Math&Code, and Chinese variants.
- Pre-training data: Approximately 950 billion tokens from SlimPajama and StarCoder were sampled at a 7:3 ratio and processed for about three epochs, totaling 3 trillion training tokens.The combined corpus was processed with the Llama tokenizer.
- Model architecture: TinyLlama uses a decoder-only Transformer with RoPE, pre-norm RMSNorm, SwiGLU, and grouped-query attention.Grouped-query attention uses 32 query heads and 4 key-value groups to reduce memory-bandwidth overhead and speed inference.
- Speed optimization: FlashAttention-2, fused operations, xFormers, and FSDP3 improve computational throughput, training scalability, and memory efficiency.FSDP3 supports multi-GPU and multi-node training, while fused SwiGLU reduces memory footprint and enables larger batch sizes.
- Speed optimization: 24,000 tokens per second per A100-40G GPU was achieved, with training efficiency compared against Pythia-1.0B and MPT-1.3B.The comparison measures GPU hours required to train 300 billion tokens.
- Version 1.1 training process: Version 1.1 reduces pre-training from 3 trillion to 2 trillion tokens while reporting marginal downstream improvement, and supplies general, Math&Code, and Chinese variants.The variants target general applications, mathematical and coding tasks, and Chinese text understanding, respectively.
3 Results
TinyLlama is evaluated against similarly sized open-source models across commonsense reasoning, problem-solving, and Chinese understanding tasks. It outperforms baselines on many commonsense tasks, shows stronger problem-solving performance, and achieves strong Chinese-task results.
- Commonsense reasoning tasks: TinyLlama models outperform baselines on many commonsense reasoning tasks and obtain the highest averaged scores.The evaluation uses zero-shot settings across Hellaswag, OpenBookQA, WinoGrande, ARC-Easy, ARC-Challenge, BoolQ, and PIQA.
- Problem-solving tasks: TinyLlama demonstrates better problem-solving skills than existing models on the InstructEval benchmark.The benchmark covers MMLU, BBH, DROP, and HumanEval under 5-shot, 3-shot, or zero-shot evaluation settings.
- Chinese understanding tasks: TinyLlama v1.1 Chinese outperforms other models in most Chinese understanding and reasoning tasks.The evaluation covers xwinograd, xstorycloze, XNLI, and XCOPA, while TinyLlama v1.0 and Math&Code also perform strongly.
4 Conclusion
The paper presents TinyLlama as an open-source small-scale language model with compact architecture and promising performance. It releases training resources and examines multi-stage training and data scheduling to support further research.
- 4 Conclusion: TinyLlama combines a compact architecture with promising performance for potential mobile-device applications and lightweight language-model experimentation.The authors frame these uses as enabling end-user applications and testing innovative ideas.
- 4 Conclusion: The released resources include pre-training code, intermediate checkpoints, and data-processing details.The release is intended to promote transparency in open-source language-model pre-training.
- 4 Conclusion: The TinyLlama v1.1 series verifies the effectiveness of multi-stage training and data scheduling.The authors present this verification as an additional contribution of the work.
A Data sampling ratio for TinyLlama v1.1
The supplied passages identify tables for the data-sampling ratios of TinyLlama variants, including Math&Code and Chinese models, but provide no ratio values.
- A Data sampling ratio for TinyLlama v1.1: Table 6 presents the data-sampling ratio for TinyLlama-Math&Code.The passage identifies the table but does not report its numerical entries.
- A Data sampling ratio for TinyLlama v1.1: Table 7 presents the data-sampling ratio for TinyLlama-v1.1-Chinese.The passage identifies the table but does not report its numerical entries.