# 第 7 课：Pretrain 训练入口如何串起来
这一课只解决一个问题：`train_pretrain.py` 是如何把 tokenizer、模型、数据、loss、反向传播和保存权重串成一次预训练的。
## 目录
- [0. 本节主线](#l07-mainline)
- [1. 本节要懂的 6 个原理](#l07-principles)
- [2. 变量流转](#l07-flow)
- [3. 原理一：训练脚本先把超参数变成 `args`](#l07-args)
- [4. 原理二：`MiniMindConfig` 决定模型骨架](#l07-config)
- [5. 原理三：`init_model` 同时准备 tokenizer 和 model](#l07-init-model)
- [6. 原理四：`PretrainDataset` 决定训练样本长什么样](#l07-pretrain-dataset)
- [7. 原理五：`train_epoch` 是真正训练一步的地方](#l07-train-epoch)
- [8. 原理六：checkpoint 分两类](#l07-checkpoint)
- [9. 实验验证](#l07-experiment)
- [10. 本节检查](#l07-check)
- [11. 下一课](#l07-next)
## 0. 本节主线
Pretrain 训练入口的本质是：
```text
读取命令行参数
-> 创建 MiniMindConfig
-> init_model 得到 model 和 tokenizer
-> PretrainDataset 产出 input_ids / labels
-> DataLoader 组 batch
-> model(input_ids, labels=labels) 计算 loss
-> loss.backward()
-> optimizer.step()
-> 定期保存 .pth 和 resume checkpoint
```
所以本节的核心不是“模型内部怎么计算 attention”，而是先看清楚训练脚本把哪些对象接在一起。
## 1. 本节要懂的 6 个原理
| 原理 | 要理解什么 | 源码证据 |
|---|---|---|
| 训练脚本先把超参数变成 `args` | 命令行参数决定模型尺寸、数据路径、训练轮数、保存位置 | `trainer/train_pretrain.py:83-107` |
| `MiniMindConfig` 决定模型骨架 | `hidden_size`、`num_hidden_layers`、`use_moe` 会进入模型构造 | `trainer/train_pretrain.py:114-117` |
| `init_model` 同时准备 tokenizer 和 model | tokenizer 负责数据编码，model 负责 forward/loss | `trainer/train_pretrain.py:133-138`, `trainer/trainer_utils.py:119-131` |
| `PretrainDataset` 决定训练样本长什么样 | 每条样本返回 `(input_ids, labels)` | `dataset/lm_dataset.py:37-55` |
| `train_epoch` 是真正训练一步的地方 | forward、loss、backward、clip、step 都在这里发生 | `trainer/train_pretrain.py:24-49` |
| checkpoint 分两类 | `out/*.pth` 给后续加载权重，`checkpoints/*_resume.pth` 给断点续训 | `trainer/train_pretrain.py:61-70`, `trainer/trainer_utils.py:73-116` |
学完本节，你应该能把预训练脚本读成一条链路，而不是一堆分散的配置项。
## 2. 变量流转
从命令行到一次参数更新：
```text
args.hidden_size / args.num_hidden_layers
-> lm_config
-> init_model(lm_config, args.from_weight)
-> model, tokenizer
args.data_path / args.max_seq_len
-> PretrainDataset(...)
-> DataLoader(...)
-> input_ids, labels
input_ids, labels
-> model(input_ids, labels=labels)
-> res.loss
-> loss.backward()
-> optimizer.step()
```
要特别注意这三个变量：
```text
input_ids: [batch_size, max_seq_len]
labels: [batch_size, max_seq_len]
logits: [batch_size, max_seq_len, vocab_size]
```
## 3. 原理一：训练脚本先把超参数变成 `args`
### 原理讲解
训练脚本不是一上来就训练模型，而是先把所有会影响训练行为的参数集中放到 `args` 里。
这里的参数可以分成四类：
```text
模型尺寸：hidden_size, num_hidden_layers, use_moe
数据形状：data_path, max_seq_len, batch_size
优化行为：learning_rate, epochs, accumulation_steps, grad_clip
保存/恢复：save_dir, save_weight, from_weight, from_resume
```
读这一段时，不需要背每个默认值。你要理解的是：后面的 `lm_config`、`PretrainDataset`、`optimizer`、checkpoint 路径，都来自这些参数。
### 源码证据 A：命令行参数定义
文件：`trainer/train_pretrain.py:83-107`
看它是为了理解：训练脚本暴露了哪些旋钮，以及这些旋钮后面会控制什么。
源码：
```python
parser = argparse.ArgumentParser(description="MiniMind Pretraining")
parser.add_argument("--save_dir", type=str, default="../out", help="模型保存目录")
parser.add_argument('--save_weight', default='pretrain', type=str, help="保存权重的前缀名")
parser.add_argument("--epochs", type=int, default=2, help="训练轮数")
parser.add_argument("--batch_size", type=int, default=32, help="batch size")
parser.add_argument("--learning_rate", type=float, default=5e-4, help="初始学习率")
parser.add_argument('--hidden_size', default=768, type=int, help="隐藏层维度")
parser.add_argument('--num_hidden_layers', default=8, type=int, help="隐藏层数量")
parser.add_argument('--max_seq_len', default=340, type=int, help="训练的最大截断长度（中文1token≈1.5~1.7字符）")
parser.add_argument("--data_path", type=str, default="../dataset/pretrain_t2t_mini.jsonl", help="预训练数据路径")
parser.add_argument('--from_weight', default='none', type=str, help="基于哪个权重训练，为none则从头开始")
parser.add_argument('--from_resume', default=0, type=int, choices=[0, 1], help="是否自动检测&续训（0=否，1=是）")
args = parser.parse_args()
```
这段代码说明：
- `hidden_size` 和 `num_hidden_layers` 后面会进入模型配置。
- `data_path` 和 `max_seq_len` 后面会进入数据集。
- `from_weight='none'` 表示预训练默认从随机初始化开始。
- `from_resume` 不是加载普通权重，而是恢复训练状态。
### 理解到这一步就够
你应该能看到一个参数，并判断它大概会影响哪一部分：
```text
--max_seq_len 影响 Dataset 产出的序列长度；
--hidden_size 影响模型参数量；
--from_weight 影响是否从已有 .pth 加载模型权重。
```
暂时不用看：
- `wandb_project` 怎么上报日志。
- `use_compile` 的性能优化细节。
- DDP 多卡启动参数。
## 4. 原理二：`MiniMindConfig` 决定模型骨架
### 原理讲解
`MiniMindConfig` 可以理解为模型的结构说明书。
它不是权重本身，而是告诉模型：
```text
hidden_size 多宽
num_hidden_layers 有多少层
use_moe 是否用 MoE
vocab_size 输出词表多大
num_attention_heads 有多少个 attention head
```
在 `train_pretrain.py` 里，脚本只显式传了 `hidden_size`、`num_hidden_layers`、`use_moe`。其他配置使用 `MiniMindConfig` 的默认值。也就是说，训练脚本没有把所有模型细节都写在自己这里，而是把大部分结构默认值交给配置类。
### 源码证据 A：训练脚本创建配置
文件：`trainer/train_pretrain.py:114-117`
看它是为了理解：命令行参数如何变成模型配置对象。
源码：
```python
os.makedirs(args.save_dir, exist_ok=True)
lm_config = MiniMindConfig(
hidden_size=args.hidden_size,
num_hidden_layers=args.num_hidden_layers,
use_moe=bool(args.use_moe)
)
ckp_data = lm_checkpoint(lm_config, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume==1 else None
```
这段代码说明：
- `lm_config` 是后面创建模型和 checkpoint 命名的共同依据。
- `hidden_size`、`num_hidden_layers`、`use_moe` 是这一节需要先盯住的三个结构参数。
- 如果 `from_resume==1`，脚本会先尝试读取 resume checkpoint。
### 源码证据 B：配置类的默认值
文件：`model/model_minimind.py:10-27`
看它是为了理解：训练脚本没有传的结构参数从哪里来。
源码：
```python
class MiniMindConfig(PretrainedConfig):
model_type = "minimind"
def __init__(self, hidden_size=768, num_hidden_layers=8, use_moe=False, **kwargs):
super().__init__(**kwargs)
self.hidden_size = hidden_size
self.num_hidden_layers = num_hidden_layers
self.use_moe = use_moe
self.dropout = kwargs.get("dropout", 0.0)
self.vocab_size = kwargs.get("vocab_size", 6400)
self.num_attention_heads = kwargs.get("num_attention_heads", 8)
self.num_key_value_heads = kwargs.get("num_key_value_heads", 4)
self.head_dim = kwargs.get("head_dim", self.hidden_size // self.num_attention_heads)
self.intermediate_size = kwargs.get("intermediate_size", math.ceil(hidden_size * math.pi / 64) * 64)
```
这段代码说明：
- 默认词表大小是 `6400`，要和 tokenizer 的词表规模匹配。
- 默认 `head_dim = hidden_size // num_attention_heads`。
- `intermediate_size` 是由 `hidden_size` 推出来的 FFN 中间维度。
### 理解到这一步就够
你应该能说清楚：
```text
args 只是命令行参数；
lm_config 才是模型构造真正使用的结构配置。
```
暂时不用看：
- RoPE/YaRN 的具体数学。
- MoE 的 router 和 expert 细节。
- 每一层参数量的精确计算。
## 5. 原理三：`init_model` 同时准备 tokenizer 和 model
### 原理讲解
训练时需要两个东西同时存在：
```text
tokenizer: 把文本变成 input_ids
model: 接收 input_ids，输出 logits/loss
```
`init_model` 把这两个对象一起返回，是因为 Dataset 要用 tokenizer，而训练 step 要用 model。
这里还有一个关键分支：
```text
from_weight == 'none' -> 从随机初始化开始
from_weight != 'none' -> 从 out/{from_weight}_{hidden_size}.pth 加载权重
```
预训练默认 `from_weight='none'`，所以是从头学语言规律。SFT 阶段通常会从预训练权重继续。
### 源码证据 A：训练入口调用 `init_model`
文件：`trainer/train_pretrain.py:133-138`
看它是为了理解：模型、tokenizer、dataset、optimizer 在哪里被创建。
源码：
```python
model, tokenizer = init_model(lm_config, args.from_weight, device=args.device)
train_ds = PretrainDataset(args.data_path, tokenizer, max_length=args.max_seq_len)
train_sampler = DistributedSampler(train_ds) if dist.is_initialized() else None
scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype == 'float16'))
optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate)
```
这段代码说明：
- `init_model` 先给训练准备模型和 tokenizer。
- `PretrainDataset` 依赖 tokenizer，因为它要把文本编码成 token id。
- `optimizer` 直接绑定 `model.parameters()`，也就是会更新模型所有可训练参数。
### 源码证据 B：`init_model` 内部做了什么
文件：`trainer/trainer_utils.py:119-131`
看它是为了理解：加载 tokenizer、创建模型、可选加载权重这三步的顺序。
源码：
```python
def init_model(lm_config, from_weight='pretrain', tokenizer_path='../model', save_dir='../out', device='cuda'):
tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
model = MiniMindForCausalLM(lm_config)
if from_weight!= 'none':
moe_suffix = '_moe' if lm_config.use_moe else ''
weight_path = f'{save_dir}/{from_weight}_{lm_config.hidden_size}{moe_suffix}.pth'
weights = torch.load(weight_path, map_location=device)
model.load_state_dict(weights, strict=False)
get_model_params(model, lm_config)
Logger(f'Trainable Params: {sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e6:.3f}M')
return model.to(device), tokenizer
```
这段代码说明：
- tokenizer 默认从 `../model` 目录加载。
- 模型类是 `MiniMindForCausalLM`。
- `from_weight` 决定是否去 `../out` 找 `.pth` 权重。
- 返回前会打印总参数量和可训练参数量。
### 理解到这一步就够
你应该能回答：
```text
为什么本地缺少 out/pretrain_768.pth 时，from_weight='pretrain' 会失败？
因为 init_model 会拼出这个路径并 torch.load。
```
但在 `train_pretrain.py` 默认参数里：
```text
from_weight='none'
```
所以预训练默认不会加载已有权重。
暂时不用看：
- HuggingFace `AutoTokenizer` 的内部加载细节。
- `strict=False` 的所有兼容场景。
## 6. 原理四：`PretrainDataset` 决定训练样本长什么样
### 原理讲解
训练脚本本身不关心原始 jsonl 里文本怎么切。它只要求 DataLoader 每次给它：
```text
input_ids
labels
```
这两个张量都来自 `PretrainDataset`。
对预训练来说，样本逻辑很直接：
```text
读取 text
-> tokenizer 编码
-> 手动加 BOS/EOS
-> pad 到 max_length
-> labels = input_ids.clone()
-> padding label 改成 -100
```
所以预训练阶段学的是：给定前面的 token，预测下一个 token。普通文本区域基本都参与 loss，padding 不参与。
### 源码证据 A：`PretrainDataset.__getitem__`
文件：`dataset/lm_dataset.py:47-55`
看它是为了理解：一条 jsonl 文本如何变成训练脚本需要的 `(input_ids, labels)`。
源码：
```python
def __getitem__(self, index):
sample = self.samples[index]
tokens = self.tokenizer(str(sample['text']), add_special_tokens=False, max_length=self.max_length - 2, truncation=True).input_ids
tokens = [self.tokenizer.bos_token_id] + tokens + [self.tokenizer.eos_token_id]
input_ids = tokens + [self.tokenizer.pad_token_id] * (self.max_length - len(tokens))
input_ids = torch.tensor(input_ids, dtype=torch.long)
labels = input_ids.clone()
labels[input_ids == self.tokenizer.pad_token_id] = -100
return input_ids, labels
```
这段代码说明：
- `input_ids` 和 `labels` 长度一样。
- `labels = input_ids.clone()` 是预训练目标的核心。
- `-100` 只用于忽略 padding token。
- Dataset 返回的就是训练循环里 `for step, (input_ids, labels)` 接到的东西。
### 理解到这一步就够
你应该能说清楚：
```text
train_pretrain.py 没有直接读取 jsonl 内容；
它通过 PretrainDataset 间接得到已经编码好的 input_ids / labels。
```
暂时不用看：
- `datasets.load_dataset` 的缓存实现。
- 大数据集的分片和流式加载。
## 7. 原理五：`train_epoch` 是真正训练一步的地方
### 原理讲解
训练循环里最重要的动作只有五个：
```text
1. 从 loader 取 input_ids / labels
2. 把 batch 放到 device
3. forward 得到 loss
4. backward 计算梯度
5. optimizer.step 更新参数
```
你读训练循环时，不要先陷入日志、保存、wandb。先找这五个动作。
这里还有两个训练技巧：
```text
loss = loss / accumulation_steps
```
表示梯度累积时，每个小 batch 的 loss 先缩小，避免累积后梯度过大。
```text
clip_grad_norm_
```
表示更新前裁剪梯度，防止梯度爆炸。
### 源码证据 A：forward、loss、backward、step
文件：`trainer/train_pretrain.py:24-49`
看它是为了理解：一次训练 step 到底发生在哪几行。
源码：
```python
def train_epoch(epoch, loader, iters, start_step=0, wandb=None):
for step, (input_ids, labels) in enumerate(loader, start=start_step + 1):
input_ids = input_ids.to(args.device)
labels = labels.to(args.device)
lr = get_lr(epoch * iters + step, args.epochs * iters, args.learning_rate)
for param_group in optimizer.param_groups:
param_group['lr'] = lr
with autocast_ctx:
res = model(input_ids, labels=labels)
loss = res.loss + res.aux_loss
loss = loss / args.accumulation_steps
scaler.scale(loss).backward()
if step % args.accumulation_steps == 0:
scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
scaler.step(optimizer)
scaler.update()
optimizer.zero_grad(set_to_none=True)
```
这段代码说明：
- `input_ids` 和 `labels` 来自 DataLoader。
- `model(input_ids, labels=labels)` 会触发模型内部的 causal LM loss。
- `loss.backward()` 不是直接调用，而是包在 `scaler.scale(loss).backward()` 里，兼容混合精度。
- `optimizer.step()` 也通过 `scaler.step(optimizer)` 执行。
### 源码证据 B：模型 forward 计算 causal LM loss
文件：`model/model_minimind.py:245-253`
看它是为了理解：训练循环里的 `res.loss` 从哪里来。
源码：
```python
def forward(self, input_ids, attention_mask=None, past_key_values=None, use_cache=False, logits_to_keep=0, labels=None, **kwargs):
hidden_states, past_key_values, aux_loss = self.model(input_ids, attention_mask, past_key_values, use_cache, **kwargs)
slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
logits = self.lm_head(hidden_states[:, slice_indices, :])
loss = None
if labels is not None:
x, y = logits[..., :-1, :].contiguous(), labels[..., 1:].contiguous()
loss = F.cross_entropy(x.view(-1, x.size(-1)), y.view(-1), ignore_index=-100)
return MoeCausalLMOutputWithPast(loss=loss, aux_loss=aux_loss, logits=logits, past_key_values=past_key_values, hidden_states=hidden_states)
```
这段代码说明：
- `labels is not None` 时才会计算 loss。
- `logits[..., :-1, :]` 对齐 `labels[..., 1:]`，这就是 next-token prediction。
- `ignore_index=-100` 会忽略 Dataset 标出来的不训练位置。
### 理解到这一步就够
你应该能把一次训练 step 讲成：
```text
batch 进模型
-> logits 和 loss 出来
-> loss 反传成梯度
-> optimizer 根据梯度改参数
```
暂时不用看：
- AMP 的数值缩放细节。
- `aux_loss` 在 MoE 里的具体来源。
- `get_lr` 的完整调度曲线。
## 8. 原理六：checkpoint 分两类
### 原理讲解
训练脚本会保存两类文件，它们用途不同：
```text
out/pretrain_768.pth
```
这是普通模型权重，后面 SFT 或推理可以拿来加载。
```text
checkpoints/pretrain_768_resume.pth
```
这是断点续训文件，除了模型权重，还保存 optimizer、epoch、step、scaler 等训练状态。
初学时最容易混淆的是：
```text
from_weight 加载 out 里的普通权重
from_resume 加载 checkpoints 里的续训状态
```
这两个不是一回事。
### 源码证据 A：训练过程中保存普通权重和 resume checkpoint
文件：`trainer/train_pretrain.py:61-70`
看它是为了理解：一次保存动作会写出哪两类文件。
源码：
```python
if (step % args.save_interval == 0 or step == iters) and is_main_process():
model.eval()
moe_suffix = '_moe' if lm_config.use_moe else ''
ckp = f'{args.save_dir}/{args.save_weight}_{lm_config.hidden_size}{moe_suffix}.pth'
raw_model = model.module if isinstance(model, DistributedDataParallel) else model
raw_model = getattr(raw_model, '_orig_mod', raw_model)
state_dict = raw_model.state_dict()
torch.save({k: v.half().cpu() for k, v in state_dict.items()}, ckp)
lm_checkpoint(lm_config, weight=args.save_weight, model=model, optimizer=optimizer, scaler=scaler, epoch=epoch, step=step, wandb=wandb, save_dir='../checkpoints')
model.train()
```
这段代码说明：
- `torch.save(..., ckp)` 保存普通 `.pth` 权重。
- `lm_checkpoint(...)` 另外保存断点续训信息。
- 保存路径里的 `768` 来自 `lm_config.hidden_size`。
### 源码证据 B：resume checkpoint 的内容
文件：`trainer/trainer_utils.py:73-116`
看它是为了理解：为什么 resume 文件不只是模型权重。
源码：
```python
resume_data = {
'model': state_dict,
'optimizer': optimizer.state_dict(),
'epoch': epoch,
'step': step,
'world_size': dist.get_world_size() if dist.is_initialized() else 1,
'wandb_id': wandb_id
}
...
torch.save(resume_data, resume_tmp)
os.replace(resume_tmp, resume_path)
...
if os.path.exists(resume_path):
ckp_data = torch.load(resume_path, map_location='cpu')
return ckp_data
return None
```
这段代码说明：
- resume 文件保存了 `optimizer`，所以能恢复优化器状态。
- resume 文件保存了 `epoch` 和 `step`，所以能从中断位置继续。
- 普通 `.pth` 通常只够加载模型参数，不够精确恢复训练现场。
### 理解到这一步就够
你应该能区分：
```text
我要继续训练中断任务：看 from_resume / checkpoints。
我要基于已有模型开始新阶段：看 from_weight / out。
```
暂时不用看：
- 多卡 world_size 变化时 step 如何换算。
- `.tmp` 文件和 `os.replace` 的原子保存细节。
## 9. 实验验证
### 实验 A：CPU 跑一个 tiny pretrain step
这个实验验证：
```text
MiniMindConfig
-> MiniMindForCausalLM
-> PretrainDataset
-> DataLoader batch
-> forward loss
-> backward
-> optimizer.step
```
它不会训练出有效模型，只用来确认一次训练 step 的对象流转。
运行：
```bash
cd /home/sun/minimind
PYTHONDONTWRITEBYTECODE=1 python course/labs/trace_pretrain_step.py
```
记录：
```text
input_ids.shape =
labels.shape =
logits.shape =
loss =
grad_norm_before_clip =
max_first_param_delta_after_step =
```
你应该看到：
```text
input_ids.shape = (2, 48)
labels.shape = (2, 48)
logits.shape = (2, 48, 6400)
loss 是一个正数
max_first_param_delta_after_step 大于 0
```
最后一项大于 0，说明 optimizer 确实更新了模型参数。
### 实验 B：改变模型尺寸看参数量变化
运行：
```bash
PYTHONDONTWRITEBYTECODE=1 python course/labs/trace_pretrain_step.py --hidden_size 96 --num_hidden_layers 3
```
记录：
```text
params =
logits.shape =
loss =
```
你应该看到：
- `params` 变大。
- `logits.shape` 的最后一维仍然是 `vocab_size`。
- `hidden_size` 影响模型内部宽度，但最终输出仍要预测词表里的 token。
## 10. 本节检查
如果你真懂了本节，应该能不看答案说清楚：
1. `args`、`lm_config`、`model` 三者分别是什么关系。
2. `train_pretrain.py` 默认为什么是从头预训练，而不是加载已有权重。
3. `PretrainDataset` 返回的 `(input_ids, labels)` 在训练循环哪一行被接住。
4. `res.loss` 是训练脚本自己算的，还是模型 forward 里算的。
5. `from_weight` 和 `from_resume` 分别会去找哪类文件。
6. 为什么 `out/pretrain_768.pth` 和 `checkpoints/pretrain_768_resume.pth` 不能简单当成同一个东西。
## 11. 下一课
第 8 课进入 `MiniMindConfig` 与参数规模：我们会把 `hidden_size`、`num_hidden_layers`、`vocab_size`、`num_attention_heads` 和参数量关系讲清楚。