106 lines
2.3 KiB
Markdown
106 lines
2.3 KiB
Markdown
# Day 3 Deep Dive: Pretrain 训练循环
|
||
|
||
## 深挖问题 1：训练循环每一步发生了什么？
|
||
|
||
代码位置：`trainer/train_pretrain.py::train_epoch`
|
||
|
||
流程：
|
||
|
||
```text
|
||
for input_ids, labels in loader:
|
||
move to device
|
||
update lr
|
||
autocast forward
|
||
loss = res.loss + res.aux_loss
|
||
loss /= accumulation_steps
|
||
backward
|
||
if accumulation boundary:
|
||
unscale
|
||
grad clip
|
||
optimizer step
|
||
scaler update
|
||
zero grad
|
||
log
|
||
save
|
||
```
|
||
|
||
你要能解释每一行为什么存在。
|
||
|
||
## 深挖问题 2：学习率怎么变？
|
||
|
||
代码位置：`trainer/trainer_utils.py::get_lr`
|
||
|
||
```python
|
||
return lr * (0.1 + 0.45 * (1 + cos(pi * current_step / total_steps)))
|
||
```
|
||
|
||
含义：
|
||
|
||
- 初始接近 `lr`。
|
||
- 最终接近 `0.1 * lr`。
|
||
- 中间 cosine 平滑下降。
|
||
|
||
工业注意：
|
||
|
||
- 这个实现没有 warmup。
|
||
- 小模型和小数据通常能跑，但大规模训练一般会加 warmup。
|
||
- 如果 early loss 抖动严重，可以考虑降低 lr 或加 warmup。
|
||
|
||
## 深挖问题 3：checkpoint 保存了什么？
|
||
|
||
代码位置：`trainer/trainer_utils.py::lm_checkpoint`
|
||
|
||
保存两类文件：
|
||
|
||
- `out/pretrain_768.pth`：主要用于推理/下游训练。
|
||
- `checkpoints/pretrain_768_resume.pth`：用于断点续训。
|
||
|
||
resume checkpoint 包含：
|
||
|
||
- model
|
||
- optimizer
|
||
- epoch
|
||
- step
|
||
- world_size
|
||
- wandb id
|
||
- scaler 等额外状态
|
||
|
||
为什么 optimizer 也要保存？
|
||
|
||
因为 AdamW 有动量状态。如果只恢复模型权重，优化器状态丢失，训练动态会突然改变。
|
||
|
||
## 深挖问题 4：MoE aux_loss 为什么加在 pretrain loss 上？
|
||
|
||
如果 `use_moe=1`，模型输出 `aux_loss`。训练总 loss：
|
||
|
||
```python
|
||
loss = res.loss + res.aux_loss
|
||
```
|
||
|
||
这表示主目标仍是 next-token prediction，但额外鼓励 expert 负载均衡。
|
||
|
||
## 今日实验
|
||
|
||
1. 先跑 200 step smoke test：
|
||
|
||
```bash
|
||
cd trainer
|
||
python train_pretrain.py --epochs 1 --save_interval 200 --log_interval 20
|
||
```
|
||
|
||
2. 中断训练。
|
||
|
||
3. 恢复：
|
||
|
||
```bash
|
||
python train_pretrain.py --epochs 2 --from_resume 1
|
||
```
|
||
|
||
4. 检查日志是否显示跳过 step。
|
||
|
||
## 复盘问题
|
||
|
||
- 为什么 `loss / accumulation_steps` 后再 backward？
|
||
- 为什么 `scaler.unscale_` 后再 clip grad？
|
||
- 为什么保存权重前切 `model.eval()`，保存后再 `model.train()`？
|