wehub-resource-sync
769 lines
36 KiB
Markdown
769 lines
36 KiB
Markdown
<!-- WEHUB_ZH_README -->
|
||
> [!NOTE]
|
||
> 本文档由 WeHub 基于上游 README 翻译整理，属于社区翻译，非官方中文文档。
|
||
> [English](./README.en.md) · [原始项目](https://github.com/FareedKhan-dev/train-llm-from-scratch) · [上游 README](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/HEAD/README.md)
|
||
> 原作者、版权与许可证归属以原始项目及本仓库 LICENSE 文件为准。
|
||
|
||
|
||
|
||
<div align="center">
|
||
|
||
<!-- omit in toc -->
|
||
# 从零训练大语言模型（LLM）
|
||
|
||
   [](https://fareedkhan-dev.github.io/train-llm-from-scratch/)
|
||
|
||
**我正在寻找人工智能方向的博士职位。** [GitHub](https://github.com/FareedKhan-dev)
|
||
|
||
</div>
|
||
|
||
我基于论文 [Attention is All You Need](https://arxiv.org/abs/1706.03762). 使用 PyTorch 从零实现了一个 Transformer 模型。你可以用我的脚本在单张 GPU 上训练自己的 **十亿** 或 **百万** 参数的大语言模型。
|
||
|
||
这最初是一份预训练教程。现在它涵盖了从原始文本到对齐的、具备推理风格模型的完整流程，每个算法都是用纯 PyTorch 手写的（没有 `trl`，没有 `peft`，没有 `transformers`）。整个旅程重复着一个核心思路：把文本变成数字，预测下一个 token，然后不断调整数据和损失函数，直到模型按我们的期望行事。
|
||
|
||
|
||
|
||
下面是我们将端到端走过的路径：
|
||
|
||
```
|
||
raw text -> tokens -> a Transformer -> next-token loss -> a base model
|
||
base model -> SFT -> Reward Model -> {PPO, DPO} -> GRPO -> evaluation and chat
|
||
```
|
||
|
||
下面是一个训练好的 1300 万参数大语言模型的输出，让你了解这套方法在小型模型端的表现：
|
||
|
||
```
|
||
In ***1978, The park was returned to the factory-plate that
|
||
the public share to the lower of the electronic fence that
|
||
follow from the Station's cities. The Canal of ancient Western
|
||
nations were confined to the city spot. The villages were directly
|
||
linked to cities in China that revolt that the US budget and in
|
||
Odambinais is uncertain and fortune established in rural areas.
|
||
```
|
||
|
||
<!-- omit in toc -->
|
||
## 目录
|
||
- [适用读者](#who-this-is-for)
|
||
- [前置要求与训练时间](#prerequisites-and-training-time)
|
||
- [环境搭建](#setup)
|
||
- [代码结构](#code-structure)
|
||
- [步骤 1：准备数据](#step-1-preparing-the-data)
|
||
- [步骤 2：由小组件构建模型](#step-2-the-model-built-from-small-pieces)
|
||
- [多层感知机（MLP）](#multi-layer-perceptron-mlp)
|
||
- [单头注意力（Single Head Attention）](#single-head-attention)
|
||
- [多头注意力（Multi Head Attention）](#multi-head-attention)
|
||
- [Transformer 块](#the-transformer-block)
|
||
- [完整 Transformer](#the-full-transformer)
|
||
- [步骤 3：预训练基座模型](#step-3-pretraining-the-base-model)
|
||
- [步骤 4：生成文本](#step-4-generating-text)
|
||
- [步骤 5：后训练——将基座模型变为助手](#step-5-post-training-turning-a-base-model-into-an-assistant)
|
||
- [SFT（Supervised Fine-Tuning，监督微调）](#sft-supervised-fine-tuning)
|
||
- [奖励模型（Reward Model）](#the-reward-model)
|
||
- [DPO、ORPO 和 KTO](#dpo-orpo-and-kto)
|
||
- [PPO](#ppo)
|
||
- [GRPO / RLVR](#grpo--rlvr)
|
||
- [步骤 6：评估](#step-6-evaluation)
|
||
- [步骤 7：与模型对话](#step-7-talking-to-the-model)
|
||
- [Streamlit 控制面板](#the-streamlit-control-panel)
|
||
- [文档站点](#the-documentation-site)
|
||
- [运行完整流程](#run-the-whole-thing)
|
||
- [下一步](#whats-next)
|
||
|
||
## 适用读者
|
||
|
||
我尽量让这一页能同时满足不同背景的读者：
|
||
|
||
- 如果你是 **学生**，请从上到下阅读。每段代码之前都有通俗说明——它做什么、为什么这样做；大多数代码块之后还附有你应该看到的输出。
|
||
- 如果你是 **开发者**，所需的命令和文件路径都在这里。你可以直接复制、运行，并阅读引用的源文件。
|
||
- 如果你是 **研究者**，后训练部分更值得一看：SFT、Bradley-Terry 奖励模型、带 GAE 的 PPO、DPO/ORPO/KTO，以及 GRPO——全部在同一小型 Transformer 上从零实现，并在真实公开数据集上训练。
|
||
|
||
本 README 中的每张图都采用统一配色，颜色含义如下：
|
||
|
||
- 绿色表示原始数据
|
||
- 青色表示磁盘上已存储、已分词的数据
|
||
- 蓝色表示普通处理步骤
|
||
- 黄色表示模型或训练步骤
|
||
- 橙色表示强化学习与奖励相关部分
|
||
- 红色表示损失（loss）
|
||
- 灰色表示已保存的检查点（checkpoint）
|
||
- 紫色表示最终输出或评估
|
||
|
||
## 前置要求与训练时间
|
||
|
||
你需要具备面向对象编程、神经网络和 PyTorch 的基础知识。以下是一些入门资源：
|
||
|
||
| 主题 | 视频链接 |
|
||
|---------------------|-----------------------------------------------------------|
|
||
| OOP（面向对象编程） | [OOP Video](https://www.youtube.com/watch?v=Ej_02ICOIgs) |
|
||
| Neural Network（神经网络） | [Neural Network Video](https://www.youtube.com/watch?v=Jy4wM2X21u0) |
|
||
| Pytorch | [Pytorch Video](https://www.youtube.com/watch?v=V_xro1bcAuA) |
|
||
|
||
训练需要 GPU。免费的 Colab 或 Kaggle T4 足以训练 1300 万参数模型，但无法容纳十亿参数模型。以下是粗略参考：
|
||
|
||
| GPU 名称 | 显存 | 20 亿参数 LLM 训练 | 1300 万参数 LLM 训练 | 最大实用 LLM 规模（训练） |
|
||
|--------------------------|--------|-----------------|------------------|-----------------------------------|
|
||
| NVIDIA A100 | 40 GB | ✔ | ✔ | ~6B 至 8B |
|
||
| NVIDIA V100 | 16 GB | ✘ | ✔ | ~2B |
|
||
| NVIDIA RTX 4090 | 24 GB | ✔ | ✔ | ~4B |
|
||
| NVIDIA RTX 5090 | 32 GB | ✔ | ✔ | 1300 万参数已验证，更大配置待定 |
|
||
| NVIDIA RTX 3090 | 24 GB | ✔ | ✔ | ~3.5B 至 4B |
|
||
| NVIDIA RTX 4080 | 16 GB | ✘ | ✔ | ~2B |
|
||
| NVIDIA RTX 4060 | 8 GB | ✘ | ✔ | ~1B |
|
||
| Tesla T4 | 16 GB | ✘ | ✔ | ~1.5B 至 2B |
|
||
|
||
若大型配置显存不足，预训练脚本提供可选标志（`--amp`、`--grad-checkpointing`、`--grad-accum`），可显著降低显存占用。后续会详细介绍。
|
||
|
||
## 环境搭建
|
||
|
||
克隆仓库并以可编辑模式安装。可编辑安装会把 `config`、`src`、`data_loader` 和 `ui` 加入导入路径，因此你不再需要手动设置 `PYTHONPATH`：
|
||
|
||
```bash
|
||
git clone https://github.com/FareedKhan-dev/train-llm-from-scratch.git
|
||
cd train-llm-from-scratch
|
||
pip install -e .
|
||
```
|
||
|
||
可按需安装可选扩展：
|
||
|
||
```bash
|
||
pip install -e ".[train]" # datasets + wandb, for downloading data and logging
|
||
pip install -e ".[ui]" # streamlit + pandas + altair, for the control panel
|
||
pip install -e ".[docs]" # mkdocs, for the documentation site
|
||
pip install -e ".[all]" # everything
|
||
```
|
||
|
||
项目有两套配置系统，一开始就分清它们会很有帮助：
|
||
|
||
- `config/config.py` 是面向旧版预训练脚本 `scripts/train_transformer.py` 的原始简易配置，采用普通 Python 常量。
|
||
- `config/post_training_config.py` 以及 `configs/` 中的 JSON 文件驱动其余所有流程（更大基座的预训练、SFT、奖励模型、DPO、PPO、GRPO）。每个阶段编辑一个小 JSON 文件即可，任何字段也可在命令行覆盖，例如 `--lr 2e-5 --batch_size 16`。
|
||
|
||
为快速验证，每个阶段都提供精简版 `configs/smoke/`，会缩小模型规模，使完整运行在 CPU 或单张 GPU 上数秒内完成。
|
||
|
||
## 代码结构
|
||
|
||
```bash
|
||
train-llm-from-scratch/
|
||
├── src/
|
||
│ ├── models/ # the Transformer, built from small pieces
|
||
│ │ ├── mlp.py # the feed-forward block
|
||
│ │ ├── attention.py # single head and multi head attention
|
||
│ │ ├── transformer_block.py # one block: attention + MLP + residuals
|
||
│ │ └── transformer.py # the full model: embeddings + blocks + lm_head
|
||
│ └── post_training/ # SFT, reward model, PPO, DPO, GRPO, eval, inference
|
||
├── config/
|
||
│ ├── config.py # legacy pretraining config (plain constants)
|
||
│ ├── post_training_config.py # dataclasses for every post-training stage
|
||
│ └── loader.py # merges defaults < base.json < stage.json < CLI
|
||
├── configs/ # editable JSON, one file per stage (+ smoke/)
|
||
├── data_loader/ # batch iterators for each kind of data
|
||
├── scripts/ # every runnable step lives here
|
||
├── ui/ # the Streamlit control panel
|
||
├── docs/ # the MkDocs site (theory + diagrams)
|
||
├── images/ # the diagrams in this README (+ the generator)
|
||
└── pyproject.toml # pip install -e .
|
||
```
|
||
|
||
## 第 1 步：准备数据
|
||
|
||
模型看到的只有整数。因此第一步永远是相同的：把文本转成 token id，并以训练时便于快速读取的格式存到磁盘。我们会做四次，对应后面要进行的几类训练。
|
||
|
||
|
||
|
||
四条数据流分别是：
|
||
|
||
1. 来自 [The Pile](https://huggingface.co/datasets/monology/pile-uncopyrighted), 的 **预训练文本**，以 token id 扁平数组形式存储在 HDF5 文件中。
|
||
2. 用于 SFT 的 **指令数据**（Alpaca、Dolly、GSM8K），打包成定长行，并附带掩码标明哪些 token 属于助手的回答。
|
||
3. 用于奖励模型和 DPO 的 **偏好对**（Anthropic HH-RLHF、UltraFeedback），存储为 `{prompt, chosen, rejected}`。
|
||
4. 用于 PPO 和 GRPO 的 **RL 提示**（GSM8K 及小规模算术热身），存储为 `{prompt, gold}`。
|
||
|
||
### 分词（Tokenization）
|
||
|
||
我们使用 OpenAI 的 `tiktoken` 中的 `r50k_base` 分词器，与 GPT-3 所用相同。文本会变成整数列表，并在每篇文档末尾追加特殊的 `<|endoftext|>` token（id 50256），以便模型学会一段文本何处结束、下一段何处开始。
|
||
|
||
|
||
|
||
旧路径：下载 The Pile 的一个切片并分词为 HDF5：
|
||
|
||
```bash
|
||
python scripts/data_download.py # downloads the validation file + 1 training shard
|
||
python scripts/data_preprocess.py # tokenizes to data/train/pile_train.h5 and data/val/pile_dev.h5
|
||
```
|
||
|
||
更新、更快的路径会流式读取并批量编码同一数据，直接写入扁平 token 数组：
|
||
|
||
```bash
|
||
python scripts/prepare_pretrain_data.py --split val --out data/pile_dev.h5
|
||
python scripts/prepare_pretrain_data.py --split train --num_shards 1 --out data/pile_train.h5
|
||
```
|
||
|
||
分词完成后，数据就是一长串整数。下面是为此 README 准备的验证文件的真实片段（8.76 million token）：前十个 id，以及它们解码回的内容：
|
||
|
||
```python
|
||
#### OUTPUT ####
|
||
dtype: int32 | shape: (8762951,) | total tokens: 8762951
|
||
first 10 token ids: [18610, 286, 3993, 3081, 319, 4088, 11, 4640, 2163, 11]
|
||
decoded back to text:
|
||
'Effect of sleep quality on memory, executive function, and language
|
||
performance in patients with refractory focal epilepsy ...'
|
||
```
|
||
|
||
这就是分词的核心：输入文本，输出整数扁平数组，而这些整数可直接解码回原文。
|
||
|
||
### 对话格式与损失掩码（loss mask）
|
||
|
||
预训练之后，模型必须知道谁在说话。`r50k_base` 分词器只有一个特殊 token，因此我们不再发明新 token，而是用纯文本角色标记，让模型在 SFT 中自行学习。单轮对话如下（见 `src/post_training/chat_template.py`）：
|
||
|
||
```
|
||
<|user|>
|
||
{user content}<|endoftext|><|assistant|>
|
||
{assistant content}<|endoftext|>
|
||
```
|
||
|
||
数学与推理任务中，我们要求助手按固定结构展示推导过程，因为后续强化学习奖励会检查答案标签内的数字：
|
||
|
||
```
|
||
<think>step by step reasoning ...</think><answer>42</answer>
|
||
```
|
||
|
||
关键技巧是 **loss mask**。编码对话时，我们同时构建 0/1 掩码：仅在助手 token（以及结束该轮的 `<|endoftext|>`）上为 1。这样 SFT 训练的是写回答，而不是复述提示。下面是构建 id 与对齐掩码的代码：
|
||
|
||
```python
|
||
def encode_chat(messages, add_generation_prompt=False):
|
||
ids, mask = [], []
|
||
for m in messages:
|
||
role = m["role"]
|
||
# Role header is always masked out (we never train the model to emit it).
|
||
header_ids = _encode_ordinary(_header_for(role))
|
||
ids.extend(header_ids)
|
||
mask.extend([0] * len(header_ids))
|
||
|
||
content_ids = _encode_ordinary(m["content"])
|
||
is_completion = role == "assistant"
|
||
ids.extend(content_ids)
|
||
mask.extend([1 if is_completion else 0] * len(content_ids)) # train on assistant only
|
||
|
||
ids.append(EOT_ID) # turn terminator
|
||
mask.append(1 if is_completion else 0) # learn to stop
|
||
return ids, mask
|
||
```
|
||
|
||
下面是本仓库打印的真实渲染对话与验证器奖励示例：
|
||
|
||
```python
|
||
#### OUTPUT ####
|
||
rendered chat:
|
||
<|user|>
|
||
What is 13 + 29?<|endoftext|><|assistant|>
|
||
<think>13 + 29 = 42</think><answer>42</answer><|endoftext|>
|
||
|
||
extract_answer("<answer>42</answer>") -> 42.0
|
||
reward_gsm8k("<answer>42</answer>", 42.0) -> 1.2 # correct AND well formatted
|
||
reward_gsm8k("<answer>7</answer>", 42.0) -> 0.2 # wrong, but it used the format
|
||
```
|
||
|
||
下面是一条真实的打包 SFT 行，展示仅助手 token 参与训练（该行 512 个 token 中，掩码为 1 的有 48 个）：
|
||
|
||
```python
|
||
#### OUTPUT ####
|
||
tokens shape: (2131, 512) | loss_mask shape: (2131, 512)
|
||
row 0: trained (mask=1) tokens = 48 / 512
|
||
row 0 decoded:
|
||
<|user|>
|
||
What is the world's oldest annual marathon based on the reference text below? ...
|
||
<|assistant|>
|
||
The Boston Marathon is the world's oldest annual marathon, beginning on April 19th 1897.
|
||
```
|
||
|
||
各阶段的数据准备脚本如下：
|
||
|
||
```bash
|
||
python scripts/prepare_sft_data.py # Alpaca + Dolly + GSM8K -> sft_packed.h5
|
||
python scripts/prepare_preference_data.py # HH-RLHF + UltraFeedback -> preferences.jsonl
|
||
python scripts/prepare_rl_prompts.py # GSM8K + arithmetic -> rl_prompts.jsonl
|
||
```
|
||
|
||
## 第 2 步：由小组件组装模型
|
||
|
||
Transformer 若写成一整块代码会显得吓人，因此我们拆成四个小组件再堆叠。每个组件都是一个小型 `nn.Module`。我们从最底层开始。
|
||
|
||
### 多层感知机（Multi Layer Perceptron，MLP）
|
||
|
||
MLP 是每个块中负责逐 token「思考」的部分。它接收每个 token 向量，先扩展到四倍维度，施加 `ReLU`，再压缩回原尺寸。扩展为投影回去之前混合特征留出空间。
|
||
|
||
|
||
|
||
```python
|
||
class MLP(nn.Module):
|
||
"""A simple Multi-Layer Perceptron with one hidden layer."""
|
||
def __init__(self, n_embed):
|
||
super().__init__()
|
||
self.hidden = nn.Linear(n_embed, 4 * n_embed) # expand to 4x
|
||
self.relu = nn.ReLU() # non-linearity
|
||
self.proj = nn.Linear(4 * n_embed, n_embed) # project back down
|
||
|
||
def forward(self, x):
|
||
x = self.relu(self.hidden(x))
|
||
x = self.proj(x)
|
||
return x
|
||
```
|
||
|
||
`__init__` 配置两个线性层与激活函数。`forward` 按顺序执行它们。输入与输出形状相同，为 `(B, T, n_embed)`，因此块可直接堆叠而无需变形。代码位于 `src/models/mlp.py`。
|
||
|
||
### 单头注意力（Single Head Attention）
|
||
|
||
注意力让某个 token 能关注其他 token。每个头从输入构建三种视图：query（我在找什么）、key（我包含什么）、value（若被选中要传递什么）。我们对每个 query 与每个 key 打分，缩放分数，用因果掩码遮蔽未来，经 softmax 将分数变为权重，再对 value 做加权求和。
|
||
|
||
|
||
|
||
```python
|
||
class Head(nn.Module):
|
||
"""A single attention head with causal masking."""
|
||
def __init__(self, head_size, n_embed, context_length):
|
||
super().__init__()
|
||
self.key = nn.Linear(n_embed, head_size, bias=False)
|
||
self.query = nn.Linear(n_embed, head_size, bias=False)
|
||
self.value = nn.Linear(n_embed, head_size, bias=False)
|
||
# a lower-triangular matrix used to mask out future positions
|
||
self.register_buffer('tril', torch.tril(torch.ones(context_length, context_length)))
|
||
|
||
def forward(self, x):
|
||
B, T, C = x.shape
|
||
k = self.key(x)
|
||
q = self.query(x)
|
||
scale_factor = 1 / math.sqrt(C)
|
||
attn_weights = q @ k.transpose(-2, -1) * scale_factor # (B, T, T) scores
|
||
attn_weights = attn_weights.masked_fill(self.tril[:T, :T] == 0, float('-inf')) # no peeking ahead
|
||
attn_weights = F.softmax(attn_weights, dim=-1)
|
||
v = self.value(x)
|
||
out = attn_weights @ v # weighted sum of values
|
||
return out
|
||
```
|
||
|
||
因果掩码（causal mask）正是使其成为语言模型的关键：位置 `t` 只能关注位置 `0..t`，而绝不能关注它试图预测的未来。代码位于 `src/models/attention.py`。
|
||
|
||
### 多头注意力（Multi Head Attention）
|
||
|
||
单个注意力头学习一种关系。我们希望并行运行多个头，以便模型能同时追踪多种模式（例如，代词及其所指的名词）。我们运行 `n_head` 个注意力头，拼接它们的输出，并将结果再通过一个线性层。
|
||
|
||
|
||
|
||
```python
|
||
class MultiHeadAttention(nn.Module):
|
||
def __init__(self, n_head, n_embed, context_length):
|
||
super().__init__()
|
||
self.heads = nn.ModuleList(
|
||
[Head(n_embed // n_head, n_embed, context_length) for _ in range(n_head)]
|
||
)
|
||
self.proj = nn.Linear(n_embed, n_embed) # mixes the heads back together
|
||
|
||
def forward(self, x):
|
||
x = torch.cat([h(x) for h in self.heads], dim=-1) # concat along the feature dim
|
||
x = self.proj(x)
|
||
return x
|
||
```
|
||
|
||
每个注意力头在大小为 `n_embed // n_head` 的较小子空间中工作，因此拼接后的结果恰好回到 `n_embed`。最终的投影层让各个头能够相互通信。
|
||
|
||
### Transformer 块
|
||
|
||
现在我们将注意力与 MLP 合并为一个块。该块使用 pre-norm 残差连接：我们先做归一化，运行一个子层，再将结果加回输入。这种「加回去」（即残差）使梯度能够流经深层堆叠而不至于消失。
|
||
|
||
|
||
|
||
```python
|
||
class Block(nn.Module):
|
||
def __init__(self, n_head, n_embed, context_length):
|
||
super().__init__()
|
||
self.ln1 = nn.LayerNorm(n_embed)
|
||
self.attn = MultiHeadAttention(n_head, n_embed, context_length)
|
||
self.ln2 = nn.LayerNorm(n_embed)
|
||
self.mlp = MLP(n_embed)
|
||
|
||
def forward(self, x):
|
||
x = x + self.attn(self.ln1(x)) # attention sub-layer + residual
|
||
x = x + self.mlp(self.ln2(x)) # MLP sub-layer + residual
|
||
return x
|
||
```
|
||
|
||
将 `x = x + self.attn(self.ln1(x))` 理解为「关注其他 token，然后将所学内容加回到自身」。MLP 那一行对逐 token 的思考采用同样的思路。代码位于 `src/models/transformer_block.py`。
|
||
|
||
### 完整 Transformer
|
||
|
||
最后我们将所有部分封装在一起。Token id 通过嵌入表变为向量，我们添加位置嵌入以使模型了解 token 顺序，运行块堆叠，最后再归一化一次，并投影到词汇表大小的分数（称为 logits）。如果传入 targets，模型还会返回交叉熵（cross-entropy）损失。
|
||
|
||
|
||
|
||
```python
|
||
class Transformer(nn.Module):
|
||
def __init__(self, n_head, n_embed, context_length, vocab_size, N_BLOCKS):
|
||
super().__init__()
|
||
self.token_embed = nn.Embedding(vocab_size, n_embed)
|
||
self.position_embed = nn.Embedding(context_length, n_embed)
|
||
self.attn_blocks = nn.ModuleList(
|
||
[Block(n_head, n_embed, context_length) for _ in range(N_BLOCKS)]
|
||
)
|
||
self.layer_norm = nn.LayerNorm(n_embed)
|
||
self.lm_head = nn.Linear(n_embed, vocab_size)
|
||
self.register_buffer('pos_idxs', torch.arange(context_length))
|
||
|
||
def forward(self, idx, targets=None):
|
||
x = self.forward_hidden(idx) # token + position embeddings, then the blocks + final norm
|
||
logits = self.lm_head(x) # (B, T, vocab_size)
|
||
loss = None
|
||
if targets is not None:
|
||
B, T, C = logits.shape
|
||
# reshape (not view): the target slice is not contiguous, so .view() fails on CPU
|
||
flat_logits = logits.reshape(B * T, C)
|
||
targets = targets.reshape(B * T).long()
|
||
loss = F.cross_entropy(flat_logits, targets)
|
||
return logits, loss
|
||
```
|
||
|
||
有一个小细节值得指出：我们在 targets 上使用 `.reshape` 而非 `.view`。target 批次是数据的一个非连续切片，而 `.view` 在 CPU 上无法处理这种情况。`.reshape` 两种情况都能处理，其他方面完全相同。完整模型（包括 `forward_hidden`——奖励头与价值头稍后也会复用——以及 `generate`）位于 `src/models/transformer.py`。
|
||
|
||
构建模型时会打印其参数量。本仓库使用的三种规模如下：
|
||
|
||
```python
|
||
#### OUTPUT ####
|
||
13M small config (n_embed=128, n_head=8, n_blocks=1): 13,142,656 params
|
||
this tutorial's base (n_embed=512, n_head=8, n_blocks=8): 77,031,552 params
|
||
post-training default (n_embed=1024, n_head=16, n_blocks=24): 406,359,168 params
|
||
```
|
||
|
||
## 第 3 步：预训练基础模型
|
||
|
||
预训练是最耗时的环节。我们读取随机的 token 窗口，要求模型在每个位置预测下一个 token，用交叉熵衡量其错误程度，并微调权重。我们将此重复数千次。
|
||
|
||
|
||
|
||
最简单的版本是原始的 `scripts/train_transformer.py`，它读取 `config/config.py` 并在单块 GPU 上训练。要训练 1300 万参数的模型，请在 `config/config.py` 中设置以下值：
|
||
|
||
```python
|
||
VOCAB_SIZE = 50304
|
||
CONTEXT_LENGTH = 128
|
||
N_EMBED = 128
|
||
N_HEAD = 8
|
||
N_BLOCKS = 1
|
||
```
|
||
|
||
then run:
|
||
|
||
```bash
|
||
python scripts/train_transformer.py
|
||
```
|
||
|
||
对于长时间训练，可以定期保存检查点并在中断后恢复：
|
||
|
||
```bash
|
||
python scripts/train_transformer.py --checkpoint-every 1000 --keep-last 3
|
||
python scripts/train_transformer.py --resume latest
|
||
```
|
||
|
||
如果更大的配置无法放入内存，可开启可选的内存节省选项（默认全部关闭，因此默认行为不会改变）：
|
||
|
||
```bash
|
||
python scripts/train_transformer.py --amp --grad-checkpointing --grad-accum 8
|
||
```
|
||
|
||
更大规模的现代路径是 `scripts/pretrain_base.py`。配方相同，但加入了训练中等规模基础模型所需的要素：跨 GPU 的 DistributedDataParallel、bf16 autocast、梯度累积、带 warmup 的余弦学习率调度，以及定期保存检查点。单 GPU 或多 GPU，命令形式相同：
|
||
|
||
```bash
|
||
# one GPU
|
||
python scripts/pretrain_base.py
|
||
# both GPUs
|
||
torchrun --standalone --nproc_per_node=2 scripts/pretrain_base.py
|
||
```
|
||
|
||
循环的核心很简单。每一步拉取一个批次，在 bf16 下运行前向传播，为梯度累积缩放损失，反向传播，裁剪梯度，并执行优化器步进：
|
||
|
||
```python
|
||
for micro in range(cfg.grad_accum):
|
||
xb, yb = next(batch_iter)
|
||
with amp_autocast(cfg.amp_dtype, ctx.device):
|
||
_, loss = model(xb, yb)
|
||
loss = loss / cfg.grad_accum # so the accumulated gradient is the full-batch mean
|
||
loss.backward()
|
||
torch.nn.utils.clip_grad_norm_(model.parameters(), cfg.grad_clip)
|
||
optimizer.step()
|
||
```
|
||
|
||
训练过程中会打印步数、损失、学习率和吞吐量；在评估时还会打印 dev 损失和峰值 GPU 内存：
|
||
|
||
```
|
||
#### OUTPUT ####
|
||
Model parameters: 77,031,552 (~77M) | world_size=2
|
||
Effective batch = 24*4*2 = 192 seqs/step
|
||
step 0 | loss 11.1393 | lr 7.50e-06 | 0 tok/s
|
||
step 20 | loss 8.6159 | lr 1.57e-04 | 148,936 tok/s
|
||
step 100 | loss 6.3108 | lr 6.00e-04 | 150,609 tok/s
|
||
[eval] step 100 | train 6.2501 | dev 6.1745
|
||
step 500 | loss 4.5317 | lr 5.39e-04 | 137,334 tok/s
|
||
[eval] step 500 | train 4.7066 | dev 4.6499
|
||
step 1000 | loss 4.0100 | lr 3.48e-04 | 131,348 tok/s
|
||
[eval] step 1000 | train 4.1200 | dev 4.1419
|
||
step 1500 | loss 3.6483 | lr 1.45e-04 | 123,781 tok/s
|
||
[eval] step 1500 | train 3.8393 | dev 3.8985
|
||
step 1900 | loss 3.7725 | lr 6.36e-05 | 151,488 tok/s
|
||
[eval] step 1900 | train 3.7345 | dev 3.7607
|
||
Done. Final checkpoint -> /ephemeral/ckpts/base_pretrained.pt
|
||
```
|
||
|
||
### 损失曲线
|
||
|
||
以下是我为本 README 在 2×L40 GPU 上训练的 7700 万参数基础模型的真实训练损失与 dev 损失。损失起始于接近 `ln(vocab_size)`，约为 10.8（均匀猜测的模型的损失），并随模型学习文本统计规律而下降：
|
||
|
||
|
||
|
||
本次运行初始损失为 **11.14**（略高于均匀猜测基线 `ln(50304) = 10.83`），在 2000 步后降至约 **3.73 训练 / 3.76 验证**，在两张 L40 上训练速度约为每秒 **130,000 至 150,000** 个 token。损失曲线下降，正是预训练（pretraining）的全部故事：模型正慢慢将语言模式压缩进其权重之中。
|
||
|
||
## 步骤 4：生成文本
|
||
|
||
训练好的模型会预测下一个 token 的概率分布。要生成文本，我们从该分布中采样一个 token，将其追加到序列末尾，再把更长的序列送回模型。如此重复，直到得到足够多的 token。
|
||
|
||
```python
|
||
def generate(self, idx, max_new_tokens):
|
||
for _ in range(max_new_tokens):
|
||
idx_cond = idx[:, -self.context_length:] # never look back further than the context window
|
||
logits, _ = self(idx_cond)
|
||
logits = logits[:, -1, :] # only the last position matters for the next token
|
||
probs = F.softmax(logits, dim=-1)
|
||
idx_next = torch.multinomial(probs, num_samples=1)
|
||
idx = torch.cat((idx, idx_next), dim=1)
|
||
return idx
|
||
```
|
||
|
||
从已保存的检查点运行：
|
||
|
||
```bash
|
||
python scripts/generate_text.py --model_path models/transformer_B.pt --input_text "The" --max_new_tokens 100
|
||
```
|
||
|
||
这个 1300 万参数的模型已经能产出真实词汇和大体正确的语法，这正是从小规模起步令人鼓舞之处。
|
||
|
||
## 步骤 5：后训练——将基座模型变成助手
|
||
|
||
基座模型可以续写文本，却无法按指令行事或有意识地推理。这需要后训练（post-training）。好消息是模型本身始终不变。我们复用完全相同的 Transformer 骨干网络，每个阶段只改两件事：数据和损失函数。
|
||
|
||
|
||
|
||
让所有这一切都能装进一个小仓库的关键设计思路是：**包装，而非重写**。教学用的 `Transformer` 只增加一个额外方法 `forward_hidden`，该方法返回 `lm_head` 之前的隐状态。奖励头、价值头以及所有对数概率计算，都围绕这一个方法组合而成。`src/models/` 中无需重写任何内容。
|
||
|
||
### SFT（监督微调，Supervised Fine-Tuning）
|
||
|
||
SFT 教会基座模型以对话格式作答。它仍是下一个 token 预测，唯有一点不同：损失只对助手 token 计算，使用的是步骤 1 中构建的掩码。
|
||
|
||
|
||
|
||
```python
|
||
def sft_loss(logits, tokens, loss_mask):
|
||
logits = logits[:, :-1, :] # predict token t+1 from position t
|
||
targets = tokens[:, 1:]
|
||
mask = loss_mask[:, 1:].to(logits.dtype)
|
||
|
||
V = logits.size(-1)
|
||
ce = F.cross_entropy(logits.reshape(-1, V).float(), targets.reshape(-1).long(), reduction="none")
|
||
ce = ce.view(targets.shape) * mask # zero out the prompt positions
|
||
return ce.sum() / mask.sum().clamp(min=1.0) # average over assistant tokens only
|
||
```
|
||
|
||
准备数据并训练：
|
||
|
||
```bash
|
||
python scripts/prepare_sft_data.py --context_length 1024
|
||
torchrun --standalone --nproc_per_node=2 scripts/train_sft.py
|
||
```
|
||
|
||
损失函数代码在 `src/post_training/sft.py`，训练器在 `scripts/train_sft.py`。
|
||
|
||
### 奖励模型
|
||
|
||
要做强化学习，就需要一个数字来衡量答案好坏。获取它的方法之一，是在人类偏好对上训练奖励模型。我们在 SFT 骨干之上加一个小型线性头，从最后一个真实 token 读出一个标量，并用 Bradley-Terry 损失训练，使被选答案的得分始终高于被拒答案。
|
||
|
||
|
||
|
||
```python
|
||
def bradley_terry_loss(chosen_rewards, rejected_rewards):
|
||
"""Mean -log sigmoid(chosen - rejected) over a batch of preference pairs."""
|
||
return -F.logsigmoid(chosen_rewards - rejected_rewards).mean()
|
||
```
|
||
|
||
```bash
|
||
python scripts/prepare_preference_data.py --source both
|
||
torchrun --standalone --nproc_per_node=2 scripts/train_reward.py
|
||
```
|
||
|
||
核心指标是偏好准确率（preference accuracy），即留出配对中被模型判为更高分的被选答案所占比例。本次在 7974 对真实偏好对上运行，达到 **0.574**（高于 0.5 的随机基线）；若使用更大模型和更多数据，该值可升至 0.65 至 0.75。
|
||
|
||
```
|
||
#### OUTPUT ####
|
||
Reward model from sft.pt | 7974 pairs | total_steps=996
|
||
[eval] step 250 | test_acc 0.539 | margin 0.006
|
||
[eval] step 750 | test_acc 0.576 | margin 0.063
|
||
Done RM. test_acc 0.574 margin 0.063 -> reward.pt
|
||
```
|
||
|
||
代码位于 `src/post_training/reward_model.py` 与 `src/post_training/reward_train.py`。
|
||
|
||
### DPO、ORPO 与 KTO
|
||
|
||
DPO 完全跳过奖励模型与 RL 循环。它直接作用于偏好对，比较策略使被选答案（相对于冻结的 SFT 参考副本）比被拒答案更可能的程度。
|
||
|
||
|
||
|
||
```python
|
||
def dpo_loss(policy_chosen_logps, policy_rejected_logps,
|
||
ref_chosen_logps, ref_rejected_logps, beta=0.1):
|
||
pi_logratios = policy_chosen_logps - policy_rejected_logps
|
||
ref_logratios = ref_chosen_logps - ref_rejected_logps
|
||
logits = pi_logratios - ref_logratios
|
||
loss = -F.logsigmoid(beta * logits).mean()
|
||
return loss, ...
|
||
```
|
||
|
||
```bash
|
||
torchrun --standalone --nproc_per_node=2 scripts/train_dpo.py --loss_type dpo
|
||
# --loss_type orpo reference free, folds SFT and alignment into one stage
|
||
# --loss_type kto works from an unpaired desirable / undesirable signal
|
||
```
|
||
|
||
本次运行中，DPO 在留出配对上达到 **0.574** 的隐式奖励准确率（implicit-reward accuracy），即策略比冻结参考更偏好被选回复的比例。三种目标函数均在 `src/post_training/dpo.py` 中。
|
||
|
||
### PPO
|
||
|
||
PPO 是经典的 RLHF 循环。对每个提示，模型生成一个答案（rollout），我们对其评分（用奖励模型或 GSM8K 答案校验器），对偏离参考模型施加按 token 计的小惩罚，用 GAE 估计每个 token 的好坏，再执行若干步裁剪梯度更新。
|
||
|
||
|
||
|
||
两个关键组件——优势估计与裁剪策略损失——都很简短：
|
||
|
||
```python
|
||
def ppo_policy_loss(new_logp, old_logp, advantages, mask, clip=0.2):
|
||
ratio = torch.exp(new_logp - old_logp)
|
||
surr1 = ratio * advantages
|
||
surr2 = torch.clamp(ratio, 1.0 - clip, 1.0 + clip) * advantages # the clip keeps the step small
|
||
loss = -masked_mean(torch.min(surr1, surr2), mask)
|
||
return loss, ...
|
||
```
|
||
|
||
```bash
|
||
python scripts/prepare_rl_prompts.py
|
||
torchrun --standalone --nproc_per_node=2 scripts/train_ppo.py --reward_source verifier
|
||
# --reward_source rm to use the trained reward model instead of the answer checker
|
||
```
|
||
|
||
Actor-critic 通过小型价值头（`src/post_training/value_head.py`）共享骨干，GAE、裁剪策略损失与裁剪价值损失均在 `src/post_training/ppo.py` 中。
|
||
|
||
### GRPO / RLVR
|
||
|
||
GRPO 是 2025 年、DeepSeek-R1 风格的方法。它舍弃价值网络。对每个提示，它采样一整组答案，用可验证奖励（最终数字是否与标准答案一致）评分，并以该组自身的均值与标准差作为基线。某答案的优势，就是它比同组其他答案好多少。
|
||
|
||
|
||
|
||
```python
|
||
def group_advantages(rewards, group_size, eps=1e-4):
|
||
r = rewards.view(-1, group_size)
|
||
mean = r.mean(dim=1, keepdim=True)
|
||
std = r.std(dim=1, keepdim=True)
|
||
adv = (r - mean) / (std + eps) # how much better than the rest of my group
|
||
return adv.reshape(-1)
|
||
```
|
||
|
||
```bash
|
||
torchrun --standalone --nproc_per_node=2 scripts/train_grpo.py --group_size 8
|
||
```
|
||
|
||
先运行一段简短算术课程，让模型在面对完整 GSM8K 之前能获得一些非零奖励来学习。组相对优势、裁剪替代目标与 k3 KL 惩罚均在 `src/post_training/grpo.py` 中。
|
||
|
||
## 步骤 6：评估
|
||
|
||
贯穿所有阶段的单一指标是 greedy（贪心）GSM8K 准确率。我们给模型一道数学题，让它生成回答，从 `<answer>` 标签中提取数字，并与标准答案核对。
|
||
|
||
|
||
|
||
```bash
|
||
for s in base_pretrained sft dpo ppo grpo; do
|
||
python scripts/eval_post_training.py --ckpt models/$s.pt --label $s --limit 200 --append logs/table.jsonl
|
||
done
|
||
python scripts/eval_post_training.py --table logs/table.jsonl
|
||
```
|
||
|
||
这会构建核心对比，让你能在同一坐标轴上追踪每个阶段：Base、SFT、DPO、PPO 和 GRPO，全部使用相同的 greedy decoding（贪心解码）和相同的可验证奖励：解析 `<answer>` 标签内的数字并与标准答案核对。同一命令可在任意 checkpoint 上运行，因此你可以填入自己运行的表格；基础模型越大、预训练算力越多，这些分数就越高。
|
||
|
||
### SFT 之后的变化
|
||
|
||
SFT 最明显的效果是行为变化。基础模型只知道如何续写文本，因此当你给它一个问题时，它只会继续写更多文字。SFT 之后，模型学会了对话格式：它会在训练时使用的 `<think>...</think><answer>...</answer>` 结构内作答，而这正是奖励模型和后续 RL 阶段所优化的结构。这种学到的格式是之后每个阶段所依赖的基础。你可以用 `scripts/chat.py` 与任意阶段的 checkpoint 对话，直接观察这一点——这正是下一节的内容。
|
||
|
||
## 第 7 步：与模型对话
|
||
|
||
`scripts/chat.py` 可加载任意 checkpoint，从 checkpoint 自身读取模型维度，并让你与之对话。它对指令模型应用 chat template（对话模板），或将基础模型视为原始续写。
|
||
|
||
|
||
|
||
```bash
|
||
# instruction-tuned models
|
||
python scripts/chat.py --ckpt models/sft.pt --prompt "What is 13 + 29?"
|
||
python scripts/chat.py --ckpt models/grpo.pt --prompt "..." --greedy
|
||
# base model, raw continuation
|
||
python scripts/chat.py --ckpt models/base_pretrained.pt --raw --prompt "Once upon a time"
|
||
# interactive, no --prompt
|
||
python scripts/chat.py --ckpt models/sft.pt
|
||
```
|
||
|
||
生成复用与训练和评估相同的经过验证的核心，因此你在对话中看到的内容正是 RL 阶段所优化的。采样由 `--temperature`、`--top_p`、`--top_k` 或 `--greedy` 控制。代码位于 `src/post_training/inference.py`。
|
||
|
||
## Streamlit 控制面板
|
||
|
||
如果你更愿意点击而非输入命令，有一个小型控制面板可以启动每个阶段、实时查看 loss、运行评估，并与 checkpoint 对话：
|
||
|
||
```bash
|
||
pip install -e ".[ui]"
|
||
streamlit run ui/app.py
|
||
```
|
||
|
||
每个阶段对应一页（Data、Pretrain、SFT、Reward、DPO、PPO、GRPO、Evaluate、Chat），每页都是基于你可手动编辑的同一套 JSON 配置的表单。
|
||
|
||
## 文档站点
|
||
|
||
文档站点为每个阶段提供更深入的说明，包括理论、图示、真实代码以及各指标的含义：
|
||
|
||
**https://fareedkhan-dev.github.io/train-llm-from-scratch/**
|
||
|
||
本地运行：
|
||
|
||
```bash
|
||
pip install -e ".[docs]"
|
||
mkdocs serve
|
||
```
|
||
|
||
还有一个 Foundations（基础）章节，讲解本代码假定你已了解的概念（tokenization、decoder-only Transformer、attention、objectives、optimization 和 generation）。
|
||
|
||
## 运行完整流程
|
||
|
||
基础模型完成预训练且数据准备就绪后，一个脚本即可运行整个 post-training（后训练）流程，并打印跨阶段对比表：
|
||
|
||
```bash
|
||
bash scripts/run_posttraining.sh # uses both GPUs via torchrun
|
||
NPROC=1 bash scripts/run_posttraining.sh # single GPU
|
||
```
|
||
|
||
若要对秒级完成的小模型进行快速端到端检查，每个阶段都提供了 smoke config（冒烟配置）：
|
||
|
||
```bash
|
||
python tests/test_post_training_smoke.py # core math, on CPU
|
||
python scripts/train_sft.py --config configs/smoke/sft.json # a real (tiny) training run
|
||
```
|
||
|
||
## 下一步
|
||
|
||
我建议你先训练 1300 万参数模型，观察它能产出通顺的词，然后用 memory flags 逐步放大 `n_embed` 和 `n_blocks`，直到触及 GPU 上限。之后逐阶段走完 post-training 流程，观察 GSM8K 分数的变化。每个阶段的代码都足够精简，一次通读即可，且它们共用同一套模型。
|
||
|
||
若想深入了解某个特定阶段，文档站点为每个阶段都设有专题页面。
|
||
|
||
<hr>
|
||
|
||
想聊点什么？[My Linkedin](https://www.linkedin.com/in/fareed-khan-dev/)
|
||
|
||
## Star History
|
||
|
||
[](https://star-history.com/#FareedKhan-dev/train-llm-from-scratch&Date)
|