让数据跟上,也让训练能恢复

判断输入等待,保存真正能续跑的状态,核对样本消费与恢复结果。

第 04 单元 · 约 80 分钟(含练习) · 产出:数据供给实验与恢复对账

从哪里接上:第 03 课的训练步

设备在等输入,进程却没有报错

第 03 课的一次更新计算正确,训练仍然很慢。看时间线发现 GPU 每做完一批就空一会儿:下一批数据还没准备好。模型代码不是唯一的生产者;读取、解压、解析、随机增强、组 batch 和传输都可能在它前面排队。先分清“没有数据可算”和“数据到了却算得慢”,才知道应该改哪一段。

输入流水线可以看作生产者、有限队列和消费者。生产者把原始记录变成模型需要的张量,队列保存已准备的 batch,训练步是消费者。预取是提前准备后面的 batch,让准备与计算在条件允许时重叠。它隐藏等待,并没有让存储或 CPU 凭空变快。队列太浅可能接不住波动,太深则占内存、增加在途数据,也让恢复时更难说明哪些样本真正完成了更新。

先算稳态速度,再决定预取多少

四个 batch,每批准备 6ms、训练 10ms,忽略传输并假设准备与训练可完全重叠。串行要 4×(6+10)=64ms;流水线先等 6ms,然后连续训练四批,共 46ms。节约的是三个重叠区间,不是把第一次输入等待也删掉。无限长的理想稳态每批受较慢阶段限制,这里是 10ms。

四批数据准备与计算重叠,46毫秒完成准备阶段依次为0到6、6到12、12到18、18到24毫秒;训练阶段为6到16、16到26、26到36、36到46毫秒。图按每毫秒12像素绘制,有足够输入缓冲。理想教学流水线:同一批准备完成后,才可开始计算准备1234训练12340616263646ms重叠要求独立资源和足够缓冲;这不是本机实测结果。
图 04.1 · 横向位置是时间,纵向区分阶段。不要把两行耗时再相加。

如果准备改为 14ms,四批串行需 96ms,理想流水线也要 14×4+10=66ms。GPU 会反复等生产者;增加队列只能吸收短时波动,无法长期超过 14ms/批的供给能力。相反,如果计算很慢而队列一直满,增加数据工作进程可能只增加争抢和内存。

定位时同时记录取 batch 等待、准备时间、训练时间、队列深度与样本吞吐。先缓存一份已预处理输入做对照:若等待消失,继续拆读取与变换;若没有,检查传输或同步。小文件很多时,请求开销可能比数据量更关键;把记录打成 shard(数据分片文件)能改变读取模式,但也会影响随机访问、并行粒度和故障重读范围。必须固定同一数据版本和采样规则再比较。

恢复的对象是一段训练过程,不只是权重

checkpoint 是持久保存、可用于续跑的训练快照。第 03 课解释过参数和优化器状态;现在还要保存更新步数、学习率计划的位置,以及样本消费的位置。RNG 是随机数生成器,它的状态决定之后的随机顺序或随机操作。只保存最初 seed 不一定能回到中途的随机位置。

先选择一致性边界:本实验只在完成一次优化器更新后保存。此时参数、动量与 step 必须属于同一步,样本游标指向下一批尚未应用更新的数据。若选择在梯度累积中途保存,还得保存部分梯度、累积计数和归一化分母,否则恢复后并不是同一更新。正式系统还应按实现记录学习率调度器、混合精度缩放器和各进程随机状态。

完整恢复要回答的四个问题
问题快照中的信息漏掉会怎样
正在学什么模型、参数、数据版本与配置恢复到另一任务或另一模型
已经怎样更新优化器状态、步数、学习率状态下一次更新方向或尺度改变
接下来读谁样本顺序、已提交消费位置、采样状态重复训练或跳过样本
之后如何随机相关 RNG 状态后续打乱或随机变换不一致

读过,不等于已经训练过

假设 batch 0–9 已完成更新,batch 10–12 已被预取进内存。此时故障,安全快照的下一消费位置是 10,不是 13。重启后重新读取 10–12 是重建尚未提交的工作,不算重复更新;若错误地存成 13,就静默漏掉了三个 batch。反过来,用第 10 次更新后的参数配第 8 次更新后的游标,也会把两批样本重复应用。

异步保存也不能边更新同一份状态边序列化:应先获得属于一致时刻的快照,再后台写出。持久层要能区分正在写的临时文件与可恢复的完整版本,校验完整性后再公布入口;本机文件重命名的办法不能不经验证直接推广到对象存储。恢复时找不到可信完整版本,应回退到上一个已验证点,而不是拼接不同批次的分片。

