许可证：apache-2.0
数据集：

• cerebras/SlimPajama-627B
  语言：
• en

TinyLlama-1.1B-v1.1

• 代码库: github.com/jzhang38/TinyLlama
• 技术报告: arxiv.org/pdf/2401.02385

我们采用了与Llama 2完全相同的架构和分词器，这意味着TinyLlama可以即插即用地应用于许多基于Llama的开源项目中。此外，TinyLlama仅有11亿参数，体积小巧，能够满足多种对计算和内存资源有限制的应用需求。

概述

在本项目中，我们不仅训练了一个单一的TinyLlama模型，还首先在1.5万亿token的语料库上训练TinyLlama，以获得基础的语言能力。随后，我们通过三种不同的数据采样方式持续预训练，将这一模型转化为三个不同的变体。这一过程的示意图如下所示。

预训练

由于这些问题（bug1、bug2），我们尝试重新训练TinyLlama以提供更好的模型。我们使用2万亿token进行训练，并将预训练分为三个阶段：1）基础预训练，2）特定领域的持续预训练，3）冷却阶段。

基础预训练

在这一初始阶段，我们仅使用SlimPajama数据集训练模型，以培养其常识推理能力。在基础预训练期间，模型训练了1.5万亿token。由于我们使用的集群每个节点配备4块A100-40G显卡，且仅在节点内分片模型权重，因此本次训练只能将批大小设置为约180万。

特定领域的持续预训练

在这一阶段，我们引入了三种不同的语料库：SlimPajama（与第一阶段相同）、数学与代码（Starcoder和Proof Pile）以及中文（Skypile）。这种方法使我们能够开发出具备专项能力的三种变体模型。

在此阶段的前约60亿token中，我们线性增加了领域特定语料库（不包括SlimPajama，其采样比例与第一阶段保持一致）的采样比例。这种逐步增加采样的策略旨在逐步调整预训练数据的分布，确保训练过程更加稳定。在采样比例调整阶段结束后，我们继续以稳定的采样策略预训练模型，直至达到约1.85万亿token。

冷却阶段

在预训练末期实施冷却阶段已成为实现更好模型收敛的关键技术。然而，由于我们在一开始就采用了余弦学习率策略，因此很难像MiniCPM或DeepSeek那样调整学习率进行冷却。因此，我们尝试通过调整批大小来实现冷却。具体来说，在冷却阶段，我们将批大小从180万增加到720万，同时保持原始的余弦学习率调度。

TinyLlama模型家族

经过广泛而详细的预训练过程后，我们现在发布了三个专项版本的模型：

1. TinyLlama_v1.1：标准版本，适用于通用场景。
2. TinyLlama_v1.1_Math&Code：具备更强的数学和代码能力。
3. TinyLlama_v1.1_Chinese：具备良好的中文理解能力。

数据

以下是各阶段的数据分布情况：

TinyLlama_v1.1

TinyLlama_v1.1_math_code

TinyLlama_v1.1_chinese

使用方法

需安装transformers>=4.31。更多信息请查看TinyLlama GitHub页面。

from transformers import AutoTokenizer  
import transformers  
import torch  

model = "TinyLlama/TinyLlama_v1.1"  
tokenizer = AutoTokenizer.from_pretrained(model)  
pipeline = transformers.pipeline(  
    "text-generation",  
    model=model,  
    torch_dtype=torch.float16,  
    device_map="auto",  
)  

sequences = pipeline(  
    'The TinyLlama project aims to pretrain a 1.1B Llama model on 3 trillion tokens. With some proper optimization, we can achieve this within a span of "just" 90 days using 16 A100-40G GPUs 🚀🚀. The training has started on 2023-09-01.',  
    do_sample=True,  
    top_k=10,  
    num_return_sequences=1,  
    repetition_penalty=1.5,  
    eos_token_id=tokenizer.eos_token_id,  
    max_length=500,  
)  
for seq in sequences:  
    print(f"结果: {seq['generated_text']}")  

评估