FareedKhan-dev/train-llm-from-scratch
GitHub: FareedKhan-dev/train-llm-from-scratch
基于 PyTorch 从零实现 Transformer 并完成 LLM 全流程训练（预训练到对齐）的开源教学项目。
Stars: 8449 | Forks: 1164
step by step reasoning ... 42
```
这里的关键技巧是**损失掩码**。在对一段对话进行编码时，我们还会构建一个 0/1 掩码，只有助手回答部分的 token（以及标志对话结束的 `<|endoftext|>`）对应的位置才为 1。这样，SFT 就能训练模型去写出答案，而不是机械地重复提示词。以下是构建 id 及其对齐掩码的确切代码：
```
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
```
以下是一段真实的渲染对话及正在运行的验证器奖励，打印自本仓库：
```
#### OUTPUT ####
rendered chat:
<|user|>
What is 13 + 29?<|endoftext|><|assistant|>
13 + 29 = 42 42 <|endoftext|>
extract_answer("42 ") -> 42.0
reward_gsm8k("42 ", 42.0) -> 1.2 # correct AND well formatted
reward_gsm8k("7 ", 42.0) -> 0.2 # wrong, but it used the format
```
而这里展示的是一行真实的打包后的 SFT 数据行，展示了如何只对助手的 token 进行训练（在这一行 512 个 token 中，有 48 个 token 的掩码为 1）：
```
#### 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.
```
这些阶段的数据准备脚本如下：
```
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`。我们从最底层开始构建。
### 多层感知机 (MLP)
MLP 是每个块中负责对每个 token 进行“思考”的部分。它接收每个 token 的向量，将其扩展到原来大小的四倍，应用 `ReLU` 激活函数，然后再将其压缩回来。这种扩展为层在投影回来之前提供了混合特征的空间。
```
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`。
### 单头注意力
注意力机制让一个 token 能够“看到”其他 token。每个头都会为输入构建三种视图：query（我在寻找什么）、key（我包含什么信息）以及 value（如果被选中，我将传递什么信息）。我们将每个 query 与每个 key 进行匹配打分，缩放这些分数，用因果掩码隐藏未来信息，通过 softmax 将分数转化为权重，最后对 value 进行加权求和。
```
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
```
因果掩码正是使其成为语言模型的关键：位置 `t` 只能关注位置 `0..t`，永远无法看到它正试图预测的未来信息。相关代码位于 `src/models/attention.py`。
### 多头注意力
一个头只能学习一种关系。我们希望有很多个头并行工作，这样模型就能同时追踪多种模式（例如，代词和它所指代的名词）。我们运行 `n_head` 个头，将它们的输出拼接起来，并将结果传入最后一个线性层。
多头注意力](images/05_multi_head_attention.png)
```
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 残差连接：我们先进行归一化，运行一个子层，然后将结果加回到输入中。这种“加回”（残差）操作使得梯度能够流经深层堆叠的网络而不至于消失。
```
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。如果我们传入了目标值，模型还会返回交叉熵损失。
```
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
```
有一个小细节值得指出：我们在目标值上使用了 `.reshape` 而不是 `.view`。因为目标的批次数据是数据集的一个非连续切片，而在 CPU 上 `.view` 拒绝处理非连续张量。`.reshape` 能处理这两种情况，且在其他方面表现完全一致。完整的模型，包括 `forward_hidden`（稍后被奖励头和价值头复用）以及 `generate` 方法，都位于 `src/models/transformer.py` 中。
构建模型时，它会打印出参数量。以下是本仓库使用的三种规模：
```
#### 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` 中设置这些值：
```
VOCAB_SIZE = 50304
CONTEXT_LENGTH = 128
N_EMBED = 128
N_HEAD = 8
N_BLOCKS = 1
```
然后运行：
```
python scripts/train_transformer.py
```
对于长时间的训练，你可以保存定期的 checkpoint，并在中断后恢复训练：
```
python scripts/train_transformer.py --checkpoint-every 1000 --keep-last 3
python scripts/train_transformer.py --resume latest
```
如果更大的配置无法装入内存，可以开启可选的内存优化选项（默认全部关闭，因此默认行为永远不会改变）：
```
python scripts/train_transformer.py --amp --grad-checkpointing --grad-accum 8
```
更大、更现代的路径是 `scripts/pretrain_base.py`。这是相同的配方，但包含了训练中型基础模型所需的功能：跨 GPU 的 DistributedDataParallel、bf16 autocast、梯度累加、带预热的余弦学习率计划以及定期保存的 checkpoint。无论是单 GPU 还是多 GPU，命令的形式都一样：
```
# one GPU
python scripts/pretrain_base.py
# both GPUs
torchrun --standalone --nproc_per_node=2 scripts/pretrain_base.py
```
循环的核心代码很简短。每一步都会获取一个批次，在 bf16 下执行前向传播，为梯度累加缩放损失，执行反向传播，裁剪梯度，最后更新优化器：
```
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()
```
在训练过程中，它会打印当前步数、损失、学习率和吞吐量；在评估阶段，它还会打印验证集的损失和 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 万参数基础模型的真实训练和验证损失曲线。损失初始值接近 `ln(vocab_size)`，大约是 10.8（这是一个进行均匀猜测的模型的损失），随着模型学习到文本的统计规律，损失会逐渐下降：
这次运行从 **11.14** 的损失开始（略高于 `ln(50304) = 10.83` 的均匀猜测基准线），经过 2000 步后降至约 **3.73 训练损失 / 3.76 验证损失**，在两张 L40 显卡上的训练吞吐量约为 **每秒 130,000 到 150,000 个 token**。曲线持续下降讲述了预训练的全部故事：模型正在慢慢地将语言的规律压缩到它的权重之中。
## 第 4 步：生成文本
训练好的模型会预测下一个 token 的概率分布。为了生成文本，我们从该分布中采样一个 token，将其追加到序列末尾，然后将这个更长的序列再次输入模型。我们重复这个过程，直到获得了足够多的 token。
```
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
```
从保存的 checkpoint 运行它：
```
python scripts/generate_text.py --model_path models/transformer_B.pt --input_text "The" --max_new_tokens 100
```
1300 万参数的模型已经能够生成真实的单词和基本正确的语法，这正是从小规模起步的鼓舞人心之处。
## 第 5 步：后训练，将基础模型转变为助手
基础模型能够续写文本，但它无法遵循指令或主动进行推理。这就需要后训练了。好消息是模型本身从未改变。我们复用完全相同的 Transformer 骨架，在每个阶段只改变两件事：数据和损失函数。
SFT 教会基础模型以对话格式进行回答。它依然是下一个 token 的预测任务，但有一点改变：损失只计算在助手的 token 上，使用的是我们在第 1 步中构建的掩码。
```
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
```
准备数据并开始训练：
```
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 损失对其进行训练，使得被选中的回答得分始终高于被拒绝的回答。
```
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()
```
```
python scripts/prepare_preference_data.py --source both
torchrun --standalone --nproc_per_node=2 scripts/train_reward.py
```
最核心的评估指标是偏好准确率，即在留出数据对中，模型给选中回答打分更高的比例。在这次基于 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 完全跳过了奖励模型和强化学习循环。它直接在偏好对上工作，比较策略模型让选中回答出现的概率（相对于冻结的 SFT 模型的参考副本）比让拒绝回答出现的概率高多少。
```
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, ...
```
```
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** 的隐式奖励准确率（即策略模型比冻结的参考模型更偏好“选中回答”的比例）。这三种目标函数的代码都位于 `src/post_training/dpo.py`。
### PPO
PPO 是经典的 RLHF 循环。对于每个提示词，模型会写出一个回答（一次 rollout），我们对它进行打分（使用奖励模型或 GSM8K 答案检查器），针对偏离参考模型的行为添加微小的逐 token 惩罚，使用 GAE 估计每个 token 的价值，并进行几次裁剪后的梯度更新步骤。
两个核心支撑组件——优势估计和裁剪策略损失——其实非常简短：
```
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, ...
```
```
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 风格的方法。它抛弃了价值网络。对于每个提示词，它会采样一整组回答，使用可验证的奖励（最终的数字是否匹配标准答案）对它们进行打分，并使用该组自身的平均值和标准差作为基准。一个回答的优势仅仅取决于它比自己所在组的表现好多少。
```
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)
```
一段简短的算术课程会首先运行，这样在模型面对完整的 GSM8K 数据之前，能获得一些非零的奖励用来学习。组相对优势、裁剪代理损失以及 k3 KL 惩罚均位于 `src/post_training/grpo.py`。
## 第 6 步：评估
能将所有阶段串联在一起的那个唯一数字，就是贪心解码下的 GSM8K 准确率。我们给模型一道数学题，让它生成回答，从 `` 标签中提取出数字，并将其与标准答案进行核对。
```
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
```
这里会生成核心的对比表格，让你能在单一维度上追踪每个阶段：基础模型、SFT、DPO、PPO 和 GRPO，它们全部采用相同的贪心解码和相同的可验证奖励：解析 `` 标签内的数字并与标准答案核对。同一条命令可以在任何 checkpoint 上运行，因此你可以为你自己的运行结果填写表格。基础模型越大、你赋予它的预训练计算量越多，这些分数就越高。
### SFT 之后发生了什么变化
SFT 最明显的效果是行为上的改变。基础模型只知道如何续写文本，所以当你给它一个问题时，它只会写出更多的文本。经过 SFT 之后，模型学会了对话格式：它会按照训练时学到的 `... ... ` 结构给出回答，而这正是奖励模型和后续强化学习阶段随后要优化的结构。这种习得的格式是之后每个阶段赖以构建的基础。你可以使用 `scripts/chat.py` 与任何阶段的 checkpoint 对话，并直观地观察这一点，这就是下一节的内容。
## 第 7 步：与模型对话
`scripts/chat.py` 会加载任何 checkpoint，从 checkpoint 本身读取模型维度，并让你能够与之对话。它对指令模型应用对话模板，或者将基础模型视为纯粹的文本续写。
```
# 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
```
文本生成复用了与训练和评估阶段相同的、经过测试的核心代码，因此你在聊天中看到的行为正是强化学习阶段所优化出的结果。采样过程由 `--temperature`、`--top_p`、`--top_k` 或 `--greedy` 控制。相关代码位于 `src/post_training/inference.py`。
## Streamlit 控制面板
如果你更喜欢点击操作而不是敲击命令，这里有一个小型控制面板。它可以启动每个阶段、实时观察损失变化、运行评估，并与 checkpoint 进行聊天：
```
pip install -e ".[ui]"
streamlit run ui/app.py
```
它为每个阶段（数据、预训练、SFT、奖励、DPO、PPO、GRPO、评估、聊天）提供了一个独立的页面，每个页面都是基于你原本需要手动编辑的相同 JSON 配置生成的表单。
## 文档站点
每个阶段在文档站点上都有更深入的解读，包含理论、图表、真实代码以及每个指标的含义：
**https://fareedkhan-dev.github.io/train-llm-from-scratch/**
要在本地运行：
```
pip install -e ".[docs]"
mkdocs serve
```
此外还有一个“基础”版块，解释了本代码假定你已经掌握的概念（分词、decoder-only Transformer、注意力、目标函数、优化以及生成）。
## 运行整个流程
一旦基础模型完成了预训练并且数据准备好了，只需一个脚本即可运行整个后训练链，并打印出跨阶段的对比表格：
```
bash scripts/run_posttraining.sh # uses both GPUs via torchrun
NPROC=1 bash scripts/run_posttraining.sh # single GPU
```
为了在几秒钟内完成的小型模型上进行快速的端到端检查，每个阶段都提供了冒烟测试配置：
```
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 万参数的模型开始，看到它能够生成合理的单词后，再通过内存优化标志逐步放大 `n_embed` 和 `n_blocks`，直到达到你的 GPU 极限。之后，一次一个阶段地走完后训练的流程，观察 GSM8K 的准确率是如何变化的。每个阶段的代码都足够精简，适合一口气读完，而且它们都共享同一个模型。
如果你想深入了解某个特定阶段，文档站点上为每一个阶段都提供了专门的页面。
想聊聊吗？[我的 LinkedIn](https://www.linkedin.com/in/fareed-khan-dev/) ## Star 历史 [](https://star-history.com/#FareedKhan-dev/train-llm-from-scratch&Date)
# 从零开始训练 LLM
   [](https://fareedkhan-dev.github.io/train-llm-from-scratch/)
**我正在寻找 AI 方面的博士职位**。[GitHub](https://github.com/FareedKhan-dev)
我使用 PyTorch 从头实现了 transformer 模型，基于论文 [Attention is All You Need](https://arxiv.org/abs/1706.03762)。你可以使用我的脚本，在单 GPU 上训练你自己的**十亿**或**百万**参数级别的 LLM。
这最初只是一个预训练教程。现在它涵盖了从原始文本到对齐的推理风格模型的全过程，每种算法都是用原生 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 万参数 LLM 的输出，让你直观感受这个小规模的起点能实现什么效果：
```
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.
```
## 目录
- [适用人群](#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)
- [多头注意力](#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 (监督微调)](#sft-supervised-fine-tuning)
- [奖励模型](#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 中的每张图表都采用了相同的颜色规范，因此每种颜色都有特定含义：
- 绿色代表原始数据
- 青色代表存储在磁盘上的已分词数据
- 蓝色代表普通的处理步骤
- 黄色代表模型或训练步骤
- 橙色代表强化学习与奖励部分
- 红色代表损失
- 灰色代表已保存的 checkpoint
- 紫色代表最终输出或评估结果
## 前置条件与训练时间
你需要对面向对象编程、神经网络以及 PyTorch 有基本的了解。以下是一些帮助你入门的资源：
| 主题 | 视频链接 |
|---------------------|-----------------------------------------------------------|
| OOP | [OOP 视频](https://www.youtube.com/watch?v=Ej_02ICOIgs) |
| 神经网络 | [神经网络视频](https://www.youtube.com/watch?v=Jy4wM2X21u0) |
| Pytorch | [Pytorch 视频](https://www.youtube.com/watch?v=V_xro1bcAuA) |
你需要一张 GPU 来进行训练。免费的 Colab 或 Kaggle T4 足以运行 1300 万参数的模型，但无法容纳十亿参数的模型。以下是一个大致的参考指南：
| GPU 名称 | 显存 | 2B LLM 训练 | 13M LLM 训练 | 最大实际 LLM 规模 (训练) |
|--------------------------|--------|-----------------|------------------|-----------------------------------|
| NVIDIA A100 | 40 GB | ✔ | ✔ | ~6B 到 8B |
| NVIDIA V100 | 16 GB | ✘ | ✔ | ~2B |
| NVIDIA RTX 4090 | 24 GB | ✔ | ✔ | ~4B |
| NVIDIA RTX 5090 | 32 GB | ✔ | ✔ | 已验证 13M，更大配置待定 |
| 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 |
如果大型配置出现显存溢出 (OOM)，预训练脚本提供了一些可选的标志（`--amp`、`--grad-checkpointing`、`--grad-accum`），可以大幅降低显存消耗。稍后我们会详细介绍。
## 环境配置
克隆仓库并以可编辑模式安装。可编辑安装会将 `config`、`src`、`data_loader` 和 `ui` 放入你的引用路径中，这样你就不必再手动设置 `PYTHONPATH` 了：
```
git clone https://github.com/FareedKhan-dev/train-llm-from-scratch.git
cd train-llm-from-scratch
pip install -e .
```
根据你需要的模块，还可以安装以下可选的扩展包：
```
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 上几秒钟内就能完成。
## 代码结构
```
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，并将这些 id 以训练时能快速读取的格式存储到磁盘上。这一过程我们要重复四次，分别对应我们后续要进行的四种训练。
这四种数据流分别是：
1. **预训练文本** 来自 [The Pile](https://huggingface.co/datasets/monology/pile-uncopyrighted)，以扁平的 token id 数组形式存储在 HDF5 文件中。
2. **指令数据** (Alpaca, Dolly, GSM8K) 用于 SFT，被组装成固定长度的行，并带有一个用于标识哪些 token 属于助手回答的掩码。
3. **偏好对** (Anthropic HH-RLHF, UltraFeedback) 用于奖励模型和 DPO，存储为 `{prompt, chosen, rejected}` 格式。
4. **RL 提示词** (GSM8K 和用于热身的小型算术题) 用于 PPO 和 GRPO，存储为 `{prompt, gold}` 格式。
### 分词
我们使用 OpenAI `tiktoken` 中的 `r50k_base` 分词器，这也是 GPT-3 当年使用的分词器。文本被转换为整数列表，我们在每个文档末尾追加一个特殊的 `<|endoftext|>` token（id 为 50256），以便模型学习一段文本在哪里结束，下一段从哪里开始。
对于传统路径，可以下载 The Pile 的一个切片并将其分词存入 HDF5：
```
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 数组中：
```
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 准备的验证集文件（876 万个 token）的真实预览，展示了前 10 个 id 及其解码回的文本内容：
```
#### 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 ...'
```
这就是一次输出中展示的分词核心概念：输入文本，输出扁平的整数数组，而这些整数能直接解码回原始单词。
### 对话格式与损失掩码
在预训练之后的所有阶段中，模型必须知道是谁在说话。`r50k_base` 分词器只有一个特殊 token，因此我们没有发明新的特殊 token，而是使用了纯文本的角色标记，模型只需在 SFT 阶段学习它们即可。单轮对话看起来像这样（参见 `src/post_training/chat_template.py`）：
```
<|user|>
{user content}<|endoftext|><|assistant|>
{assistant content}<|endoftext|>
```
对于数学和推理任务，我们要求助手按照固定的结构展示其思考过程，因为后续的强化学习奖励会检查 answer 标签内的数字：
```
想聊聊吗？[我的 LinkedIn](https://www.linkedin.com/in/fareed-khan-dev/) ## Star 历史 [](https://star-history.com/#FareedKhan-dev/train-llm-from-scratch&Date)
标签：DLL 劫持, Kubernetes, PyTorch, RLHF, SFT, Transformer, 人工智能, 凭据扫描, 大语言模型, 模型训练, 用户模式Hook绕过, 逆向工具