TinyLlama：在万亿令牌上训练的极简语言模型

FreeGuideOnline 最新 2026-06-28

bash pip install transformers torch accelerate


### 使用 Hugging Face 加载模型
TinyLlama 的 Chat 版本针对对话场景进行了微调，使用方式与 Llama 2 类似。基础模型则适用于文本补全任务。

```python
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"

tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.float16,  # 推荐使用半精度提升速度
    device_map="auto"
)

# 构建对话提示词模板（Chat 版本需要遵循特定格式）
messages = [
    {
        "role": "system",
        "content": "You are a helpful AI assistant.",
    },
    {"role": "user", "content": "Explain quantum computing in simple terms."},
]

prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

# 生成回复
outputs = model.generate(
    **inputs,
    max_new_tokens=256,
    temperature=0.7,
    do_sample=True,
)
response = tokenizer.decode(outputs[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)
print(response)

极简推理：仅需命令行

除了 Python API，你还可以通过 Ollama 或 llama.cpp 在终端中直接与 TinyLlama 对话。

使用 Ollama（推荐初学者）

  1. 安装 Ollama：从 ollama.com 下载并安装。
  2. 拉取 TinyLlama 模型： 
    ollama pull tinyllama
  3. 运行聊天： 
    ollama run tinyllama

此时你会进入交互式命令行，直接输入问题即可获得回答，速度极快且完全离线。

使用 llama.cpp

  1. 克隆仓库并编译： 
    git clone https://github.com/ggerganov/llama.cpp
cd llama.cpp && make
  2. 下载 TinyLlama 的 GGUF 量化版本（例如 TheBloke 提供的文件）： 
    wget https://huggingface.co/TheBloke/TinyLlama-1.1B-Chat-v1.0-GGUF/resolve/main/tinyllama-1.1b-chat-v1.0.Q4_K_M.gguf
  3. 运行推理： 
    ./main -m tinyllama-1.1b-chat-v1.0.Q4_K_M.gguf -p "请介绍你自己" -n 256

深入理解 TinyLlama 的训练

数据与训练策略

TinyLlama 的成功很大程度上归功于 “数据重复与长期训练”策略。训练过程分为三个阶段：

  1. 第一阶段：使用 SlimPajama 数据集（去除低质量内容后的 RedPajama 子集），共约 1.5 万亿 token，训练了约 80% 的总步数。
  2. 第二阶段：引入 Starcoderdata（代码数据）和 SlimPajama 的 Python 子集，加强模型对编程语言的理解，此阶段训练 2000 亿 token。
  3. 第三阶段：使用高质量指令数据（如 Alpaca、Dolly、OpenOrca 等）进行监督微调，产出最终的 Chat 版本。

整个训练在 16 个节点（每个节点 4 块 A100-40G）上进行，总计约 3 万亿 token。团队没有使用任何特殊的训练技巧，仅仅依靠标准的下一个 token 预测损失，证明了只要数据足够多，小模型也能持续进步。

架构细节

  • 层数：22 层 Transformer 解码器
  • 隐藏维度：2048
  • 注意力头数：32，采用分组查询注意力（GQA），键值头数为 4
  • 上下文长度：2048（实际推理时可适当扩展，但训练时设定为 2048）
  • 词表大小：32000（基于 SentencePiece BPE）
  • 位置编码：旋转位置编码（RoPE）

这种设计完全复刻了 Llama 2 7B 的架构，仅通过减少层数和隐藏维度来降低参数量，保证了架构层面的成熟和高效。

TinyLlama 的性能表现

在多项标准基准测试中，TinyLlama 1.1B 展现出了超越 Pythia-1.4B、OPT-1.3B 等同级模型的综合能力。尤其在常识推理（HellaSwag、WinoGrande）和代码生成（HumanEval）任务上表现出色。与更大的模型相比，它在部分任务上甚至接近或超过了 Llama 2 7B 的初版。

此外，由于模型体积小，在批量推理场景下吞吐量极高，单块 A10 GPU 即可轻松支持上千并发请求，非常适合作为本地知识库、聊天机器人或智能助手的底层引擎。

微调你自己的 TinyLlama

使用 LoRA 进行高效微调

模型大小只有 1.1B，使得在普通消费级 GPU（甚至具有 8GB 显存的显卡）上就能完成全参数微调。但更节省资源的方式是使用 LoRA（低秩适配）。

以下是一个使用 PEFT 库的快速微调示例：

from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments, Trainer
from peft import LoraConfig, get_peft_model
from datasets import load_dataset

model_name = "TinyLlama/TinyLlama-1.1B-intermediate-step-1431k-3T"  # 基础预训练模型
tokenizer = AutoTokenizer.from_pretrained(model_name)
tokenizer.pad_token = tokenizer.eos_token

model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.bfloat16,
    device_map="auto"
)

# 配置 LoRA
lora_config = LoraConfig(
    r=8,
    lora_alpha=16,
    target_modules=["q_proj", "v_proj", "k_proj", "o_proj"],
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM"
)
model = get_peft_model(model, lora_config)
model.print_trainable_parameters()  # 可训练参数量极小

# 准备数据集（以 alpaca 格式为例）
dataset = load_dataset("tatsu-lab/alpaca", split="train[:1000]")

def format_prompt(example):
    return {
        "text": f"Below is an instruction that describes a task. Write a response that appropriately completes the request.\n\n### Instruction:\n{example['instruction']}\n\n### Response:\n{example['output']}"
    }

dataset = dataset.map(format_prompt)
dataset = dataset.map(lambda samples: tokenizer(samples["text"], truncation=True, max_length=512), batched=True)

# 设置训练参数
training_args = TrainingArguments(
    output_dir="./tinyllama-lora",
    per_device_train_batch_size=4,
    gradient_accumulation_steps=4,
    num_train_epochs=3,
    logging_steps=10,
    save_strategy="epoch",
    learning_rate=2e-4,
    bf16=True,
    remove_unused_columns=False,
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=dataset,
    data_collator=None,
)

trainer.train()
trainer.save_model()

微调完成后，你可以像正常模型一样加载并测试效果。

部署与优化技巧

量化压缩

为了进一步缩小模型体积、提升推理速度，TinyLlama 非常适合进行量化。社区提供了多种量化版本：

  • GPTQ：适用于 GPU，通过算法优化权值精度。
  • GGUF：适用于 CPU 推理，配合 llama.cpp 使用。
  • AWQ：最新的激活感知量化方法，保持精度同时加速。

使用 optimum 或 llama.cpp 的量化工具可以轻松将原模型（约 2.2GB）压缩至仅 600MB 左右（4-bit），性能损失极小。

使用 vLLM 或 Text Generation Inference 进行高性能部署

如果你需要构建 API 服务，可使用 vLLM 框架，它对 TinyLlama 有原生支持，能够实现极高吞吐量。

from vllm import LLM, SamplingParams

llm = LLM(model="TinyLlama/TinyLlama-1.1B-Chat-v1.0")
sampling_params = SamplingParams(temperature=0.8, top_p=0.95, max_tokens=256)
outputs = llm.generate(["你好，今天天气不错。"], sampling_params)
print(outputs[0].outputs[0].text)