GPU训练中断后怎么恢复：检查点与数据并行状态处理

GPU训练跑到一半断电或进程被杀，最直接的办法是从最近一次检查点恢复。
但检查点里不只有模型权重，恢复时如果漏掉优化器状态或数据加载位置，训练效果可能变差。
下面按单卡和多卡两种场景说明操作。

先确认检查点里存了什么

一个能完整续跑的检查点通常包含四类内容：模型参数、优化器状态、当前epoch和step、以及随机数状态（可选）。
如果只保存了模型参数，恢复后优化器动量会重置，loss可能出现小幅跳变，短期训练影响不大，长训练建议完整保存。

保存端参考写法：

import torch

torch.save({
    'epoch': epoch,
    'step': global_step,
    'model': model.state_dict(),
    'optimizer': optimizer.state_dict(),
    'scheduler': scheduler.state_dict(),
}, f'ckpt_epoch{epoch}.pt')

恢复前先用ls -lh和torch.load检查文件是否完整，避免加载到写了一半的损坏文件。

单卡恢复：加载权重和优化器状态

单卡场景最简单，关键是加载后把模型切到训练模式，并确认优化器状态确实被读入。

ckpt = torch.load('ckpt_epoch10.pt', map_location='cuda:0')
model.load_state_dict(ckpt['model'])
optimizer.load_state_dict(ckpt['optimizer'])
scheduler.load_state_dict(ckpt['scheduler'])
start_epoch = ckpt['epoch'] + 1
model.train()

加载完成后打印一个参数值或优化器state长度，确认不是空字典。
如果optimizer.state_dict()['state']为空，说明检查点里没有保存优化器状态，需要回退到只加载模型权重的策略。

多卡数据并行：并行度不一致怎么办

用DataParallel或DistributedDataParallel保存的检查点，参数名可能带module.前缀。
换卡数或换加载方式时，直接load_state_dict会报missing keys。
常见处理方式是用map_location指定设备，再对键名做一次清理：

state = ckpt['model']
new_state = {k.replace('module.', ''): v for k, v in state.items()}
model.load_state_dict(new_state, strict=True)

如果恢复后卡数和保存时不同，比如从4卡变2卡，数据并行的梯度平均方式会变，建议先小学习率跑几十步观察loss是否稳定，不要直接按原学习率续跑。
分布式训练还要保证world_size和保存时一致，否则优化器状态里的参数分组可能对不上。

恢复后先做一次短验证

加载完不要立刻跑完整训练。
先固定随机种子跑几十个step，观察loss是否从检查点附近继续下降，同时抽查一个batch的数据索引是否和保存的step衔接。
如果loss突然明显升高，优先检查优化器状态是否加载成功、学习率调度器是否跳到了正确位置。

另外，检查点文件建议保留最近2到3个，并按epoch命名，避免最新文件损坏后无回退点。
存储路径尽量放在独立数据盘，不要和系统盘混用。

常见疑问

检查点加载报CUDA out of memory，通常是map_location没指定，先把权重加载到CPU再转GPU。
多卡恢复后loss异常，先确认保存和恢复的并行策略是否一致，再检查优化器状态是否为空。
只保存了模型权重时，恢复后可以正常推理，但继续训练建议重新积累优化器状态。