企业官网建设流程全解析

1. 为什么训练到一半崩了，模型权重却救不回来

做过深度学习训练的人大概都经历过这种窒息时刻：模型跑到第 37 个 epoch，loss 曲线正朝着理想方向收敛，突然显存爆了、机器被抢占、或者一个莫名其妙的 CUDA error 直接把进程干掉。等你重新登录，发现除了终端里残留的几行日志，什么都没留下——权重没了，优化器状态没了，学习率调度器的进度也没了，只能从头再来。

这就是Checkpoint（检查点）机制存在的意义。在 PyTorch 的语境里，checkpoint 不只是"把模型参数存下来"这么简单，它是一套完整的训练状态快照方案，涵盖模型权重、优化器动量、学习率调度器步数、随机数种子状态，甚至混合精度训练的 scaler 状态。少存一样，恢复训练时就可能出现"看起来接上了、实际训练行为已经变了"的隐性偏差。

这篇内容面向的是已经能跑通基础训练脚本、但还没系统梳理过 checkpoint 设计的开发者。我会把 PyTorch 里 checkpoint 的几种典型形态拆开讲：最朴素的state_dict保存、完整的训练状态恢复、大模型场景下的激活值重计算（gradient checkpointing，注意它和训练存档是两回事但经常被混为一谈）、以及分布式训练下的分片存档。中间会穿插我自己踩过的坑，比如优化器状态没存导致恢复后 loss 突然跳变、map_location用错导致 GPU 显存被白白占满这类问题。

先把一个容易混淆的概念摆正：PyTorch 里"checkpoint"这个词至少指代两个完全不同的东西。一个是训练存档（training checkpoint），就是把训练状态落盘以便断点续训；另一个是梯度检查点（gradient checkpointing），用计算换显存，前向时不保存中间激活、反向时重新算一遍。两者名字一样，用途天差地别，但社区里经常有人问"我开了 checkpoint 为什么还是要存模型"，就是把这两个搞混了。下面会分别讲清楚。

2. state_dict 到底存了什么，为什么不能只存模型

2.1 从一次"恢复后 loss 暴涨"说起

早些年我写训练脚本，保存逻辑就一行torch.save(model.state_dict(), path)，恢复时model.load_state_dict(torch.load(path))，看起来天衣无缝。直到有一次做对比实验，中途机器重启，我从第 20 个 epoch 的存档恢复，结果 loss 从 0.8 直接跳到 2.3，训了好几个 epoch 才慢慢降回去。当时第一反应是数据加载顺序变了，查了半天 dataloader，最后才发现问题出在优化器状态没存。

Adam 这类自适应优化器内部维护着每个参数的一阶动量（exp_avg）和二阶动量（exp_avg_sq）。你只恢复模型权重、不恢复优化器状态，相当于让一个跑了两万步的模型突然换上一套全新的、零动量的优化器。Adam 在初期对梯度的估计极不稳定，更新步长会明显偏大，表现出来就是 loss 突跳。SGD with momentum 同理，动量缓冲区清零后，前几步的更新方向会偏离原来的轨迹。

所以一个"能真正续训"的 checkpoint，至少得包含这几样：

2.2 一个可直接抄的保存/恢复模板

下面这套结构是我目前项目里稳定在用的，字段划分清晰，恢复时按需取用：

import torch import random import numpy as np def save_checkpoint(path, model, optimizer, scheduler, scaler, epoch, best_metric): torch.save({ "epoch": epoch, "model": model.state_dict(), "optimizer": optimizer.state_dict(), "scheduler": scheduler.state_dict() if scheduler else None, "scaler": scaler.state_dict() if scaler else None, "best_metric": best_metric, "rng": { "torch": torch.get_rng_state(), "numpy": np.random.get_state(), "python": random.getstate(), }, }, path) def load_checkpoint(path, model, optimizer=None, scheduler=None, scaler=None, device="cuda"): ckpt = torch.load(path, map_location=device) model.load_state_dict(ckpt["model"]) if optimizer and "optimizer" in ckpt: optimizer.load_state_dict(ckpt["optimizer"]) if scheduler and ckpt.get("scheduler"): scheduler.load_state_dict(ckpt["scheduler"]) if scaler and ckpt.get("scaler"): scaler.load_state_dict(ckpt["scaler"]) if "rng" in ckpt: torch.set_rng_state(ckpt["rng"]["torch"]) np.random.set_state(ckpt["rng"]["numpy"]) random.setstate(ckpt["rng"]["python"]) return ckpt.get("epoch", 0), ckpt.get("best_metric", None)

