模型权重保存

模型权重保存是训练与部署之间的衔接点：本篇讲透 PyTorch 原生机制——torch.save 的 Pickle 序列化本质、state_dict 与完整模型两种模式、断点续训的保存清单，并给出训练态/推理态格式的全库版图导航。

格式版图：训练态与推理态

Checkpoint 按用途二分，决定了「保存什么」与「选什么格式」：

维度	训练态	推理态
目的	断点续训	部署、分发、推理
内容	权重 + 优化器状态 + 调度器 + RNG + 训练进度	仅权重 + 结构配置
绑定	与训练框架、并行策略强绑定	通用度高，跨框架流通
代表	PyTorch 原生 `.pt`、Megatron/DeepSpeed/FSDP 分片	HF 格式、safetensors、GGUF、ONNX

本篇主体覆盖原生训练态格式；其余版图的细节分发给各专题笔记：

原生单文件（.pt/.pth，本篇）：单卡训练、DDP、小模型推理
分片检查点：Megatron、DeepSpeed ZeRO、FSDP/DCP 的目录结构与分片加载 → 大模型检查点
推理分发格式：HF 目录格式、safetensors、GGUF 量化 → 推理模型格式
跨框架交换：静态图导出部署 → ONNX 模型导出
下载来源：从 Hub 拉取预训练权重 → HuggingFace Hub

torch.save 的本质：Pickle 序列化

torch.save 不是自定义二进制协议，而是基于 Pickle 的 Python 对象序列化：当前默认产出 zip 归档，内含 data.pkl（pickle 化的对象结构，记录张量的 key、dtype、shape 与存储偏移）和 data/0、data/1 等原始字节文件（张量数据本体）。张量数据不经过 pickle 逐元素编码，因此大张量的保存开销可控。

这一本质决定了两条主线：

任意 Python 对象都能存——dict、优化器状态、epoch 计数，甚至整个模型对象；
反序列化即可能执行代码——pickle 协议允许载荷携带可调用的构造逻辑，是安全风险的根源。

.pt 与 .pth 扩展名无语义差别，仅是命名习惯。

两种保存模式

模式	保存内容	评价
完整模型 `torch.save(model, ...)`	序列化整个 `nn.Module`，含结构引用	不推荐
state_dict `torch.save(model.state_dict(), ...)`	有序字典：参数名 → Tensor	行业标准

完整模型模式在 pickle 中记录模型类的模块路径，加载时要求同一份类定义代码处于相同位置，跨项目、跨版本、代码重构后均会断链，仅适合临时快照。

state_dict 模式只存「参数名 → 张量」映射（torch.nn 的 state_dict() 与 load_state_dict() 即为此接口），结构与代码解耦，加载时先由当前代码实例化模型、再灌入权重，兼容性与可移植性最佳。torch.optim 的优化器同样暴露 state_dict()，二者搭配构成断点续训的基础。

断点续训保存什么

仅存权重不足以恢复训练——Adam 动量/方差、学习率调度进度、数据流位置、随机状态任一缺失，续训曲线就会断裂。标准做法是打包一个 checkpoint 字典：

torch.save({
    "model": model.state_dict(),          # 权重参数
    "optimizer": optimizer.state_dict(),  # Adam 一阶/二阶动量
    "scheduler": scheduler.state_dict(),  # LR 调度进度
    "epoch": epoch,
    "step": global_step,
    "rng": torch.get_rng_state(),         # 随机状态：对齐 dropout 与数据打乱
    "cuda_rng": torch.cuda.get_rng_state_all(),
}, f"ckpt_step{global_step}.pt")

恢复时逐项回灌：load_state_dict 灌回模型与优化器，torch.set_rng_state / set_rng_state_all 灌回随机状态，DataLoader 则根据 epoch/step 重定位数据流。

加载注意事项

state = torch.load("weights.pt", map_location="cpu", weights_only=True)
model.load_state_dict(state)  # strict 默认 True，key 不匹配直接报错
model.eval()                  # 推理态：关闭 dropout / BN 统计更新

结构强匹配：加载侧模型类的层名称、结构必须与保存时一致，load_state_dict 默认 strict=True，多余或缺失 key 均抛错；微调改结构后需显式过滤 key。
设备映射：GPU 上保存的权重默认恢复到原设备，跨机或换卡时必须传 map_location（如 "cpu"、"cuda:0"、或字典映射），否则无对应 GPU 时直接失败或显存爆掉。
安全风险：pickle 反序列化可执行任意代码，来路不明的 .pt 严禁加载。PyTorch 2.6 起 torch.load 默认 weights_only=True，仅允许张量与基础容器类型，天然免疫恶意载荷；加载含自定义对象的旧检查点才需显式放宽为 False。公开分发的模型应改用 safetensors（见推理模型格式）。
版本兼容：高版本 PyTorch 保存的文件在低版本可能无法解析（序列化格式与 pickle 协议随版本演进），生产环境保持保存/加载主版本一致；跨版本迁移优先走 safetensors 等中性格式。
适用场景：单卡训练、DDP（每 rank 权重一致，rank 0 保存一份即可，参见分布式训练总览）；参数量超出单卡显存或采用 TP/PP/ZeRO 分片时，原生单文件不再适用，转大模型检查点。

最佳实践

区分用途裁剪内容：仅推理时删掉优化器与 RNG 状态再存档，体积可省一半以上；需续训则完整保留，分片场景严禁删任何 rank 文件。
保存即原子：先写临时文件再 os.replace，训练中途被 kill 不会留下半个文件。
定期转安全格式：训练态 .pt 仅内部流通；对外发布导出 safetensors，双保险规避 pickle 风险。
格式转换走官方工具：分片合并、框架间迁移一律用对应框架的转换脚本，手动改文件名或拼接权重必然错乱。
传输后校验：大文件下载/拷贝后核对 SHA256，损坏的分片在加载晚期才报错，排查代价极高。

知识花园

探索