理解 LLM 训练，最好的入口不是一上来就看复杂的分布式训练框架，而是先看一个足够小、足够完整的训练脚本。
MiniMind 的 trainer/train_pretrain.py 就是这样的入口。它保留了现代 LLM 训练中最核心的流程：
准备参数 -> 初始化环境 -> 构造模型 -> 构造 tokenizer 和 dataset -> 创建 optimizer -> 前向计算 loss -> 反向传播 -> 梯度累积 -> 梯度裁剪 -> 参数更新 -> 保存 checkpoint
这篇文章只抓主线，不陷入所有细枝末节。
1. 训练脚本先做什么
一个训练脚本的第一步通常是解析参数。
比如：
batch_size learning_rate epochs max_seq_len hidden_size num_hidden_layers accumulation_steps grad_clip
这些参数决定了三件事：
模型有多大； 数据怎么喂； 训练怎么更新。
对于新手来说，最重要的是先分清：
| 参数类型 | 例子 | 作用 |
|---|---|---|
| 模型结构参数 | hidden_size、num_hidden_layers | 决定模型形状。 |
| 数据参数 | max_seq_len、batch_size | 决定输入长度和 batch。 |
| 优化参数 | learning_rate、epochs、grad_clip | 决定如何训练。 |
如果训练和推理时模型结构参数不一致，权重就加载不上。
2. 构造模型配置
MiniMind 会先构造 MiniMindConfig：
lm_config = MiniMindConfig(
hidden_size=args.hidden_size,
num_hidden_layers=args.num_hidden_layers,
use_moe=args.use_moe
)
这个 config 决定了后面模型的主要结构：
hidden_size num_hidden_layers num_attention_heads num_key_value_heads vocab_size intermediate_size use_moe
可以把 MiniMindConfig 理解成模型的建筑图纸。
训练脚本不会手写每一层怎么连，而是把配置交给：
model = MiniMindForCausalLM(lm_config)
模型内部再根据 config 创建 embedding、Transformer blocks、lm_head 等组件。
3. Dataset 和 DataLoader
预训练阶段用的是 PretrainDataset。
它把原始文本变成：
x = input_ids[:-1] y = input_ids[1:]
训练时 DataLoader 会不断吐出 batch：
for step, (X, Y, loss_mask) in enumerate(train_loader):
...
这里可以简单理解为：
X：模型输入 Y：预测目标 loss_mask：哪些位置参与 loss
对于 causal LM，模型训练的本质是：
用 X 的每个位置预测 Y 的每个位置。
4. forward：loss 在哪里产生
训练循环里会调用模型：
res = model(X, labels=Y) loss = res.loss
这一步做了几件事：
input_ids -> embedding -> 多层 MiniMindBlock -> final norm -> lm_head -> logits -> CrossEntropyLoss
logits 的形状通常是：
[batch_size, seq_len, vocab_size]
也就是说，每个 token 位置都会输出一个词表大小的分布。
如果 vocab_size=6400，那么每个位置都会给出 6400 个 token 的分数。
loss 会比较：
当前位置预测的下一个 token 分布 vs 真实的下一个 token id
这就是 next-token prediction。
5. 一个 batch 是一次训练所有 token 吗
这是初学时很容易误解的地方。
推理时，模型通常一个 token 一个 token 生成：
生成第 1 个 token -> 把它接回输入 -> 生成第 2 个 token -> ...
但训练时不是这样。
训练时，一个样本长度为 N，模型会并行计算所有位置的预测：
位置 0 预测位置 1 位置 1 预测位置 2 位置 2 预测位置 3 ...
所以一个 batch 的一次 forward，会同时训练很多 token 位置。
这也是为什么训练可以高效并行，而推理 decode 阶段必须自回归。
6. 梯度累积：为什么 loss = loss / accumulation_steps
MiniMind 训练代码里有一行：
loss = loss / args.accumulation_steps
这是梯度累积的标准写法。
假设显存只能放下小 batch：
micro batch size = 4
但我们希望等价于大 batch：
global batch size = 16
就可以累积 4 次梯度：
第 1 次 forward/backward：只累积梯度，不更新 第 2 次 forward/backward：只累积梯度，不更新 第 3 次 forward/backward：只累积梯度，不更新 第 4 次 forward/backward：optimizer.step()
如果每次 loss 不除以 4，那么累积后的梯度会变成原来的 4 倍。
所以要写：
loss = loss / accumulation_steps
这样 4 次 backward 累加起来，才等价于一个大 batch 的平均 loss。
7. 梯度裁剪：防止更新过猛
训练中还有一个参数：
grad_clip
它用于梯度裁剪。
深度模型训练时，偶尔会出现梯度过大的情况。如果直接用这个梯度更新参数，模型可能发生不稳定，loss 突然爆掉。
梯度裁剪的直觉是：
如果整体梯度范数超过阈值，就按比例缩小。
常见代码类似：
torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
它不是把每个梯度元素硬截断，而是控制整体梯度向量的范数。
可以理解为：
方向尽量保留； 步子不要迈太大。
8. aux_loss：MoE 的额外损失
如果模型启用了 MoE，forward 返回里可能有：
res.aux_loss
这个 loss 不是语言建模 loss，而是 MoE router 的负载均衡损失。
MoE 中有多个 expert。如果 router 总是把 token 分给同一个 expert，就会出现：
一个专家很忙； 其他专家几乎没训练； 模型容量被浪费。
aux loss 的作用是鼓励 token 更均衡地分配到不同 expert。
所以最终 loss 可能是：
语言建模 loss + MoE aux loss
如果没有启用 MoE，这部分通常就是 0。
9. checkpoint：训练产物是什么
训练一段时间后，脚本会保存权重：
out/pretrain_768.pth out/full_sft_768.pth out/agent_768.pth
这些文件只保存模型参数，不是完整工程。
推理时需要同时具备：
模型结构代码； 模型配置； tokenizer； 权重文件。
所以用 eval_llm.py 加载 out 权重时，要保证：
hidden_size 一致； num_hidden_layers 一致； use_moe 一致； tokenizer 路径正确。
否则就会出现权重 shape 不匹配。
10. 小结
train_pretrain.py 可以看成 LLM 训练主线的最小闭环：
Dataset 产出 X/Y -> 模型 forward 得到 logits -> logits 和 Y 计算 loss -> backward 得到梯度 -> 梯度累积和裁剪 -> optimizer 更新参数 -> 保存 checkpoint
先看懂这个闭环，再去看 SFT、DPO、PPO、GRPO，会轻松很多。
因为后面的训练方法虽然目标更复杂，但都绕不开这条主线：
数据如何组织； 模型如何 forward； loss 如何定义； 参数如何更新。