这里有个细节值得单独说：torch.get_rng_state()返回的是 CPU 上的随机数生成器状态。如果你用了 CUDA 上的随机操作（比如某些 dropout 实现、或者你自己写的 kernel），严格来说还得存torch.cuda.get_rng_state_all()。不过实践中大多数标准层的随机性都走 CPU generator，存 CPU 状态已经能覆盖绝大部分复现需求。追求极致可复现的话，把 CUDA 状态也一并存上，代价只是存档大几 KB。

2.3 为什么 state_dict 是"引用式"的，存的时候要小心

model.state_dict()返回的是一个有序字典，里面的 tensor 是和模型参数共享内存的引用，不是拷贝。这意味着如果你在保存之后、落盘之前又对模型做了一步更新，存下来的可能是更新后的值——虽然torch.save是同步写盘，一般不会出问题，但在多线程或者异步保存的场景下这就是个隐患。

更常见的坑是：有人为了省显存，保存时用model.state_dict()直接传给torch.save，然后在另一个线程里继续训练。如果保存过程恰好和参数更新重叠，可能存出一个"半新半旧"的不一致快照。稳妥做法是在保存前做一次浅拷贝：

state = {k: v.detach().clone() for k, v in model.state_dict().items()} torch.save(state, path)

代价是保存瞬间多占一份模型大小的内存。模型不大时无所谓，几十 GB 的大模型就得考虑别的策略了，后面第 5 节会讲分片保存。

3. 梯度检查点：用算力换显存的那笔账怎么算

3.1 它和训练存档完全是两码事

前面提过，gradient checkpointing（梯度检查点）经常被和训练存档混淆。它的本质是：前向传播时只保存少数几个"检查点"层的激活值，其余中间激活全部丢弃；反向传播需要用到某个中间激活时，从最近的检查点重新做一次前向把它算出来。这样显存占用大幅下降，代价是反向时多算一遍前向，训练变慢。

PyTorch 提供了torch.utils.checkpoint.checkpoint和checkpoint_sequential两个入口。用起来很简单：

from torch.utils.checkpoint import checkpoint class Block(nn.Module): def __init__(self, dim): super().__init__() self.ffn = nn.Sequential( nn.Linear(dim, dim * 4), nn.GELU(), nn.Linear(dim * 4, dim), ) def forward(self, x): return checkpoint(self.ffn, x, use_reentrant=False)

use_reentrant=False这个参数值得强调。老版本的 checkpoint 默认走 reentrant 实现，它对输入有要求（至少一个输入需要requires_grad=True），而且在某些嵌套场景下行为和预期不符。PyTorch 1.11 之后提供了非 reentrant 版本，兼容性更好、支持更多边界情况，新代码建议直接显式指定use_reentrant=False。

3.2 显存省多少，速度慢多少

这笔账得算清楚，不然容易"省了显存但训练慢到无法接受"。假设一个 Transformer block 的激活值占A字节，你有L层。不开启 checkpoint 时，反向传播需要保留所有层的激活，显存开销约L * A。开启后，如果每隔k层设一个检查点，只需保存L/k份激活，反向时每段重算一次前向，额外计算量约等于多做一次前向。

经验数据：对典型的 Transformer 结构，开启梯度检查点通常能把激活显存降到原来的 1/3 到 1/5，训练速度下降 20% 到 35%。这个 trade-off 在显存吃紧、batch size 上不去的时候非常划算——毕竟 batch size 翻倍带来的收敛收益，往往比多花 30% 时间更值。

但有个反直觉的点：不是所有层都值得设检查点。Embedding 层、LayerNorm 这类激活值小、计算便宜的操作，设检查点省不了多少显存，反而白白增加重算开销。我一般只对 attention 和 FFN 这类激活值大的模块开 checkpoint，用小模块包一层判断：

def forward(self, x): if self.gradient_checkpointing and self.training: return checkpoint(self._forward_impl, x, use_reentrant=False) return self._forward_impl(x)

注意self.training这个判断——推理时不需要梯度，开 checkpoint 纯属浪费。很多框架的gradient_checkpointing_enable()内部就做了这个判断，自己手写的话别忘了。

3.3 一个容易忽略的副作用：dropout 和 BN 的行为

