【Bug已解决】Degraded performance when resuming from checkpoint 解决方案

一、现象长什么样

训练跑了一段时间，存了 checkpoint 然后从 checkpoint 恢复继续训练，发现恢复后吞吐明显下降、每 step 变慢：

恢复前： 120 samples/sec 恢复后： 78 samples/sec （明显下降，无报错）

甚至恢复后长时间爬不回原来的速度。常见的"嫌疑"很多，但核心是：恢复动作本身引入了某种持续的性能拖累。

最小判据：

触发：从 checkpoint 恢复训练后，持续吞吐下降 现象：每 step 变慢，无报错 根因：恢复破坏了某个性能前提（DataLoader 状态 / 编译缓存 / 持久 worker） 影响：恢复后训练变慢，整体时长被拉长

最迷惑的是：恢复"功能上正确"（loss 接续得上），只是"变慢"——典型的 silent 性能回归，容易被当成"数据 / 网络波动"。

二、背景

恢复训练会重建一批运行时状态，任何一步没恢复原样，都可能留下持续的性能拖累。高频原因：

1. DataLoader 状态未恢复（采样器 epoch / RNG）：恢复后若RandomSampler/DistributedSampler的 epoch 和 RNG 没还原，数据顺序被打乱、或从头重读，导致磁盘缓存未命中（之前预热好的 page cache 失效），每个 step 都要重新从磁盘读数据 -> CPU 侧变慢 -> GPU 等数据 -> 吞吐降。这是最常见、也最隐蔽的。
2. num_workers/ 持久 worker 被重置：恢复时若 DataLoader 被重新构造且persistent_workers=False，worker 进程被销毁重建，恢复后的前若干 step 都在"重新 import / 预热 worker"，拖慢整体。
3. torch.compile缓存失效：恢复后若模型结构/设备有细微变化（比如参数被重新加载到新 tensor 对象），dynamo 的编译缓存失效，重新编译（recompile storm），恢复后前 N 步极慢。
4. CUDA Graph 失效：若用了 CUDA Graph，恢复后参数对象变了，graph 需重建，重建期间慢。
5. 优化器状态放大：若 optimizer state 恢复得不对（如放大了某些 buffer），每步 optimizer step 变重。

根因是"恢复动作破坏了某个性能前提，且该破坏是持续性的"。

三、根因

抽象成代码（示意）：

# 恢复时只存了模型/优化器，丢了 DataLoader 的 sampler 状态 def resume(): load_model_optim() # 模型/优化器恢复 # BUG：没恢复 sampler.set_epoch / RNG -> 数据重读，page cache 失效

根因链条：

1. 恢复只还原了模型 / 优化器权重；
2. DataLoader 的 sampler epoch / RNG 没还原 -> 数据顺序 / 起点错；
3. 磁盘 page cache 未命中（或重新 shuffle），每个 step 从磁盘读；
4. CPU 预处理变慢 -> GPU 等数据 -> 吞吐持续下降；
5. 功能正确（loss 接续）、性能下降，silent 回归。

一句话：恢复时丢了 DataLoader 的 sampler 状态 / 持久 worker / 编译缓存，破坏了性能前提。

四、最小可运行复现

用纯 Python 模拟"未恢复 sampler 状态导致缓存未命中、吞吐降"：

# repro_resume_perf.py def step_through(data_cache_ready): # 缓存命中时快，未命中时慢 return 1.0 if data_cache_ready else 3.0 # 单位时间成本 def simulate_resume(restore_sampler_state): # 训练预热后 page cache 就绪 cache_ready_before = True if not restore_sampler_state: # 没恢复 sampler -> 数据起点变 -> 缓存失效 cache_ready_before = False cost = step_through(cache_ready_before) return cost def main(): cost_bad = simulate_resume(restore_sampler_state=False) cost_good = simulate_resume(restore_sampler_state=True) print("未恢复 sampler：每 step 成本", cost_bad) print("恢复 sampler：每 step 成本", cost_good) assert cost_bad > cost_good, "复现：未恢复 sampler 状态导致变慢" if __name__ == "__main__": main()

运行输出：

未恢复 sampler：每 step 成本 3.0 恢复 sampler：每 step 成本 1.0

未恢复 sampler 状态让每 step 成本翻 3 倍，正是"恢复后变慢"的抽象。

五、解决方案（第一层：最小直接修复）

最小且必须的一步：恢复时一并恢复 DataLoader 的 sampler 状态（epoch + RNG），并保持persistent_workers=True避免 worker 重建：

# fix_layer1.py def save_checkpoint(model, optimizer, dataloader, path): sampler = dataloader.sampler ckpt = { "model": model.state_dict(), "optimizer": optimizer.state_dict(), "sampler_epoch": getattr(sampler, "epoch", 0), "sampler_rng": sampler.state_dict() if hasattr(sampler, "state_dict") else None, } torch.save(ckpt, path) def load_checkpoint(model, optimizer, dataloader, path): ckpt = torch.load(path) model.load_state_dict(ckpt["model"]) optimizer.load_state_dict(ckpt["optimizer"]) sampler = dataloader.sampler if hasattr(sampler, "epoch"): sampler.set_epoch(ckpt["sampler_epoch"]) if ckpt["sampler_rng"] and hasattr(sampler, "load_state_dict"): sampler.load_state_dict(ckpt["sampler_rng"])

