820 lines
21 KiB
Markdown
820 lines
21 KiB
Markdown
# 第 15 课：Checkpoint 与断点续训
|
||
|
||
这一课只解决一个问题：训练中断后，MiniMind 如何从上次的位置继续训练，而不是从头再来。
|
||
|
||
## 目录
|
||
|
||
- [0. 本节主线](#l15-mainline)
|
||
- [1. 原理讲解](#l15-principle)
|
||
- [2. 源码阅读顺序图](#l15-reading-order)
|
||
- [3. MiniMind 源码走读](#l15-source-walkthrough)
|
||
- [4. 本节必须会写 / 暂时不要求](#l15-must-write)
|
||
- [5. 手写模块](#l15-handwrite)
|
||
- [6. 对齐测试](#l15-alignment-test)
|
||
- [7. 阶段组装](#l15-stage-assembly)
|
||
- [8. 本节检查](#l15-check)
|
||
- [9. 下一课](#l15-next)
|
||
|
||
<a id="l15-mainline"></a>
|
||
## 0. 本节主线
|
||
|
||
MiniMind 的断点续训是：
|
||
|
||
```text
|
||
训练到某个 step
|
||
-> 保存普通权重文件到 out/
|
||
-> 保存 resume checkpoint 到 checkpoints/
|
||
-> checkpoint 里包含 model / optimizer / scaler / epoch / step
|
||
-> 下次启动时先查 resume checkpoint
|
||
-> 恢复 model / optimizer / scaler
|
||
-> 恢复 start_epoch / start_step
|
||
-> 用 SkipBatchSampler 跳过当前 epoch 已训练过的 batch
|
||
-> 从下一个 batch 继续训练
|
||
```
|
||
|
||
一句话：
|
||
|
||
```text
|
||
普通权重只够“推理或继续初始化”；resume checkpoint 才够“无缝续训”。
|
||
```
|
||
|
||
<a id="l15-principle"></a>
|
||
## 1. 原理讲解
|
||
|
||
### 1.1 为什么只保存 model 不够
|
||
|
||
如果只保存模型参数：
|
||
|
||
```text
|
||
model.state_dict()
|
||
```
|
||
|
||
你只能恢复“模型当前会怎么预测”。但训练状态不止模型参数。
|
||
|
||
一次训练能否无缝继续，至少依赖：
|
||
|
||
| 状态 | 形状/类型 | 为什么要保存 |
|
||
|---|---|---|
|
||
| model | 参数名到 tensor 的 dict | 恢复模型参数 |
|
||
| optimizer | dict | 恢复 AdamW 的动量、二阶矩等内部状态 |
|
||
| scaler | dict | 恢复 float16 AMP 的动态 loss scale |
|
||
| epoch | int | 知道从第几个 epoch 继续 |
|
||
| step | int | 知道当前 epoch 已训练到第几个 batch |
|
||
| world_size | int | DDP world size 变化时修正 step |
|
||
| wandb_id | str 或 None | 日志平台续接同一个 run |
|
||
|
||
对 AdamW 来说，optimizer 里不只是 learning rate，还有每个参数的动量估计。粗略写成：
|
||
|
||
$$
|
||
\begin{aligned}
|
||
m_t &= \beta_1 m_{t-1} + (1 - \beta_1)g_t \\
|
||
v_t &= \beta_2 v_{t-1} + (1 - \beta_2)g_t^2
|
||
\end{aligned}
|
||
$$
|
||
|
||
如果恢复训练时只加载 model，不加载 optimizer，那么 $m_t$ 和 $v_t$ 会丢失。训练还能跑，但不是从完全相同的训练状态继续。
|
||
|
||
### 1.2 MiniMind 保存两类文件
|
||
|
||
MiniMind 训练时会保存两类文件。
|
||
|
||
第一类是普通权重文件：
|
||
|
||
```text
|
||
../out/pretrain_768.pth
|
||
../out/full_sft_768.pth
|
||
```
|
||
|
||
它主要用于：
|
||
|
||
```text
|
||
推理；
|
||
作为下一阶段初始化权重；
|
||
手动加载 model state_dict。
|
||
```
|
||
|
||
第二类是 resume checkpoint：
|
||
|
||
```text
|
||
../checkpoints/pretrain_768_resume.pth
|
||
../checkpoints/full_sft_768_resume.pth
|
||
```
|
||
|
||
它主要用于：
|
||
|
||
```text
|
||
中断后继续训练；
|
||
恢复 optimizer / scaler / epoch / step；
|
||
跳过已经训练过的 batch。
|
||
```
|
||
|
||
路径命名规则可以写成：
|
||
|
||
$$
|
||
\begin{aligned}
|
||
\text{normal\_path}
|
||
&= \text{save\_dir}/(\text{weight}\_\text{hidden\_size}\_\text{moe?}.pth) \\
|
||
\text{resume\_path}
|
||
&= \text{save\_dir}/(\text{weight}\_\text{hidden\_size}\_\text{moe?}\_\text{resume}.pth)
|
||
\end{aligned}
|
||
$$
|
||
|
||
其中 `use_moe=True` 时，中间会多一个 `_moe`。
|
||
|
||
### 1.3 保存 checkpoint 时到底保存什么
|
||
|
||
MiniMind 的 resume checkpoint 是一个 dict，可以理解成：
|
||
|
||
```text
|
||
resume_data = {
|
||
"model": model_state_dict,
|
||
"optimizer": optimizer_state_dict,
|
||
"epoch": epoch,
|
||
"step": step,
|
||
"world_size": world_size,
|
||
"wandb_id": wandb_id,
|
||
"scaler": scaler_state_dict,
|
||
}
|
||
```
|
||
|
||
变量含义：
|
||
|
||
| key | 类型 | 含义 |
|
||
|---|---|---|
|
||
| `model` | dict[str, Tensor] | 模型参数 |
|
||
| `optimizer` | dict | 优化器状态 |
|
||
| `epoch` | int | 保存时所在 epoch |
|
||
| `step` | int | 保存时当前 epoch 的 batch step |
|
||
| `world_size` | int | 保存时 DDP 进程数 |
|
||
| `wandb_id` | str 或 None | 日志 run id |
|
||
| `scaler` | dict | AMP GradScaler 状态 |
|
||
|
||
这里的 `step` 不是全局 token 数，也不是 optimizer update 次数，而是当前 epoch 里的 dataloader step。它后面会用于跳过已经训练过的 batch。
|
||
|
||
### 1.4 为什么保存前要 unwrap model
|
||
|
||
训练时模型可能被包过：
|
||
|
||
```text
|
||
DistributedDataParallel(model)
|
||
torch.compile(model)
|
||
```
|
||
|
||
这时真正的模型可能藏在包装对象内部。
|
||
|
||
MiniMind 保存前会做：
|
||
|
||
```text
|
||
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()
|
||
```
|
||
|
||
意思是：
|
||
|
||
```text
|
||
如果是 DDP，取 model.module；
|
||
如果是 torch.compile，取 _orig_mod；
|
||
最后再取 state_dict。
|
||
```
|
||
|
||
教学版第一阶段可以先不支持 DDP 和 compile，但要知道原源码为什么这么写。
|
||
|
||
### 1.5 断点恢复时为什么要跳过 batch
|
||
|
||
假设第 0 个 epoch 训练到 `step=3` 后保存了 checkpoint。
|
||
|
||
这个 epoch 已经训练过：
|
||
|
||
```text
|
||
batch 1
|
||
batch 2
|
||
batch 3
|
||
```
|
||
|
||
恢复时如果直接重新从 dataloader 开头跑，就会重复训练这 3 个 batch。
|
||
|
||
MiniMind 的做法是：
|
||
|
||
```text
|
||
start_epoch = ckp_data["epoch"]
|
||
start_step = ckp_data["step"]
|
||
skip = start_step if epoch == start_epoch else 0
|
||
SkipBatchSampler(..., skip_batches=skip)
|
||
```
|
||
|
||
这样当前 epoch 会跳过前 `start_step` 个 batch，从下一个 batch 继续。
|
||
|
||
### 1.6 world_size 改变时为什么要修正 step
|
||
|
||
DDP 里每个进程看到的是全局数据的一部分。如果保存时 GPU 数和恢复时 GPU 数不同，同一个 `step` 对应的全局样本进度会变。
|
||
|
||
MiniMind 用一个简单修正：
|
||
|
||
$$
|
||
s_{\text{new}}
|
||
= \left\lfloor \frac{s_{\text{saved}} \cdot W_{\text{saved}}}{W_{\text{current}}} \right\rfloor
|
||
$$
|
||
|
||
其中：
|
||
|
||
| 符号 | 类型 | 含义 |
|
||
|---|---|---|
|
||
| $s_{\text{saved}}$ | int | checkpoint 里保存的 step |
|
||
| $W_{\text{saved}}$ | int | 保存时的 world size |
|
||
| $W_{\text{current}}$ | int | 当前恢复训练的 world size |
|
||
| $s_{\text{new}}$ | int | 修正后的 start_step |
|
||
|
||
这不是完美的数据级恢复，但能让恢复后的进度大致对齐。
|
||
|
||
教学版第一阶段可以不实现 world_size 修正；先实现单机单进程的恢复。
|
||
|
||
### 1.7 保存时为什么用 tmp 文件再 replace
|
||
|
||
MiniMind 的 `lm_checkpoint` 保存时会先写临时文件：
|
||
|
||
```text
|
||
xxx.pth.tmp
|
||
```
|
||
|
||
写完后再：
|
||
|
||
```text
|
||
os.replace(tmp_path, final_path)
|
||
```
|
||
|
||
这样做是为了避免保存过程中程序中断，留下一个半写入的 checkpoint。`os.replace` 在同一个文件系统里通常是原子替换：要么还是旧文件，要么变成新文件。
|
||
|
||
教学版可以复用这个习惯。
|
||
|
||
<a id="l15-reading-order"></a>
|
||
## 2. 源码阅读顺序图
|
||
|
||
这节源码按这个顺序读：
|
||
|
||
```text
|
||
train_pretrain.py 保存入口
|
||
-> trainer_utils.lm_checkpoint 保存模式
|
||
-> trainer_utils.lm_checkpoint 加载模式
|
||
-> train_pretrain.py 恢复 model / optimizer / scaler / epoch / step
|
||
-> train_pretrain.py 构造 skip
|
||
-> trainer_utils.SkipBatchSampler
|
||
```
|
||
|
||
对应文件：
|
||
|
||
```text
|
||
trainer/train_pretrain.py
|
||
trainer/trainer_utils.py
|
||
```
|
||
|
||
先看保存入口，是为了知道 checkpoint 在训练循环的哪个时机触发。
|
||
再看 `lm_checkpoint`，是为了知道普通权重和 resume checkpoint 分别保存什么。
|
||
最后看 `SkipBatchSampler`，是为了知道恢复后怎么避免重复训练已经处理过的 batch。
|
||
|
||
<a id="l15-source-walkthrough"></a>
|
||
## 3. MiniMind 源码走读
|
||
|
||
### 第 1 步：训练循环触发保存
|
||
|
||
File: `trainer/train_pretrain.py:61-70`
|
||
|
||
Read this to understand: MiniMind 会在固定间隔或 epoch 结束时保存权重和 resume checkpoint。
|
||
|
||
Code/config/template excerpt:
|
||
|
||
```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()
|
||
```
|
||
|
||
This code shows:
|
||
|
||
- 只有主进程保存，避免 DDP 多进程同时写同一个文件。
|
||
- 普通权重保存到 `args.save_dir`，默认是 `../out`。
|
||
- `lm_checkpoint` 额外保存 resume checkpoint，默认到 `../checkpoints`。
|
||
- 保存前切到 eval，保存后切回 train。
|
||
|
||
### 第 2 步：`lm_checkpoint` 构造路径
|
||
|
||
File: `trainer/trainer_utils.py:63-67`
|
||
|
||
Read this to understand: checkpoint 文件名如何由阶段名、hidden size 和 MoE 后缀组成。
|
||
|
||
Code/config/template excerpt:
|
||
|
||
```python
|
||
def lm_checkpoint(lm_config, weight='full_sft', model=None, optimizer=None, epoch=0, step=0, wandb=None, save_dir='../checkpoints', **kwargs):
|
||
os.makedirs(save_dir, exist_ok=True)
|
||
moe_path = '_moe' if lm_config.use_moe else ''
|
||
ckp_path = f'{save_dir}/{weight}_{lm_config.hidden_size}{moe_path}.pth'
|
||
resume_path = f'{save_dir}/{weight}_{lm_config.hidden_size}{moe_path}_resume.pth'
|
||
```
|
||
|
||
This code shows:
|
||
|
||
- `weight` 是阶段名，比如 `pretrain`、`full_sft`、`dpo`。
|
||
- `hidden_size` 是模型规模标识。
|
||
- `use_moe=True` 时文件名带 `_moe`。
|
||
- `resume_path` 比普通路径多 `_resume`。
|
||
|
||
### 第 3 步：保存模式写两个文件
|
||
|
||
File: `trainer/trainer_utils.py:69-106`
|
||
|
||
Read this to understand: `model is not None` 时，`lm_checkpoint` 进入保存模式。
|
||
|
||
Code/config/template excerpt:
|
||
|
||
```python
|
||
if model is not None:
|
||
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()
|
||
state_dict = {k: v.half().cpu() for k, v in state_dict.items()}
|
||
ckp_tmp = ckp_path + '.tmp'
|
||
torch.save(state_dict, ckp_tmp)
|
||
os.replace(ckp_tmp, ckp_path)
|
||
...
|
||
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
|
||
}
|
||
```
|
||
|
||
This code shows:
|
||
|
||
- 普通权重文件保存的是半精度 CPU state dict。
|
||
- resume checkpoint 也包含同一份 `model`。
|
||
- resume checkpoint 还包含 optimizer、epoch、step、world_size、wandb_id。
|
||
- `scaler` 是通过 `**kwargs` 传进来的额外状态。
|
||
|
||
### 第 4 步：额外状态通过 `kwargs` 存入 resume dict
|
||
|
||
File: `trainer/trainer_utils.py:93-104`
|
||
|
||
Read this to understand: 为什么 `scaler=scaler` 能被保存进 checkpoint。
|
||
|
||
Code/config/template excerpt:
|
||
|
||
```python
|
||
for key, value in kwargs.items():
|
||
if value is not None:
|
||
if hasattr(value, 'state_dict'):
|
||
raw_value = value.module if isinstance(value, DistributedDataParallel) else value
|
||
raw_value = getattr(raw_value, '_orig_mod', raw_value)
|
||
resume_data[key] = raw_value.state_dict()
|
||
else:
|
||
resume_data[key] = value
|
||
|
||
resume_tmp = resume_path + '.tmp'
|
||
torch.save(resume_data, resume_tmp)
|
||
os.replace(resume_tmp, resume_path)
|
||
```
|
||
|
||
This code shows:
|
||
|
||
- `scaler` 有 `state_dict()`，所以会保存为 `resume_data["scaler"]`。
|
||
- PPO/GRPO 里 scheduler、critic model 等也可以通过同样方式保存。
|
||
- resume checkpoint 也先写 tmp，再原子替换。
|
||
|
||
### 第 5 步：加载模式只查 resume checkpoint
|
||
|
||
File: `trainer/trainer_utils.py:107-116`
|
||
|
||
Read this to understand: `model is None` 时，`lm_checkpoint` 进入加载模式。
|
||
|
||
Code/config/template excerpt:
|
||
|
||
```python
|
||
else:
|
||
if os.path.exists(resume_path):
|
||
ckp_data = torch.load(resume_path, map_location='cpu')
|
||
saved_ws = ckp_data.get('world_size', 1)
|
||
current_ws = dist.get_world_size() if dist.is_initialized() else 1
|
||
if saved_ws != current_ws:
|
||
ckp_data['step'] = ckp_data['step'] * saved_ws // current_ws
|
||
Logger(f'GPU数量变化({saved_ws}→{current_ws})，step已自动转换为{ckp_data["step"]}')
|
||
return ckp_data
|
||
return None
|
||
```
|
||
|
||
This code shows:
|
||
|
||
- 自动续训只读 `_resume.pth`。
|
||
- 找不到 resume checkpoint 时返回 `None`。
|
||
- 如果 world size 变化，会修正 `step`。
|
||
- 加载位置是 CPU，后面再由 `load_state_dict` 放回模型/优化器。
|
||
|
||
### 第 6 步：训练脚本恢复状态
|
||
|
||
File: `trainer/train_pretrain.py:140-147`
|
||
|
||
Read this to understand: checkpoint 读出来以后，具体恢复哪些对象。
|
||
|
||
Code/config/template excerpt:
|
||
|
||
```python
|
||
start_epoch, start_step = 0, 0
|
||
if ckp_data:
|
||
model.load_state_dict(ckp_data['model'])
|
||
optimizer.load_state_dict(ckp_data['optimizer'])
|
||
scaler.load_state_dict(ckp_data['scaler'])
|
||
start_epoch = ckp_data['epoch']
|
||
start_step = ckp_data.get('step', 0)
|
||
```
|
||
|
||
This code shows:
|
||
|
||
- 没有 checkpoint 时从 `epoch=0, step=0` 开始。
|
||
- 有 checkpoint 时恢复 model、optimizer、scaler。
|
||
- `start_epoch` 决定从哪个 epoch 开始循环。
|
||
- `start_step` 决定当前 epoch 跳过多少 batch。
|
||
|
||
### 第 7 步：恢复时跳过已训练 batch
|
||
|
||
File: `trainer/train_pretrain.py:156-167`
|
||
|
||
Read this to understand: `start_step` 如何进入 dataloader。
|
||
|
||
Code/config/template excerpt:
|
||
|
||
```python
|
||
for epoch in range(start_epoch, args.epochs):
|
||
train_sampler and train_sampler.set_epoch(epoch)
|
||
setup_seed(42 + epoch); indices = torch.randperm(len(train_ds)).tolist()
|
||
skip = start_step if (epoch == start_epoch and start_step > 0) else 0
|
||
batch_sampler = SkipBatchSampler(train_sampler or indices, args.batch_size, skip)
|
||
loader = DataLoader(train_ds, batch_sampler=batch_sampler, num_workers=args.num_workers, pin_memory=True)
|
||
if skip > 0:
|
||
train_epoch(epoch, loader, len(loader) + skip, start_step, wandb)
|
||
else:
|
||
train_epoch(epoch, loader, len(loader), 0, wandb)
|
||
```
|
||
|
||
This code shows:
|
||
|
||
- 只有恢复后的第一个 epoch 需要 skip。
|
||
- 后续 epoch 从头正常训练。
|
||
- `len(loader) + skip` 传给 `train_epoch`，是为了让日志里的总 step 保持原 epoch 视角。
|
||
- `start_step` 传给 `train_epoch`，让 enumerate 从 `start_step + 1` 开始。
|
||
|
||
### 第 8 步：`SkipBatchSampler` 如何跳过 batch
|
||
|
||
File: `trainer/trainer_utils.py:134-157`
|
||
|
||
Read this to understand: 跳过的是 batch，不是单条样本。
|
||
|
||
Code/config/template excerpt:
|
||
|
||
```python
|
||
class SkipBatchSampler(Sampler):
|
||
def __init__(self, sampler, batch_size, skip_batches=0):
|
||
self.sampler = sampler
|
||
self.batch_size = batch_size
|
||
self.skip_batches = skip_batches
|
||
|
||
def __iter__(self):
|
||
batch = []
|
||
skipped = 0
|
||
for idx in self.sampler:
|
||
batch.append(idx)
|
||
if len(batch) == self.batch_size:
|
||
if skipped < self.skip_batches:
|
||
skipped += 1
|
||
batch = []
|
||
continue
|
||
yield batch
|
||
batch = []
|
||
if len(batch) > 0 and skipped >= self.skip_batches:
|
||
yield batch
|
||
```
|
||
|
||
This code shows:
|
||
|
||
- `sampler` 提供样本索引。
|
||
- 每攒够 `batch_size` 个索引，形成一个 batch。
|
||
- 前 `skip_batches` 个完整 batch 会被丢弃。
|
||
- 不足一个 batch 的尾巴也会 yield，前提是已经跳过够了。
|
||
|
||
<a id="l15-must-write"></a>
|
||
## 4. 本节必须会写 / 暂时不要求
|
||
|
||
必须会写：
|
||
|
||
1. checkpoint 路径命名：
|
||
|
||
$$
|
||
\begin{aligned}
|
||
\text{suffix}
|
||
&=
|
||
\begin{cases}
|
||
\text{"\_moe"}, & \text{use\_moe=True} \\
|
||
\text{""}, & \text{use\_moe=False}
|
||
\end{cases} \\
|
||
\text{ckp\_path}
|
||
&= \text{save\_dir}/(\text{weight}\_\text{hidden\_size}\text{suffix}.pth) \\
|
||
\text{resume\_path}
|
||
&= \text{save\_dir}/(\text{weight}\_\text{hidden\_size}\text{suffix}\_\text{resume}.pth)
|
||
\end{aligned}
|
||
$$
|
||
|
||
2. 保存普通权重：
|
||
|
||
```text
|
||
unwrap model
|
||
-> state_dict
|
||
-> half
|
||
-> cpu
|
||
-> torch.save
|
||
```
|
||
|
||
3. 保存 resume checkpoint：
|
||
|
||
```text
|
||
model state_dict
|
||
optimizer state_dict
|
||
scaler state_dict
|
||
epoch
|
||
step
|
||
extra_state
|
||
```
|
||
|
||
4. 加载 resume checkpoint：
|
||
|
||
```text
|
||
如果 resume_path 存在 -> torch.load(..., map_location="cpu")
|
||
否则 -> None
|
||
```
|
||
|
||
5. `SkipBatchSampler`：
|
||
|
||
```text
|
||
按 batch_size 聚合索引
|
||
跳过前 skip_batches 个 batch
|
||
yield 后面的 batch
|
||
```
|
||
|
||
暂时不要求：
|
||
|
||
```text
|
||
1. DDP world_size 修正
|
||
2. wandb_id 恢复
|
||
3. torch.compile unwrap
|
||
4. 多模型 checkpoint，例如 PPO 的 actor/critic
|
||
5. 跨机器存储和远程对象存储
|
||
```
|
||
|
||
<a id="l15-handwrite"></a>
|
||
## 5. 手写模块
|
||
|
||
本节你要补的是：
|
||
|
||
```text
|
||
course/impl/core/train_loop.py
|
||
```
|
||
|
||
### 5.1 补 `checkpoint_paths`
|
||
|
||
接口：
|
||
|
||
```python
|
||
def checkpoint_paths(
|
||
save_dir: str | Path,
|
||
weight: str,
|
||
hidden_size: int,
|
||
use_moe: bool = False,
|
||
) -> tuple[Path, Path]:
|
||
...
|
||
```
|
||
|
||
对齐源码：
|
||
|
||
```text
|
||
trainer/trainer_utils.py:63-67
|
||
```
|
||
|
||
你要实现的行为：
|
||
|
||
```text
|
||
checkpoint_paths("checkpoints", "pretrain", 768, False)
|
||
-> checkpoints/pretrain_768.pth
|
||
-> checkpoints/pretrain_768_resume.pth
|
||
|
||
checkpoint_paths("checkpoints", "pretrain", 768, True)
|
||
-> checkpoints/pretrain_768_moe.pth
|
||
-> checkpoints/pretrain_768_moe_resume.pth
|
||
```
|
||
|
||
### 5.2 补 `CourseSkipBatchSampler`
|
||
|
||
接口：
|
||
|
||
```python
|
||
class CourseSkipBatchSampler(Sampler[list[int]]):
|
||
def __iter__(self):
|
||
...
|
||
|
||
def __len__(self) -> int:
|
||
...
|
||
```
|
||
|
||
对齐源码：
|
||
|
||
```text
|
||
trainer/trainer_utils.py:134-157
|
||
```
|
||
|
||
你要实现的行为：
|
||
|
||
```text
|
||
sampler = range(10)
|
||
batch_size = 3
|
||
skip_batches = 2
|
||
|
||
原 batch:
|
||
[0, 1, 2]
|
||
[3, 4, 5]
|
||
[6, 7, 8]
|
||
[9]
|
||
|
||
输出:
|
||
[6, 7, 8]
|
||
[9]
|
||
```
|
||
|
||
### 5.3 补 `save_course_checkpoint`
|
||
|
||
接口已经在骨架里：
|
||
|
||
```python
|
||
def save_course_checkpoint(
|
||
model: torch.nn.Module,
|
||
optimizer: torch.optim.Optimizer,
|
||
*,
|
||
save_dir: str | Path,
|
||
weight: str,
|
||
hidden_size: int,
|
||
use_moe: bool,
|
||
epoch: int,
|
||
step: int,
|
||
scaler: torch.cuda.amp.GradScaler | None = None,
|
||
extra_state: dict[str, Any] | None = None,
|
||
) -> tuple[Path, Path]:
|
||
...
|
||
```
|
||
|
||
对齐源码：
|
||
|
||
```text
|
||
trainer/trainer_utils.py:69-106
|
||
```
|
||
|
||
你要实现的行为：
|
||
|
||
- 创建 `save_dir`。
|
||
- 调用 `checkpoint_paths` 得到 `ckp_path` 和 `resume_path`。
|
||
- 保存普通权重到 `ckp_path`。
|
||
- 保存 resume dict 到 `resume_path`。
|
||
- resume dict 至少包含 `model`、`optimizer`、`epoch`、`step`。
|
||
- 如果传了 `scaler`，保存 `scaler.state_dict()`。
|
||
- 如果传了 `extra_state`，把里面的 key 合并进 resume dict。
|
||
- 返回 `(ckp_path, resume_path)`。
|
||
|
||
第一版可以先不写 tmp + replace；写完基础功能后再加。
|
||
|
||
### 5.4 补 `load_course_checkpoint`
|
||
|
||
接口：
|
||
|
||
```python
|
||
def load_course_checkpoint(
|
||
save_dir: str | Path,
|
||
weight: str,
|
||
hidden_size: int,
|
||
use_moe: bool = False,
|
||
map_location: str | torch.device = "cpu",
|
||
) -> dict[str, Any] | None:
|
||
...
|
||
```
|
||
|
||
对齐源码：
|
||
|
||
```text
|
||
trainer/trainer_utils.py:107-116
|
||
```
|
||
|
||
你要实现的行为：
|
||
|
||
- 调用 `checkpoint_paths` 得到 `resume_path`。
|
||
- 如果文件不存在，返回 `None`。
|
||
- 如果文件存在，返回 `torch.load(resume_path, map_location=map_location)`。
|
||
|
||
<a id="l15-alignment-test"></a>
|
||
## 6. 对齐测试
|
||
|
||
本节新增对齐测试：
|
||
|
||
```text
|
||
course/impl/tests/test_checkpoint_resume.py
|
||
```
|
||
|
||
运行命令：
|
||
|
||
```bash
|
||
cd /home/sun/minimind
|
||
python course/impl/tests/test_checkpoint_resume.py
|
||
```
|
||
|
||
现在还没有实现时，这个测试会因为 `NotImplementedError` 失败。等你补完后，它应该打印类似：
|
||
|
||
```text
|
||
ckp_path=checkpoints/pretrain_768.pth
|
||
resume_path=checkpoints/pretrain_768_resume.pth
|
||
skipped_batches=[[6, 7, 8], [9]]
|
||
saved_ckp_exists=True
|
||
saved_resume_exists=True
|
||
checkpoint_resume=passed
|
||
```
|
||
|
||
如果这个测试通过，说明你已经掌握了：
|
||
|
||
```text
|
||
文件名规则
|
||
resume dict 的核心字段
|
||
跳过已训练 batch
|
||
保存后再加载 checkpoint
|
||
```
|
||
|
||
<a id="l15-stage-assembly"></a>
|
||
## 7. 阶段组装
|
||
|
||
本节完成后，Pretrain 阶段会多出 checkpoint/resume 能力：
|
||
|
||
```text
|
||
course/impl/core/train_loop.py::checkpoint_paths
|
||
course/impl/core/train_loop.py::save_course_checkpoint
|
||
course/impl/core/train_loop.py::load_course_checkpoint
|
||
course/impl/core/train_loop.py::CourseSkipBatchSampler
|
||
```
|
||
|
||
之后 `course/impl/train_pretrain_impl.py` 可以按这个流程组装：
|
||
|
||
```text
|
||
启动脚本
|
||
-> load_course_checkpoint(...)
|
||
-> 如果有 checkpoint，恢复 model / optimizer / scaler / epoch / step
|
||
-> 构造 CourseSkipBatchSampler
|
||
-> train_one_epoch
|
||
-> save_course_checkpoint
|
||
```
|
||
|
||
当前 Pretrain 阶段还缺：
|
||
|
||
```text
|
||
1. 教学版 CausalLM 完整 forward/loss
|
||
2. 教学版 pretrain dataset 或 dataset 适配
|
||
3. tiny pretrain 脚本组装
|
||
```
|
||
|
||
第 15 课先补齐训练状态的保存和恢复。
|
||
|
||
<a id="l15-check"></a>
|
||
## 8. 本节检查
|
||
|
||
1. 普通权重文件和 resume checkpoint 的用途有什么区别？
|
||
2. 为什么断点续训不能只保存 `model.state_dict()`？
|
||
3. AdamW 的 optimizer state 里为什么有必要保存动量状态？
|
||
4. `epoch` 和 `step` 分别表示什么？
|
||
5. `SkipBatchSampler` 跳过的是样本还是 batch？
|
||
6. 为什么恢复后的第一个 epoch 要 skip，后续 epoch 不需要 skip？
|
||
7. `scaler.state_dict()` 是为了解决什么恢复问题？
|
||
8. 为什么保存 checkpoint 时建议先写 `.tmp` 再 `os.replace`？
|
||
|
||
<a id="l15-next"></a>
|
||
## 9. 下一课
|
||
|
||
第 16 课进入 Full SFT 训练脚本总览。
|
||
|
||
下一课要解决：
|
||
|
||
```text
|
||
SFTDataset 如何产出 input_ids / labels；
|
||
train_full_sft.py 和 train_pretrain.py 哪些地方相同；
|
||
从 pretrain 权重继续训练时 from_weight 做了什么；
|
||
SFT 的 loss 仍然如何走 causal LM forward；
|
||
教学版 SFT 脚本应该如何复用 train_loop 和 checkpoint。
|
||
```
|