重算前向时，如果模块里有 dropout，两次前向的随机 mask 不一样，会导致反向传播用的激活和前向时实际产生的激活不一致。PyTorch 的 checkpoint 实现通过保存和恢复 RNG 状态来处理这个问题，但前提是你用的是官方checkpoint函数。如果你自己手写"丢弃激活、反向重算"的逻辑而忘了处理 RNG，训练会悄悄出问题——loss 能降，但收敛质量和不开 checkpoint 时有细微差异，很难排查。

BatchNorm 更麻烦：重算前向会再次更新 running_mean/running_var，导致统计量被更新两次。所以开了梯度检查点的模型，BN 层的行为需要特别留意，实践中要么把 BN 换成 GroupNorm/LayerNorm，要么确保 checkpoint 的粒度不跨越 BN 层。这也是为什么现在大模型普遍用 LayerNorm 而不是 BN 的原因之一——对 checkpoint 友好。

4. 断点续训的完整链路：从存档到"无缝接上"

4.1 恢复顺序有讲究

很多人恢复时随手写model.load_state_dict然后optimizer.load_state_dict，顺序无所谓——对，加载顺序确实不影响结果，但设备放置和参数组映射会影响。优化器的state_dict里保存的是参数索引到状态的映射，加载时必须保证优化器持有的参数顺序和保存时一致。如果你在恢复前重建了优化器、参数顺序变了（比如给某些层加了 weight decay 分组），load_state_dict会报 key 不匹配或者静默错位。

稳妥做法是：先构建模型、再构建优化器、最后加载状态，且优化器的参数分组逻辑在保存和恢复两侧保持完全一致。我习惯把参数分组写成一个独立函数，保存和恢复都调它，避免两边逻辑漂移。

4.2 map_location 用错，显存白占

torch.load默认会把 tensor 加载到它保存时所在的设备。如果存档是在 8 卡机器上存的，你在单卡机器上加载，默认行为可能尝试往 cuda:0 到 cuda:7 上放，直接报设备不存在或者把显存占满。

正确姿势是显式指定map_location：

# 加载到指定 GPU ckpt = torch.load(path, map_location="cuda:0") # 或者先加载到 CPU，再按需搬到 GPU（更省显存峰值） ckpt = torch.load(path, map_location="cpu") model.load_state_dict(ckpt["model"]) model.to("cuda:0")

第二种方式多一步搬运，但加载瞬间的显存峰值更低。对于大模型，先 CPU 后 GPU 几乎是必须的，否则加载过程本身就可能 OOM。

还有一个隐蔽的坑：map_location只影响 tensor 的落点，不影响优化器状态里那些非 tensor 的元数据。如果你存了best_metric这类 Python 标量，它们本来就在 CPU 上，不受影响。但如果你存了自定义对象（比如某个 scheduler 的内部状态），反序列化时可能触发pickle的模块导入，跨环境恢复时容易因为类路径变化而失败。所以存档里尽量只放基础类型和 tensor，自定义对象要么转成 dict，要么在恢复侧保证类定义一致。

4.3 续训时那几个"看起来对但实际错"的细节

恢复训练不只是加载状态，还有几个配套动作容易漏：

• dataloader 的 epoch 对齐：如果你用for epoch in range(start_epoch, total_epochs)续训，但 dataloader 的 shuffle 种子没跟着恢复，数据顺序会和"没中断"时不同。严格复现的话，dataloader 的 generator 种子也要纳入 checkpoint。
• 学习率 warmup 的重入：有些 scheduler 在恢复后会把 step 计数重置，导致 warmup 重新走一遍。检查scheduler.state_dict()里是否包含last_epoch，加载后确认scheduler.last_epoch等于保存时的值。
• 梯度累积的中间状态：如果你用梯度累积，且恰好在累积到一半时存档，恢复后那半个 batch 的梯度是丢的。要么在累积边界存档，要么把累积计数也存进去。
• AMP scaler 的 growth tracker：GradScaler内部维护着 scale 值和增长计数器，不恢复的话，恢复初期 scale 会从初始值重新爬升，可能造成几步梯度下溢被跳过。

这些细节单看都不起眼，但叠在一起就是"恢复后训练曲线和原来对不上"的元凶。我的习惯是写一个resume_sanity_check函数，恢复后打印 epoch、global step、当前学习率、scaler 的 scale 值，和存档时的日志对一遍，确认无误再开跑。

5. 大模型与分布式场景：存档策略要换思路

5.1 单文件存档在几十 GB 模型上会崩

模型参数上到几十 GB 后，torch.save一个完整 state_dict 会遇到两个问题：一是保存瞬间需要把整个 state_dict 序列化到内存，峰值内存可能翻倍；二是单文件太大，写盘慢、传输难、加载也慢。