亲手比较:连续 80 步,与 19 步后续跑

下载标准库实验包,按 README 解压并进入 infra-lab。这是实际在 CPU 上运行的两参数线性回归,目标为 y=3x+2,含动量、打乱顺序与断点状态。它没有实现真实 DataLoader 预取、GPU 或分布式存储;上一节的流水线是手算模型。

python3 train.py --steps 80 --checkpoint runs/full.json --output runs/full-report.json
python3 train.py --steps 19 --checkpoint runs/step19.json --output runs/step19-report.json
python3 train.py --resume runs/step19.json --steps 80 --checkpoint runs/resumed.json --output runs/resumed-report.json

--steps 80 是累计目标更新步数,所以第三条命令再做 61 步,不是再做 80 步。打开两个最终 checkpoint,对比 parameters、optimizer、step、samplePosition、sampleOrder 与 rngState;同时核对 dataset。该确定性标准库实验在相同环境应得到完全相同的最终状态。跨设备或框架的正式训练要另外规定容差与确定性条件,不能照搬这个逐字段相等承诺。

预期数值与对应参考文件

默认 seed=7。第 19 步 w≈3.3557282153、b≈2.2900000515,MSE≈0.2521643716,samplePosition=12;第 80 步两条路径均为 w≈3.0000640811、b≈1.9999645195,MSE≈6.713×10⁻⁹,samplePosition=16。可对照 连续运行报告与恢复报告。正文四舍五入的数字只方便阅读,严格比较以完整状态为准。

再只改变恢复方式,保留同一个 19 步快照。下面故意丢弃优化器历史;另两种故障是 reset-cursor 和 reset-rng。不要期待三种错误都立刻使 loss 暴涨;某些差异要在下次重新打乱数据时才暴露。

python3 train.py --resume runs/step19.json --steps 80 --fault drop-optimizer --checkpoint runs/broken.json --output runs/broken-report.json

比较第一个分叉点,不被低 loss 蒙混过去

恢复第一步就不同:先对比参数、动量和下一批样本;到一轮数据末尾才不同:查重新打乱用的 RNG 和顺序;恢复步数超出目标:查把“累计目标”误读成“追加步数”。若只有文件日期不同,先把文件元数据与训练状态区别开。验证目标是同一状态轨迹,不能看到两次 MSE 都很低就宣布恢复等价。

记录恢复成本也要保留范围:发现故障、重新申请设备、载入、补算各花多久。checkpoint 间隔越长,潜在补算越多;保存越频繁,也会争用带宽和存储。无需本课先选出一个万能间隔,先报告保存耗时和可接受的重做工作量。

交付两份对照,而不是一个“成功恢复”截图

  1. 将准备时间改为 14ms,画四批理想流水线,并标出 GPU 的等待区间。
  2. 保存连续、恢复及一种故障的报告;对比最终完整状态,说明哪项差异最早可能改变下一次更新。
  3. 迁移题:在累积 4 个微批次的第 3 次反向之后退出,当前快照只含参数与优化器。写明应回退到哪里,或必须新增保存什么。
参考与验收

四批计算区间是 14–24、28–38、42–52、56–66ms,批间各等 4ms。故障报告必须指出丢动量、改消费位置或改随机序列的具体影响;低 loss 不能抵消状态不一致。中途累积只能恢复最近完整更新点,或连同部分梯度、计数、分母与数据/RNG 状态一起保存一致快照。验收还要附运行命令和原快照,确保他人能复现你的差异。

更新已消费 0–9,队列预取了 10–12,快照下一消费位置应是多少?

带着缺项读 checkpoint 示例

读 PyTorch 通用 checkpoint 示例,找到模型与优化器分别保存的位置。然后把本课的样本游标、RNG 和累积边界补到你自己的状态清单中;框架示例不是所有训练任务的完整恢复合同。再读 可复现性说明,列出跨环境无法直接许诺逐位一致的原因。

下一步:一份状态变成多份以后

你已能在单进程里解释消费与提交。第 05 课加入多个训练进程:它们既要平均梯度,也必须对同一次更新达成一致,否则“某个进程恢复了”不等于整个训练恢复了。

留下本课的产出

本课产出:数据供给实验与恢复对账。写下条件、证据、结果和仍不确定的事。保存不代表自动通过,下面的状态由你按验收条件自检。

记录只保存在当前浏览器。

按验收条件自检