torch.save()在大模型上特别慢是因为默认使用pickle协议递归序列化复杂对象图，CPU负载高且不可并行；PyTorch 1.6+可通过_use_new_zipfile_serialization=True启用ZIP分块直写提升2–4倍速度，或改用safetensors格式实现零开销序列化。
直接用 torch.save() 或 pickle.dump() 保存大模型，尤其是带大量参数的 LLM 或扩散模型，会卡住几秒到几分钟——这不是磁盘慢，是序列化本身在拖后腿。
为什么 torch.save() 在大模型上特别慢？
PyTorch 默认用 Python 的 pickle 协议序列化模型状态，而 pickle 对嵌套字典、张量引用、复杂对象图做深度遍历和递归序列化，CPU 负载高、不可并行、且无法跳过冗余信息（比如重复的 tensor storage pointer）。
- 模型参数超 1B 时，
torch.save()可能比实际训练还慢 - 若模型含自定义类、闭包或非标准属性，pickle 还可能失败（报
AttributeError: Can't pickle ...） - 默认协议版本（protocol=2）兼容性好但效率低；新版 protocol=5 支持 buffer protocol，但 PyTorch 并未默认启用
用 torch.save(..., _use_new_zipfile_serialization=True)
这是 PyTorch 1.6+ 内置的提速开关，它把模型拆成多个二进制块 + ZIP 容器，绕过 pickle 的对象图遍历，直接 dump tensor data。
- 对 3B 参数模型，实测比默认方式快 2–4 倍，内存峰值降低约 30%
- 必须显式传参，因为旧版兼容性考虑，它不是默认行为
- 保存后文件仍是
.pt后缀，但内部结构变成 ZIP，torch.load()会自动识别，无需改加载逻辑
示例：
torch.save(model.state_dict(), "model.pt", _use_new_zipfile_serialization=True)
立即学习“Python免费学习笔记（深入）”；
换用 safetensors 格式替代 pickle
safetensors 是 Hugging Face 主推的零开销序列化格式：不走 pickle、无代码执行风险、支持 mmap 直接加载、写入速度接近 memcpy。
- 安装：
pip install safetensors - 保存几乎无额外 CPU 开销，10B 模型通常
- 文件大小比
.pt小 5–10%，因为省去了 pickle header 和 metadata 冗余 - 注意：只支持
state_dict（即 dict of tensors），不能直接存整个 model 实例（含方法、graph）
示例：
from safetensors.torch import save_file<br>save_file(model.state_dict(), "model.safetensors")
避免保存不需要的部分
很多模型保存时顺手 dump 了 optimizer、scheduler、scaler、甚至 training logs —— 这些占空间、拖时间，且推理时完全用不到。
- 仅保存
model.state_dict()，而非model整体 - 去掉
optimizer.state_dict()等训练专属字段（除非要做断点续训） - 如果用了混合精度训练，
scaler.state_dict()也别混进去 - 检查是否有大型临时 buffer（如
model._cache）被意外包含在 state_dict 中
真正卡住的往往不是磁盘 IO，而是序列化引擎本身。safetensors 和新 ZIP 序列化能立竿见影，但前提是别把训练期的杂项一起打包进去——模型越“干净”，保存越快。