PyTorch 2.x 之后推荐用torch.distributed.checkpoint（DCP）做分片保存。它支持把 state_dict 按参数切分，每个 rank 只写自己负责的那部分，加载时也能并行读回。基本用法：

import torch.distributed.checkpoint as dcp dcp.save( state_dict={"model": model.state_dict(), "optimizer": optimizer.state_dict()}, checkpoint_id="./ckpt_dir", )

DCP 的好处是存档目录里是一堆分片文件加一个元数据索引，加载时按需读取，不需要一次性把整个模型塞进内存。代价是存档不再是单文件，迁移时得整个目录一起搬。

5.2 FSDP / DDP 下的状态字典差异

DDP 相对简单，每个 rank 的模型参数是完整的（只是梯度做了 all-reduce），所以存档时通常只在 rank 0 上保存完整 state_dict 即可。但要注意：rank 0 保存时其他 rank 要同步等待，否则可能出现 rank 0 还在写盘、其他 rank 已经进入下一轮训练，导致保存的权重和训练进度错位。标准做法是用torch.distributed.barrier()在保存前后做同步。

FSDP 就复杂了，参数被分片到各个 rank 上，单个 rank 只有一部分。PyTorch 提供了FSDP.state_dict_type上下文管理器，可以切换成全量 state_dict（聚合到 rank 0）或分片 state_dict（各存各的）。全量模式简单但内存峰值高，分片模式省内存但恢复时要保证分片布局一致。大模型场景下基本都走分片 + DCP 的组合。

5.3 存档频率与保留策略

最后一个实操层面的问题：多久存一次、存几份。存太频繁拖慢训练（尤其大模型写盘要几十秒），存太少丢的进度多。我的经验是按时间而不是按 epoch 定频率：小模型每 10 到 15 分钟存一次，大模型每 30 到 60 分钟一次，同时保留最近 3 份滚动覆盖，外加一份"最佳指标"存档单独保留。

滚动覆盖的实现很简单，用固定文件名加后缀轮转：

def rotate_save(path_prefix, keep=3, **kwargs): for i in range(keep - 1, 0, -1): src = f"{path_prefix}_{i}.pt" dst = f"{path_prefix}_{i+1}.pt" if os.path.exists(src): os.replace(src, dst) save_checkpoint(f"{path_prefix}_1.pt", **kwargs)

这样即使最新一份存档写坏了（写盘中途断电是真实会发生的），还能回退到上一份。别问我怎么知道要加这个轮转的——问就是曾经在一个跑了三天的任务上，唯一一份存档在写盘时机器挂了，文件损坏，三天白跑。

6. 几个我反复踩过的坑和对应的检查清单

把前面散落的经验收拢成一份可对照的清单，每次新项目搭训练框架时过一遍，能省掉大量事后排查：

• 只存了模型没存优化器：恢复后 loss 跳变，收敛轨迹改变。检查optimizer.state_dict()是否在存档里。
• map_location没指定：跨设备恢复时报错或显存爆。加载时永远显式指定落点。
• scheduler 的last_epoch没恢复：学习率错位，warmup 重走。加载后打印scheduler.last_epoch核对。
• 梯度检查点开在了推理阶段：白白增加计算。用self.training做门控。
• checkpoint 跨越了 BN 层：running stats 被更新两次。调整 checkpoint 粒度或换归一化方式。
• 分布式保存没做 barrier：存档和训练进度错位。保存前后加同步。
• 存档只留一份：写盘损坏就全丢。滚动保留多份。
• 自定义对象塞进存档：跨环境恢复时反序列化失败。只存基础类型和 tensor。

这些坑的共同点是：它们都不会在保存时报错，只会在恢复后以"训练行为异常"的形式暴露，而那时候你往往已经跑了好几个小时。所以 checkpoint 逻辑的正确性，最好在项目初期就用一个小模型、短训练做一次"存档-杀进程-恢复-对比"的验证，确认恢复后的 loss 曲线和没中断时能对上。这个验证花不了半小时，但能避免后面几天的返工。

我个人现在的习惯是，任何新训练脚本跑通第一个 epoch 后，立刻手动 kill 一次、从存档恢复、再跑两个 epoch，对比 loss 数值。数值对得上（允许浮点误差级别的微小差异），这套 checkpoint 逻辑才算过关。这个动作已经帮我拦下过至少三次"优化器状态漏存"和"RNG 没恢复"的问题，比事后对着异常曲线猜原因高效得多。