TinyLlama pretrains a 1.1B model on 3 trillion tokens

Researchers from the StatNLP group at Singapore University of Technology and Design published TinyLlama on January 4, 2024, a language model of 1.1B parameters pretrained on 3 trillion tokens. The authors set out to see how a small model behaves when it is trained on far more tokens than the compute-optimal scaling law of Hoffmann et al. recommends for its size, and they state that this is, to their knowledge, the first attempt to train a model of around 1B parameters on that much data.

The corpus mixes natural language from SlimPajama with code from the StarCoder training data at a sampling ratio of roughly 7 to 3, which the authors put at about 950 billion tokens once the two are merged and the GitHub part of SlimPajama is dropped to avoid duplication. TinyLlama passes over that set for approximately three epochs, which is how the 3 trillion cumulative figure is reached. Pretraining ran on 16 A100-40G GPUs, and the project repository states that the run started on September 1, 2023 and was budgeted at 90 days.

TinyLlama copies the architecture and tokenizer of Meta’s Llama 2, so it carries rotary position embeddings, RMSNorm pre-normalisation and SwiGLU activations across 22 layers with a hidden size of 2048 and a 2048-token context. It also inherits grouped-query attention, where the 32 query heads share four groups of key and value representations instead of each head holding its own, which the authors use to cut memory bandwidth and speed up inference. They report a training throughput of 24,000 tokens per second per A100-40G GPU, and measure 3,456 GPU hours to cover 300 billion tokens against 4,830 hours for Pythia-1.0B and 7,920 for MPT-1.3B in their own comparison. They attribute that speed to existing open-source components rather than anything new, naming Fully Sharded Data Parallel for spreading training across nodes, FlashAttention-2, fused implementations of layer normalisation, the cross entropy loss and the rotary position embedding, and the fused SwiGLU module from xFormers.

Across seven zero-shot commonsense reasoning benchmarks the authors’ own table puts TinyLlama at an average of 52.99, ahead of OPT-1.3B at 51.44, Pythia-1.4B at 51.33 and Pythia-1.0B at 48.30, with the widest margin on HellaSwag at 59.20 against 53.65 for OPT-1.3B. It does not lead every column, and Pythia-1.4B stays ahead on BoolQ at 63.27 against 57.83. On the InstructEval suite TinyLlama averages 19.87 against 17.72 for Pythia-1.4B and 16.95 for OPT-1.3B, with most of that difference coming from HumanEval, where it scores 9.15 against 4.27.

The authors put the device case briefly, writing in the paper that the compact architecture “can enable end-user applications on mobile devices”, and the repository describes the model as suited to applications with a restricted computation and memory footprint. It puts the 4-bit quantised weights at 637 MB and names offline real-time machine translation and drafting tokens that a larger model then checks as the uses it has in mind. Neither the paper nor the repository reports a measurement on a phone. Weights and all intermediate checkpoints are on Hugging Face under Apache 2.0, with the pre-training code in the repository.