PaperlayerView exact revision on arXiv

Source-linked AI summary

TinyLlama: An Open-Source Small Language Model

Peiyuan Zhang, Guangtao Zeng, Tianduo Wang, Wei Lu

arXiv:2401.02385v2cs.CLcs.AI

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.

    Source: §1 · p. 2

  • Method

    TinyLlama trains a 1.1B-parameter decoder-only Transformer with Llama 2’s architecture and tokenizer on blended SlimPajama and StarCoder data.

    Source: §1 · p. 2, §2.1 · p. 2

  • Results

    TinyLlama surpasses OPT-1.3B and Pythia-1.4B across various downstream tasks and outperforms baselines on many commonsense reasoning tasks.

    Source: §1 · p. 2, §3 · p. 5

  • Takeaways & Limitations

    TinyLlama’s compact architecture and promising performance support its use for end-user applications and lightweight language-model research.

    Source: §4 · p. 7

Abstract

from arXiv · show

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.Source: §1 · p. 1
  • 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.Source: §1 · p. 2
  • TinyLlama is a 1.1B-parameter decoder-only Transformer trained with up to 3 trillion tokens using Llama 2’s architecture and tokenizer.Source: §1 · p. 2
  • TinyLlama surpasses OPT-1.3B and Pythia-1.4B across various downstream tasks.Source: §1 · p. 2

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.Source: §2.1 · p. 2

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.Source: §2.2 · pp. 2–3

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.Source: §2.3 · p. 3
  • 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.Source: §2.3 · p. 3, Table 2 · p. 3

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.Source: §2.5 · p. 4

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.Source: §3 · p. 5

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.Source: §3 · pp. 5–6

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.Source: §3 · p. 6

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.Source: §4 · p. 7
  • 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.Source: §4 · p. 7
  • 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.Source: §4 · p. 7

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.Source: Table 6 · p. 10
  • Table 7 presents the data-sampling ratio for TinyLlama-v1.1-Chinese.The passage identifies the table but does not report its numerical entries.Source: Table 7 · p. 10

Related papers

Loading 2401.02385v2…