AI服务器GPU快照配置与检查点恢复方法
训练任务跑在GPU上时，所谓“快照”要分两层看：虚拟化平台的整机快照能保住系统和文件，但训练进程的内存状态、优化器状态和随机数状态不会自动保留。
要让任务中断后能继续训练，必须在代码里做检查点（checkpoint），把模型权重、优化器状态、当前epoch和步数、数据加载进度一起落盘。
先确认快照目标：保文件还是保训练状态
如果只是系统盘或数据盘需要备份，用虚拟化平台或云厂商的磁盘快照即可，恢复后文件和系统环境不变。
如果目标是训到一半断电后能接着跑，单靠磁盘快照不够，因为GPU显存里的模型参数、优化器动量、学习率调度器状态都丢了。
判断标准很简单：恢复后如果要从头开始训，说明检查点没做对；
能从上次保存的step继续，才算真正保住了训练进度。
本文只讨论后者，即训练任务层面的GPU快照配置。
检查点和快照前的准备
开始改代码前，先把下面几项确认好：
- 训练框架版本，例如PyTorch 2.x，
torch.save和torch.load的行为在不同版本间有差异，以实际环境为准。 - 保存目录要有足够空间，单个7B模型的全量检查点通常包含模型权重和优化器状态，体积可能是模型权重的3倍左右。
- 保存路径建议放在独立数据盘，例如
/data/checkpoints/，不要放在系统盘或临时目录。 - 多卡训练要确认启动方式，DDP和FSDP的状态保存方式不同，本文以DDP为例。
- 记录当前Git提交号或代码版本，恢复时保证模型结构一致。
保存检查点的核心代码
在训练循环里按固定步数保存，而不是只在epoch结束时保存。
下面是最小可用写法：
import os
import torch
save_dir = "/data/checkpoints"
os.makedirs(save_dir, exist_ok=True)
def save_checkpoint(model, optimizer, scheduler, epoch, step, path):
state = {
"model": model.state_dict(),
"optimizer": optimizer.state_dict(),
"scheduler": scheduler.state_dict(),
"epoch": epoch,
"step": step,
}
torch.save(state, path)
# 在训练循环中每500步保存一次
if global_step % 500 == 0:
ckpt_path = os.path.join(save_dir, f"ckpt_step_{global_step}.pt")
save_checkpoint(model, optimizer, scheduler, epoch, global_step, ckpt_path)多卡DDP场景下，只在主进程保存，避免多个进程同时写同一个文件：
if torch.distributed.get_rank() == 0:
save_checkpoint(model.module, optimizer, scheduler, epoch, global_step, ckpt_path)保存后建议保留最近3到5个检查点，并删除更早的，避免磁盘被写满。
删除前先确认最新检查点能正常加载。
恢复训练时怎么加载
恢复不是只加载模型权重，优化器状态和步数也要一起恢复，否则学习率调度和动量会错位。
示例：
ckpt = torch.load("/data/checkpoints/ckpt_step_5000.pt", map_location="cpu")
model.load_state_dict(ckpt["model"])
optimizer.load_state_dict(ckpt["optimizer"])
scheduler.load_state_dict(ckpt["scheduler"])
start_epoch = ckpt["epoch"]
start_step = ckpt["step"]加载到GPU时，先加载到CPU再移动到GPU，可以避免多卡环境下显存分配冲突。
如果使用DDP，确保每个进程加载相同的检查点后再开始训练。
快照配置里容易踩的坑
保存目录被快照排除：有些平台默认只快照系统盘，数据盘需要手动勾选。
保存检查点前确认目录在快照范围内，或者单独对数据盘做快照。
只保存模型权重：只存model.state_dict()能恢复推理，但恢复训练会丢失优化器状态，导致loss曲线跳变。
训练场景必须保存完整状态。
保存频率过高：每步都保存会拖慢训练，检查点文件也会迅速占满磁盘。
按时间或步数折中，例如每30分钟或每500步保存一次。
快照时训练还在写文件：如果一边训练一边对数据盘做快照，可能拍到不完整的检查点文件。
建议先暂停训练或等检查点写完再触发快照。
验证快照和检查点是否有效
保存完成后，不要等故障发生才验证。
可以做一次恢复演练：
- 复制一个检查点到临时目录，用
torch.load加载，确认没有报错。 - 启动一个短训练任务，从该检查点恢复，观察起始step是否等于保存时的step。
- 对比恢复后前几个batch的loss，和保存前最后几个batch的loss应在同一量级。
- 如果使用虚拟化平台快照，恢复一台测试机，挂载快照后检查
/data/checkpoints/目录下文件是否完整。
验证通过后，再清理旧检查点或调整保存策略。
常见疑问
GPU快照和检查点是一回事吗？
不是。
GPU快照通常指虚拟化层面对磁盘或整机的备份，检查点是训练框架把训练状态写成文件。
前者保环境，后者保训练进度，生产环境建议两者都做。
恢复后loss突然变大怎么办？
先确认优化器状态和调度器状态是否加载成功，再检查学习率是否从保存时的值继续。
如果只加载了模型权重，loss跳变是正常现象，补齐完整状态后重新恢复即可。
多卡训练保存的检查点能在单卡上恢复吗？
可以，但需要把模型结构按单卡方式实例化后再加载权重，优化器状态可能因参数分组不同需要调整。
建议恢复环境和保存环境尽量一致。
快照配置本身不复杂，关键是别把“磁盘快照”当成“训练快照”。
把检查点目录纳入快照范围，保存完整训练状态，并定期做恢复验证，才能真正做到中断后接着跑。
