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

【Bug已解决】Degraded performance when resuming from checkpoint 解决方案一、现象长什么样训练跑了一段时间存了 checkpoint 然后从 checkpoint 恢复继续训练发现恢复后吞吐明显下降、每 step 变慢恢复前 120 samples/sec 恢复后 78 samples/sec 明显下降无报错甚至恢复后长时间爬不回原来的速度。常见的嫌疑很多但核心是恢复动作本身引入了某种持续的性能拖累。最小判据触发从 checkpoint 恢复训练后持续吞吐下降 现象每 step 变慢无报错 根因恢复破坏了某个性能前提DataLoader 状态 / 编译缓存 / 持久 worker 影响恢复后训练变慢整体时长被拉长最迷惑的是恢复功能上正确loss 接续得上只是变慢——典型的 silent 性能回归容易被当成数据 / 网络波动。二、背景恢复训练会重建一批运行时状态任何一步没恢复原样都可能留下持续的性能拖累。高频原因DataLoader 状态未恢复采样器 epoch / RNG恢复后若RandomSampler/DistributedSampler的 epoch 和 RNG 没还原数据顺序被打乱、或从头重读导致磁盘缓存未命中之前预热好的 page cache 失效每个 step 都要重新从磁盘读数据 - CPU 侧变慢 - GPU 等数据 - 吞吐降。这是最常见、也最隐蔽的。num_workers/ 持久 worker 被重置恢复时若 DataLoader 被重新构造且persistent_workersFalseworker 进程被销毁重建恢复后的前若干 step 都在重新 import / 预热 worker拖慢整体。torch.compile缓存失效恢复后若模型结构/设备有细微变化比如参数被重新加载到新 tensor 对象dynamo 的编译缓存失效重新编译recompile storm恢复后前 N 步极慢。CUDA Graph 失效若用了 CUDA Graph恢复后参数对象变了graph 需重建重建期间慢。优化器状态放大若 optimizer state 恢复得不对如放大了某些 buffer每步 optimizer step 变重。根因是恢复动作破坏了某个性能前提且该破坏是持续性的。三、根因抽象成代码示意# 恢复时只存了模型/优化器丢了 DataLoader 的 sampler 状态 def resume(): load_model_optim() # 模型/优化器恢复 # BUG没恢复 sampler.set_epoch / RNG - 数据重读page cache 失效根因链条恢复只还原了模型 / 优化器权重DataLoader 的 sampler epoch / RNG 没还原 - 数据顺序 / 起点错磁盘 page cache 未命中或重新 shuffle每个 step 从磁盘读CPU 预处理变慢 - GPU 等数据 - 吞吐持续下降功能正确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_stateFalse) cost_good simulate_resume(restore_sampler_stateTrue) 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_workersTrue避免 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_workersTrue让 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( modelmodel.state_dict(), optimizeroptimizer.state_dict(), sampler_epochgetattr(sampler, epoch, 0), sampler_rngsampler.state_dict() if hasattr(sampler, state_dict) else None, dataloader_rngtorch.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 RNGResumeManager对称 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(epoch5, rng123) sampler {} restore_into(s, sampler) assert sampler[epoch] 5, sampler epoch 必须恢复否则数据起点错 def test_rng_restored(): s make_state(epoch5, rng999) 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 删掉相关测试能拦下。八、排查清单恢复后变慢时确认是否功能正确但吞吐持续下降silent 性能回归检查 checkpoint 是否只存了模型/优化器丢了 sampler epoch/RNG看 DataLoader 是否persistent_workersTrue避免 worker 重建预热检查 torch.compile / CUDA Graph 是否在恢复后 recompile参数对象变了按第五 / 六节恢复 sampler 状态 持久 worker 复用编译缓存对比恢复前后每 step 耗时定位是数据侧还是计算侧变慢把第七节的 pytest 接进 CI守护恢复状态完整。九、小结从 checkpoint 恢复后性能下降根因是恢复动作只还原了模型/优化器丢了 DataLoader 的 sampler epoch/RNG、持久 worker、编译缓存等性能前提数据起点错位导致磁盘 page cache 未命中、worker 重建预热、编译 recompile从而持续变慢。功能正确、性能 silent 回归。三层层级第一层恢复时一并恢复 sampler epoch/RNG保持persistent_workersTrue第二层用ResumableState显式列出全部可恢复状态save/load 对称第三层pytest 验证 sampler/RNG 状态被恢复锁进 CI。核心教训checkpoint 不只是模型优化器。任何影响数据读取顺序、worker 生命周期、编译缓存的运行时状态都是性能的隐式前提漏恢复任何一个都会让恢复后变慢成为难查的 silent 回归。