要点：

• 把sampler_epoch/sampler_rng一并存读，数据起点一致 -> page cache 命中；
• persistent_workers=True让 worker 跨 epoch 不重建，避免预热开销；
• 恢复后吞吐回到恢复前水平。

六、解决方案（第二层：结构性改进）

把"可恢复的全部运行时状态"做成显式清单，save/load 对称处理，避免"只存模型/优化器"的遗漏：

# fix_layer2.py from dataclasses import dataclass, field from typing import Dict, Any @dataclass class ResumableState: model: Dict optimizer: Dict sampler_epoch: int = 0 sampler_rng: Any = None dataloader_rng: Any = None compile_cache_key: str = "" class ResumeManager: def capture(self, model, optimizer, dataloader): sampler = dataloader.sampler return ResumableState( model=model.state_dict(), optimizer=optimizer.state_dict(), sampler_epoch=getattr(sampler, "epoch", 0), sampler_rng=sampler.state_dict() if hasattr(sampler, "state_dict") else None, dataloader_rng=torch.get_rng_state(), ) def restore(self, state, model, optimizer, dataloader): model.load_state_dict(state.model) optimizer.load_state_dict(state.optimizer) sampler = dataloader.sampler if hasattr(sampler, "epoch"): sampler.set_epoch(state.sampler_epoch) if state.sampler_rng and hasattr(sampler, "load_state_dict"): sampler.load_state_dict(state.sampler_rng) if state.dataloader_rng is not None: torch.set_rng_state(state.dataloader_rng)

要点：

• ResumableState显式列出所有可恢复状态（含 sampler / dataloader RNG）；
• ResumeManager对称 capture/restore，不遗漏任何性能相关状态；
• 编译缓存 key 也可纳入，恢复后复用编译结果避免 recompile。

七、解决方案（第三层：断言 / CI 守护）

写 pytest 验证"恢复后 sampler 状态一致、吞吐前提不被破坏"：

# test_resume_perf.py import pytest def make_state(epoch, rng): return {"epoch": epoch, "rng": rng} def restore_into(state, sampler): sampler["epoch"] = state["epoch"] sampler["rng"] = state["rng"] def test_sampler_epoch_restored(): s = make_state(epoch=5, rng=123) sampler = {} restore_into(s, sampler) assert sampler["epoch"] == 5, "sampler epoch 必须恢复，否则数据起点错" def test_rng_restored(): s = make_state(epoch=5, rng=999) sampler = {} restore_into(s, sampler) assert sampler["rng"] == 999, "RNG 必须恢复，否则 page cache 命中率降" def test_missing_sampler_state_is_bug(): # 没存 sampler 状态 -> 恢复后 epoch 默认 0 -> 起点错位 s = {} # 漏存 sampler = {"epoch": 0} if "epoch" in s: sampler["epoch"] = s["epoch"] assert sampler["epoch"] == 0, "漏存导致 epoch 复位 -> 性能回归"

CI 一旦有人把 sampler 状态从 checkpoint 删掉，相关测试能拦下。

八、排查清单

恢复后变慢时：

1. 确认是否"功能正确但吞吐持续下降"（silent 性能回归）；
2. 检查 checkpoint 是否只存了模型/优化器，丢了 sampler epoch/RNG；
3. 看 DataLoader 是否persistent_workers=True（避免 worker 重建预热）；
4. 检查 torch.compile / CUDA Graph 是否在恢复后 recompile（参数对象变了）；
5. 按第五 / 六节恢复 sampler 状态 + 持久 worker + 复用编译缓存；
6. 对比恢复前后每 step 耗时，定位是数据侧还是计算侧变慢；
7. 把第七节的 pytest 接进 CI，守护"恢复状态完整"。

九、小结

从 checkpoint 恢复后性能下降，根因是恢复动作只还原了模型/优化器，丢了 DataLoader 的 sampler epoch/RNG、持久 worker、编译缓存等性能前提：数据起点错位导致磁盘 page cache 未命中、worker 重建预热、编译 recompile，从而持续变慢。功能正确、性能 silent 回归。

三层层级：

• 第一层：恢复时一并恢复 sampler epoch/RNG，保持persistent_workers=True；
• 第二层：用ResumableState显式列出全部可恢复状态，save/load 对称；
• 第三层：pytest 验证 sampler/RNG 状态被恢复，锁进 CI。

核心教训：checkpoint 不只是"模型+优化器"。任何影响数据读取顺序、worker 生命周期、编译缓存的运行时状态，都是性能的隐式前提；漏恢复任何一个，都会让"恢复后变慢"成为难查的 silent 回归。