第19章 序列化：torch.save / torch.load 与权重格式

"All files are stored uncompressed. All files in the archive are aligned to 64 byte boundaries such that it is possible to mmap the entire file and get an aligned pointer to tensor data."

—— caffe2/serialize/inline_container.h 对 .pt 容器格式的注释

19.1 torch.save 的真实格式

打开一个 .pt 文件用 unzip -l 能看到（按 zip 内的实际顺序，CPU torch 2.11.0 实跑）：

Archive:  model.pt
  model/data.pkl              <- pickle 序列化的 Python 对象图 (无 tensor 数据)
  model/.format_version       <- storage 排序约定的版本
  model/.storage_alignment    <- storage 的对齐字节数
  model/byteorder             <- 字节序标记 ("little" / "big")
  model/data/0                <- 第 1 个 storage 的字节数据
  model/data/1                <- 第 2 个 storage
  ...
  model/version               <- 容器格式版本号 (PyTorchStreamWriter 收尾时写)
  model/.data/serialization_id  <- 本次序列化的唯一 ID (同样由 PyTorchStreamWriter 写)

（写前四条记录的代码在 torch/serialization.py:1260-1277，tensor 数据的循环在 :1280 起。model/ 前缀是 archive name，由 PyTorchStreamWriter 统一加；末尾的 version 与 .data/serialization_id 也是它在收尾时补上的，见 caffe2/serialize/inline_container.cc:828 的 writeEndOfFile。）

torch.save 用的是 zip 容器 + 拆分 storage 的设计：

1. 把对象树（dict / list / nn.Module）用 pickle 序列化，但每个 tensor 写到一个占位符而非数据
2. 每个 storage 的实际字节数据单独存到 zip 里一个 entry（data/0、data/1 ...）
3. 加载时先反序列化 pickle 拿对象骨架 + 占位符，再按需读 zip entry 填字节

MAGIC_NUMBER = 0x1950A86A20F9469CFC6C（torch/serialization.py:65）是老格式（_use_new_zipfile_serialization=False）的文件头标记，新 zip 格式里不再出现在文件开头 —— 校验它的 :1856 只在 legacy 加载路径上跑。

19.1.1 为什么不直接 pickle 整个 dict

朴素 pickle：把所有 tensor 的字节当 bytes 序列化进 pickle 流。问题：

• pickle 流是顺序的，没法 mmap（要全部读到内存才能解析）
• 修改 / 添加单个 tensor 要重写整个文件
• 跨进程共享时序化开销大

zip 容器格式：

• 每个 storage 在 zip 里有独立 entry，支持 mmap
• 所有 entry 都不压缩，且对齐到 64 字节边界 —— inline_container.h 的注释把这两条列为容器的硬保证，正是为了"整个文件 mmap 下来还能拿到对齐指针"
• zip 提供 random access：直接跳到第 N 个 storage
• 兼容现有工具（Python 的 zipfile 能直接查看）

代价是文件比纯 pickle 略大（zip header 开销）。但对 1GB+ 模型 ckpt 几乎不可见。

19.2 mmap 加载：超大 ckpt 几秒完成

torch.load(path, mmap=True) 让 OS 把整个 zip 文件映射到虚拟内存，实际不读到 RAM：

# 没有 mmap: 必须把整个文件读到内存
state = torch.load('llama-70b.pt')   # 280 GB, 单机 RAM 不够

# 有 mmap: OS 按需 page-in
state = torch.load('llama-70b.pt', mmap=True)   # 几乎秒返回
weights = state['model.layer1.weight']           # 用到时才真读

mmap 的几个优势：

• 多个进程加载同一个 ckpt 时共享 OS page cache，省一份内存
• 加载时间几乎为零（懒加载）
• 不连续访问的部分根本不读盘

vLLM / HuggingFace 加载大模型时大量用 mmap。生产代码里这是 ckpt 加载的标准操作。

19.3 weights_only：防止 pickle RCE

pickle 是 Python 的"任意对象序列化"格式，加载时会执行嵌入的代码。恶意 ckpt 可以借 __reduce__ 让加载方调用任意可 import 的函数（比如执行 shell 命令）。

PyTorch v2.6+ 默认 weights_only=True：只允许加载白名单类型（Tensor、Parameter、各种内置类型），其他全部拒绝。这让 ckpt 加载从"等同于 eval(file)"变成"只能按白名单重建已知类型"，接近读 JSON。

如果你的 ckpt 含自定义类（如自定义 Optimizer），要么改用白名单允许：

torch.serialization.add_safe_globals([MyOptimizer])
state = torch.load('ckpt.pt')

要么显式 weights_only=False 退回老行为（不推荐生产用）。

社区从 HuggingFace 的"safetensors 推动"开始一直在向"任何 ckpt 都不该 RCE"靠拢。weights_only 默认 True 是 PyTorch 在同一方向上的安全升级。

19.4 safetensors：HuggingFace 推动的替代格式

safetensors 不是 PyTorch 内置（要 pip install safetensors），但已经成了 HuggingFace 模型分发的事实标准。文件格式：

[8 字节: header 长度]
[N 字节: JSON header — 描述每个 tensor 的 name / dtype / shape / 字节范围]
[剩余: 紧凑 tensor 字节数据，按 header 里描述的偏移排列]

关键设计：

• 零代码执行：纯字节 + JSON，加载就是"按 header 切字节"，格式层面不存在代码执行路径
• 比 pickle 快：没有反序列化开销，直接 mmap
• 零拷贝加载：tensor 直接 view 在 mmap 内存上
• 跨语言友好：Rust / JS / Go 都能读

PyTorch 端用 safetensors：

from safetensors.torch import save_file, load_file

save_file(model.state_dict(), 'model.safetensors')
loaded = load_file('model.safetensors')

HuggingFace Hub 上的现代模型几乎都用 safetensors（如 Llama / Qwen / DeepSeek 的官方权重）。.bin（pickle）格式仍然存在但被标记为"legacy"。

为什么快：pickle 路径要反序列化对象图 + 逐 storage 构造，safetensors 路径只是"读 JSON header → 按 offset 切 mmap 出来的字节"，几乎没有 per-tensor 的 Python 开销。配上 mmap 更是把"加载"推迟到真正访问那一刻。具体倍数随磁盘 / 页缓存状态差异很大，别抄固定数字，自己在目标机器上测。

零代码执行这条则是硬性的格式属性，与性能无关 —— 这也是 vLLM 一类推理引擎偏好 safetensors 的首要原因。

19.5 model.state_dict() vs torch.save(model)

# 推荐
torch.save(model.state_dict(), 'weights.pt')

# 不推荐
torch.save(model, 'whole_model.pt')

差异：

state_dict()	整对象 pickle
只存权重数据	存整个对象（含类定义引用）
加载时要先 build model 类	加载时自动 reconstruct（v2.6+ 须显式 weights_only=False）
修改类定义不影响加载	改类定义可能加载失败
文件小	文件大（含 Python 类元数据）
跨 PyTorch 版本兼容好	跨版本风险

第 9 章 §9.6.2 讨论过这个区别。生产代码强烈建议 state_dict。整对象 pickle 只在快速保存中间状态时偶尔用。

19.6 Distributed Checkpoint (DCP)

70B 模型 ckpt ≈ 280 GB（参数）+ 560 GB（optimizer state）+ ... ≈ 1 TB。如果 rank 0 收集所有 rank 数据再写一个文件：

• rank 0 单机要 1 TB RAM 收集
• rank 0 串行写 1 TB 到磁盘几十分钟
• 加载时反过来 —— 又要几十分钟

DCP（torch.distributed.checkpoint）的解法：每个 rank 只写自己持有的数据，并行写到一个目录：

import torch.distributed.checkpoint as dcp

# 保存
state_dict = {'model': model.state_dict(), 'optim': optimizer.state_dict()}
dcp.save(state_dict, checkpoint_id='ckpt_step_1000')

# 加载 (跨任意 rank 数恢复)
dcp.load(state_dict, checkpoint_id='ckpt_step_1000')

DCP 输出的不是单文件，而是一个目录：

ckpt_step_1000/
  __0_0.distcp                <- rank 0 的数据
  __1_0.distcp                <- rank 1 的数据
  ...
  .metadata                   <- 描述每片数据的逻辑映射

文件名由 _StoragePrefix(f"__{rank}_") + 计数 + DEFAULT_SUFFIX（".distcp"）拼出（torch/distributed/checkpoint/filesystem.py:673/690/97）。注意元信息文件叫 .metadata，不是 metadata.json —— 它是 pickle.dump 出来的 Metadata dataclass（torch/distributed/checkpoint/filesystem.py:69 的 _metadata_fn 定义文件名、torch/distributed/checkpoint/filesystem.py:774 的 pickle.dump 写入），不是 JSON，别指望 json.load 打开。

每 rank 写自己的 __N_0.distcp 文件，完全并行。加载时 DCP 根据 .metadata 把每片正确路由回需要它的 rank（即便加载时 rank 数与保存时不同 —— 自动 reshard）。

这套机制把 rank 0 串行的单点写变成 N 个 rank 的并行写，是 70B+ 训练的工程基础。

19.6.1 DCP 的 reshard 能力

最强的特性：保存时 8 卡 ZeRO-3，加载时 16 卡 ZeRO-3，DCP 能自动重新切分。

实现机制：.metadata 用 逻辑张量坐标（不是物理 rank）描述每片数据 —— MetadataIndex（torch/distributed/checkpoint/metadata.py:154）的键是 fqn + offset，与谁写的它无关。加载时根据当前 rank 拓扑重新映射 —— 保存时 rank 0 持有 [0:1000] 切片，加载时如果 rank 0 现在应该持有 [0:500]，DCP 自动切。

这种"加载时 reshard"让训练流水线灵活得多 —— 训练用 64 卡，fine-tuning 切到 8 卡只需重新加载，不需要离线 resharding。

19.6.2 DCP 的 SavePlan / LoadPlan

torch/distributed/checkpoint/planner.py:105/115 的 SavePlan / LoadPlan 是 DCP 的 IR：

@dataclass(frozen=True)
class SavePlan:
    items: list[WriteItem]
    storage_data: Any = None           # backend-specific 数据 (FileSystemWriter 放 _StoragePrefix)
    planner_data: Any = None           # planner 状态
    # 标记本 rank 可否复用上一次缓存的 plan
    usable: bool = True

@dataclass
class WriteItem:
    index: MetadataIndex               # 逻辑索引 (fqn + offset)
    type: WriteItemType                # TENSOR / SHARD / BYTE_IO
    bytes_io_data: BytesIOWriteData | None = None
    tensor_data: TensorWriteData | None = None

LoadPlan 装的是 ReadItem（torch/distributed/checkpoint/planner.py:86），字段更能说明 reshard 的本质：dest_index / dest_offsets（写到本地哪个张量的哪个偏移）、storage_index / storage_offsets（从 ckpt 的哪个 chunk 的哪个偏移读）、lengths（要拷的超立方体尺寸）。一个 ReadItem 就是一次"超立方体子区间拷贝"。

DCP save 流程：

1. 每个 rank 调 planner.create_local_plan() 生成自己的 SavePlan（torch/distributed/checkpoint/default_planner.py:107）
2. coordinator rank 收集所有 SavePlan、调 planner.create_global_plan(plans)（torch/distributed/checkpoint/default_planner.py:213）做去重（_dedup_save_plans.py）+ 优化，同时产出 Metadata
3. 把全局 plan 分发回各 rank
4. 每个 rank 按 plan 调 storage_writer.write_data(plan, planner) 写自己负责的部分
5. coordinator 调 storage_writer.finish(metadata, results) 写 .metadata

LoadPlan 反向：每 rank 描述"我需要哪些 tensor"，DCP 根据 .metadata 路由到正确的物理 chunk。

这套 plan-based 设计让 DCP 能换 storage 后端：默认 FileSystemWriter（写本地 / NFS），也可换成写 S3、GCS 等对象存储的后端（PyTorch 内置的是 fsspec 版，见下文；专用后端由第三方提供）。实现 StorageWriter / StorageReader 接口即可，业务层无感知。

FileSystemReader（torch/distributed/checkpoint/filesystem.py:837）与 FileSystemWriter（:965，继承自 :574 的 _FileSystemWriter 与 BlockingAsyncStager）是默认实现 —— 写到目录里若干 .distcp 文件 + 一个 .metadata。没显式传 writer 时，_storage_setup（_storage_utils.py）会按 checkpoint_id 自动在 FileSystemWriter 与 FsspecWriter（torch/distributed/checkpoint/_fsspec_filesystem.py:96）之间挑一个 —— 这就是 checkpoint_id='s3://...' 能直接工作的原因。

19.6.3 Resharding 算法：8 卡 ckpt 加载到 16 卡

§19.6.1 提到 DCP 能 reshard，具体怎么做？看 torch/distributed/checkpoint/resharding.py（注意没有下划线前缀）与 _dedup_save_plans.py。前者只有两个函数，恰好就是整套 reshard 的算术核心：_check_shard_metadata_pair_overlap（两个 shard 的超立方体有没有交集）与 _shards_get_overlap_region_wrt_saved_tensor（交集区域在源 / 目标里各自的偏移与长度）。

场景：保存时 8 卡 FSDP，每 rank 持有 weight 的 1/8 切片；加载时 16 卡 FSDP，每 rank 应该持有 1/16 切片。算法：

graph TB
    Save["保存时 (8 卡):<br/>rank 0: weight[0:N/8]<br/>rank 1: weight[N/8:2N/8]<br/>..."]
    Save --> M[".metadata:<br/>'weight' chunks 用逻辑坐标:<br/>chunk[0:N/8], chunk[N/8:2N/8], ..."]
    M --> Load["加载时 (16 卡):<br/>每 rank 算自己需要 weight[i*N/16:(i+1)*N/16]"]
    Load --> P{逻辑切片<br/>vs 物理 chunk 关系}
    P --> Calc["rank 0 需要 weight[0:N/16]<br/>= 物理 chunk[0:N/8] 的前一半"]
    Calc --> Read[读出 rank 0 存的整个 chunk<br/>再 narrow 出前一半]
    style M fill:#fef3c7
    style Calc fill:#dbeafe

具体步骤（入口是 torch/distributed/checkpoint/state_dict_loader.py:267 的 planner.create_local_plan()，落到 default_planner.py 的 create_default_local_load_plan，再由 torch/distributed/checkpoint/planner_helpers.py 的 _create_read_items 调下面两个函数）：

1. 从 .metadata 拿到所有物理 chunk 的逻辑坐标范围
2. 对当前 rank 需求的范围，用 _check_shard_metadata_pair_overlap 找出与哪些物理 chunk 有交集
3. 为每个交集用 _shards_get_overlap_region_wrt_saved_tensor 算出 ReadItem：storage_index / storage_offsets 说从哪读，dest_index / dest_offsets 说写到哪，lengths 说拷多大
4. storage_reader.prepare_local_plan / prepare_global_plan 是留给后端对这批 ReadItem 做 IO 重排的钩子；默认的 FileSystemReader 原样返回，只在 read_data 里按文件分组读，"按 offset 排序"在源码里还是一条 TODO

特殊情况处理：

• 新 rank 数 > 旧：每物理 chunk 被多个 rank 共享读（只读，不冲突）
• 新 rank 数 < 旧：每 rank 读多个物理 chunk 拼起来
• 切分维度改变：交集算术本身是按超立方体做的、并不限定维度，能正确加载；但把 dim-0 分片重排成 dim-1 分片会让每个目标 shard 都与全部源 chunk 相交 —— 每个 rank 都要读该张量的全量再切，IO 量成倍放大

reshard 的用时取决于要读多少字节、跨多少文件、走什么存储，没有可搬运的固定数字。关键性质是：不需要先 gather 成全量张量再切，直接按 chunk 交集重新映射。

理解 reshard 让你看到 DCP 不是"分布式 save/load"那么简单，而是"逻辑坐标系统 + 物理 chunk 路由"的工程。这是大模型生命周期管理（pretrain → fine-tune → distill）的关键基础设施。

19.6.4 Async Checkpoint：训练不被 ckpt 阻塞

dcp.async_save（torch/distributed/checkpoint/state_dict_saver.py:220）让 ckpt 写入与训练并发。注意它的 docstring 仍标着 .. warning:: This feature is experimental and subject to change.，且要求 最后一次 ckpt 之后必须 close：

import torch.distributed.checkpoint as dcp

# 同步版本: 训练等 ckpt 写完才继续
dcp.save(state_dict, ...)
# next training step starts...

# 异步版本: 训练立刻继续, 后台慢慢写
future = dcp.async_save(state_dict, ...)
# next training step starts immediately
# ...
future.result()  # 偶尔等一次, 确认上次 ckpt 写完

它的签名里 async_checkpointer_type 默认是 AsyncCheckpointerType.THREAD（后台线程），也可以切到进程实现（_async_process_executor.py）以躲开 GIL；async_stager 则决定 staging 放在哪（默认 CPU 内存）。

实现机制：

graph LR
    T[训练主线程] --> CP[snapshot 整个 state_dict<br/>1. 把 GPU tensor copy 到 CPU（默认同步完成）<br/>2. 用同一份 CPU 内存的 ref 创建 staging dict]
    CP --> RT[训练继续<br/>下一 step]
    CP --> BG[后台线程<br/>写 staging dict 到磁盘]
    BG --> Done[完成 future]
    style CP fill:#fef3c7
    style BG fill:#dcfce7

关键点：staging 阶段必须同步（要确保后台写的是当前 step 状态、不是几 step 后的）。staging 把 GPU 数据拷到 CPU 内存，量级是一次 D2H 拷贝；之后训练继续，后台线程 / 进程慢慢把 CPU 数据写盘，量级是一次磁盘（或对象存储）写。后者通常慢得多，这就是 async 的全部收益来源。

同步 ckpt 的吞吐损失约等于 写盘时间 / (ckpt 间隔步数 × 单步时间)；async 把分子从"写盘时间"换成"D2H 时间"。间隔越密、单步越快，收益越明显。

代价：staging 要额外一份与 ckpt 等大的 CPU 内存。如果 CPU 内存紧张，不能开 async。

理解 async ckpt 让你看 v2.x 大模型训练的真实工程取舍：宁愿多吃 CPU 内存也不让 ckpt 写阻塞训练，因为 GPU 时间贵得多。

19.6.5 自定义 DCP storage backend：写一个云存储后端

StorageWriter / StorageReader 是 DCP 的扩展点。第三方实现这两个抽象类即可让 DCP 写到任意存储。简化示例：

StorageWriter（torch/distributed/checkpoint/storage.py:27 起）要求实现的方法是：reset、set_up_storage_writer、prepare_local_plan、prepare_global_plan、write_data、finish、类方法 validate_checkpoint_id，外加可选的 storage_meta。StorageReader（:168 起）对称：reset、read_metadata、set_up_storage_reader、prepare_local_plan、prepare_global_plan、read_data、validate_checkpoint_id。

下面把抽象方法都写全了（prepare_*_plan 只做透传），用一个假的 boto3 桩在 CPU torch 2.11.0 上实跑过保存流程；真 S3 未测。

下面是读者示例，不是 PyTorch 源码原文——MyS3Writer 是示范用的第三方实现：

import io, pickle
import boto3
import torch
import torch.distributed.checkpoint as dcp
from torch.futures import Future
from torch.distributed.checkpoint.storage import StorageWriter, WriteResult

class MyS3Writer(StorageWriter):
    def __init__(self, bucket: str, prefix: str):
        super().__init__()
        self.bucket = bucket
        self.prefix = prefix
        self.s3 = boto3.client('s3')

    def reset(self, checkpoint_id=None):
        self.checkpoint_id = checkpoint_id

    def set_up_storage_writer(self, is_coordinator, *args, **kwargs):
        # rank 0 (coordinator) 创建 prefix 目录
        if is_coordinator:
            self.s3.put_object(Bucket=self.bucket,
                               Key=f"{self.prefix}/{self.checkpoint_id}/")

    def prepare_local_plan(self, plan):
        return plan

    def prepare_global_plan(self, plans):
        return plans

    def write_data(self, plan, planner):
        results = []
        for item in plan.items:
            data = planner.resolve_data(item)      # torch.Tensor 或 io.BytesIO
            if isinstance(data, torch.Tensor):
                buf = io.BytesIO()
                torch.save(data.detach().cpu().contiguous(), buf)
                data = buf
            body = data.getvalue()
            # 同一 fqn 的不同分片 offset 不同, key 里必须带上 offset, 否则互相覆盖
            offset = "_".join(map(str, item.index.offset or ()))
            key = f"{self.prefix}/{self.checkpoint_id}/{item.index.fqn}@{offset}"
            self.s3.put_object(Bucket=self.bucket, Key=key, Body=body)
            results.append(WriteResult(index=item.index,
                                       size_in_bytes=len(body),
                                       storage_data=key))
        fut = Future()
        fut.set_result(results)                    # 返回已完成、携带 WriteResult 的 future
        return fut

    def finish(self, metadata, results):
        # index -> 对象 key 的映射记进 Metadata, 加载时 StorageReader 靠它找数据
        metadata.storage_data = {wr.index: wr.storage_data
                                 for wrs in results for wr in wrs}
        # 写 .metadata (Metadata 是 dataclass, 内置实现用 pickle 而非 JSON)
        self.s3.put_object(
            Bucket=self.bucket,
            Key=f"{self.prefix}/{self.checkpoint_id}/.metadata",
            Body=pickle.dumps(metadata),
        )

    @classmethod
    def validate_checkpoint_id(cls, checkpoint_id):
        return True

# 用法: checkpoint_id 要传, DCP 靠它调 reset()
dcp.save(state_dict, checkpoint_id='step_1000',
         storage_writer=MyS3Writer('my-bucket', 'training/run_1'))

实战中 PyTorch 已经内置 fsspec-based 的 FsspecWriter / FsspecReader（torch/distributed/checkpoint/_fsspec_filesystem.py:96/154），装了 fsspec + 对应后端（s3fs、gcsfs 等）就能直接用 s3:// / gs:// 的 checkpoint_id。自家写 backend 仅在特殊需求时用：

• 自家对象存储（如内部 ceph cluster）
• 加密存储（写前 encrypt）
• 跨区域复制（写多个 region 提高可用性）
• 进度监控（每 chunk 写完发 metric 到监控系统）

理解 storage backend 的接口让你看 DCP 不是绑死本地文件系统的，是真正的 pluggable。生产场景如果对 ckpt 存储有特殊需求，扩展这层就够、不需要 fork DCP 自身。

19.6.6 fsspec：直接读写云存储

先纠正一个常见误解：torch.save / torch.load 本身不认识 s3:// 这种 URI —— torch/serialization.py 里没有任何 fsspec 集成，传字符串路径就是当本地路径开。要写云存储，自己拿 file-like 对象：

import fsspec, torch

with fsspec.open('s3://my-bucket/ckpts/step_1000.pt', 'wb') as f:
    torch.save(state, f)          # torch.save 接受 file-like 对象
with fsspec.open('s3://my-bucket/ckpts/step_1000.pt', 'rb') as f:
    state = torch.load(f)

真正内置 fsspec 的是 DCP：torch/distributed/checkpoint/_fsspec_filesystem.py 提供 FsspecWriter / FsspecReader，而 _storage_utils.py 的 _storage_setup 会按 checkpoint_id 自动挑 —— FileSystemWriter.validate_checkpoint_id 不认的路径就交给 FsspecWriter。

实战注意点：

• 认证：要先配好 boto3 / gcloud CLI 让 fsspec 用上 credentials
• 大文件：s3fs 会自动走 multipart upload，分段大小由它自己的配置决定；上传耗时基本就是 ckpt 字节数 / 出口带宽
• list / glob 慢：对象存储的 list 是分页 API，目录里文件一多就慢；DCP 直接读 .metadata 拿文件清单，不做 globbing
• 一致性：S3 自 2020 年底起对 PUT 后的 GET / LIST 提供 strong read-after-write consistency，老代码里那种 "save 后 sleep 几秒再通知" 的 workaround 已经不必要；但自建对象存储不一定有这个保证，按自家存储的文档办

所以这一行是能直接跑的（前提是装了 fsspec + s3fs 并配好凭证）：

dcp.save(state_dict, checkpoint_id='s3://bucket/ckpt/step_1000')

让每 rank 直接写对象存储、不需要本地落盘 + 后台同步。云原生训练的标准模式。

19.7 训练侧的 checkpoint 工程

19.7.1 断点续训：随机状态完整保存

光保存 model + optim 还不够。完整可恢复 ckpt 必须包括：

ckpt = {
    'model': model.state_dict(),
    'optim': optimizer.state_dict(),
    'scheduler': lr_scheduler.state_dict(),
    'dataloader': loader.state_dict(),       # 需 torchdata 的 StatefulDataLoader; 原生 Sampler / DataLoader 没有 state_dict, 见 §11.23
    'epoch': epoch,
    'step': step,
    'rng_states': {
        'cpu': torch.get_rng_state(),
        'cuda': torch.cuda.get_rng_state_all(),  # 所有 GPU 各自
        'numpy': np.random.get_state(),
        'python': random.getstate(),
    },
    'amp_scaler': scaler.state_dict() if scaler else None,
}
torch.save(ckpt, 'full_ckpt.pt')

为什么 RNG 也要存？两个场景：

• augment 复现：RandomCrop 用 torch 的随机数决定 crop 位置，rng 不一致 → 同一 sample 生成不同 augmented view → 不可复现
• dropout / noise injection：reproducibility 要求

torch.set_rng_state 加载时恢复 CPU RNG；torch.cuda.set_rng_state_all 恢复每个 GPU 的 RNG。但 DataLoader worker 的 RNG 是另一套——每个 worker 的种子是 base_seed + worker_id（torch/utils/data/_utils/worker.py:258），而 base_seed 是创建迭代器时从主进程 RNG（或 loader.generator）抽的（torch/utils/data/dataloader.py:706）。所以只要在重建迭代器之前恢复好主进程 RNG，worker 的种子也就跟着复现。

实战：HuggingFace Trainer / Accelerate 内置完整 RNG 保存。如果你自写训练循环，要么照抄上面的字典、要么直接用这些库。漏一个 RNG 类型就会导致"resume 后 loss 跳变"——很难诊断的 bug。

19.7.2 Optimizer state：被忽视的 ckpt 大头

很多人保存 ckpt 只想到 model.state_dict，但Adam 的 optimizer state 是 fp32 参数体积的 2 倍，是 bf16 权重体积的 4 倍。

Adam 优化器每个参数维护：

• exp_avg（一阶矩）：zeros_like(p)，与参数同 shape、同 dtype
• exp_avg_sq（二阶矩）：同上
• step（当前步数）：0 维 tensor

混合精度训练里 optimizer 管的是 fp32 主权重，Adam state 随之也是 fp32：70B × 4 字节 × 2 = 560 GB；而导出的 bf16 权重只有 140 GB。只看这两项，Adam state 占 4/5。

# state_dict 的真实结构 (Adam)
{
    'state': {
        0: {'exp_avg': tensor[1024,1024], 'exp_avg_sq': tensor[1024,1024], 'step': tensor(1000.)},
        1: {'exp_avg': ..., 'exp_avg_sq': ..., 'step': tensor(1000.)},
        ...                    # 每个参数对应一个 dict
    },
    'param_groups': [...]      # 优化器配置 (lr, betas, weight_decay 等)
}

param_groups 用整数索引引用 state 中的参数。这种"间接寻址"让 optimizer 与 model 解耦，但加载时必须保证 model 参数顺序与保存时一致——否则 state[0] 错配到 weight[1]，训练立刻崩盘。

DCP 也支持优化器 state ckpt 与 reshard。FSDP-2 的 optimizer state 按 rank sharded → DCP 保存每 rank 自己那 1/N → reload 时按当前 rank 数重新 reshard。这是大模型训练能从 1024 卡 fine-tune 到 64 卡推理的根本前提。

8-bit Optimizer（如 bitsandbytes 的 8-bit Adam）：把 fp32 的 exp_avg / exp_avg_sq 量化到 8-bit 存储，这两项的体积按位宽算降到 1/4 —— 上面那 560 GB 变成 140 GB。代价是优化器状态的数值精度下降，要靠 block-wise 量化 + 动态 scale 补回来。

19.7.3 PEFT / LoRA 的 ckpt 增量保存

LoRA / Adapter / Prefix-tuning 等参数高效微调（PEFT）的核心是"只训练少量参数"。ckpt 也只需保存这些参数：

# 全量 fine-tune: state_dict 含所有 7B 参数
torch.save(model.state_dict(), 'full_ft.pt')   # ~14 GB

# LoRA: 只存 LoRA 模块的参数
lora_state_dict = {k: v for k, v in model.state_dict().items() if 'lora' in k}
torch.save(lora_state_dict, 'lora.pt')          # ~50 MB

peft 库（HuggingFace）做了完整封装：model.save_pretrained() 自动只保存 LoRA 权重 + adapter_config.json（描述秩、target_modules 等）。

加载流程：

1. 先加载 base model（如 Llama-7B 全量权重）
2. 加载 adapter_config.json 知道在哪些 module 上加 LoRA
3. 加载 lora.pt 把 LoRA 权重塞进对应 module

这种"base model + small adapter"模式的工程意义：

• 存储成本：1 个 base + 100 个 adapter（每个 50 MB）= 14 GB + 5 GB = 19 GB；vs 100 个全量 = 1.4 TB
• 分发：用户可单独下载 adapter，不重复下 base
• 服务：vLLM 等支持 base + 多 adapter 共存推理（Multi-LoRA Serving），只切换 adapter 不重 load base

对应到 PyTorch 序列化层面：state_dict 是 dict，支持任意子集 —— 这是 LoRA 这种增量保存的基础设施。load_state_dict(strict=False) 允许部分 key 不存在 → 只覆盖匹配的 key。这层设计让 PEFT 能优雅落地。

19.7.4 ckpt rollback：训练失败的恢复策略

长时间训练偶尔出现 NaN loss / divergence / OOM 等问题，需要 rollback 到更早 ckpt。生产标准做法：

graph LR
    Step1[step 1000<br/>ckpt A] --> Step2[step 2000<br/>ckpt B]
    Step2 --> Step3[step 3000<br/>ckpt C]
    Step3 --> Step4[step 4000<br/>NaN loss!]
    Step4 -.rollback.-> Step3
    Step3 --> Step5[step 4000 retry<br/>不同 lr / batch]
    style Step4 fill:#fee2e2
    style Step5 fill:#dcfce7

工程做法：

• 保留多个最近 ckpt：rolling window，如最近 5 个 step 的 ckpt
• NaN detection：每 step 后检查 loss.isnan()，立刻 abort
• 自动 rollback：检测到 NaN → 自动 load 上一个 ckpt → 调小 lr 重试 → 继续

LLM 训练框架普遍提供其中的一部分，比如 HuggingFace Trainer 的 save_total_limit（保留最近几个）与 resume_from_checkpoint（续训）；NaN 检测后的自动回滚通常要自己接。

对应到 PyTorch 序列化层面：ckpt 路径要含 step 号方便排序：

ckpt_step_1000/
ckpt_step_2000/
ckpt_step_3000/
ckpt_step_4000/    <- corrupted, skip

rollback 脚本逻辑：

import glob

def find_latest_valid():
    ckpts = sorted(glob.glob('ckpt_step_*'),
                   key=lambda x: int(x.split('_')[-1]))
    for c in reversed(ckpts):
        if validate(c):    # checksum + 加载 + 跑一个 batch 不出 NaN
            return c
        print(f"corrupted: {c}, trying older")
    return None

这种"层层回滚找最近完好版本"是大模型训练的运维标准动作。理解这层让你看到 ckpt 不只是"备份"，是 training infra 的核心 control plane。

19.7.5 一个 70B 训练 ckpt 配置参考

把全章的工程实践合起来，给一个 70B 训练 ckpt 的参考骨架（list_s3_dir、init_state_dict、validate_loaded 等是占位函数）：

import torch
import torch.distributed.checkpoint as dcp
from datetime import timedelta

CKPT_INTERVAL = 200            # 每 200 step 保存一次
CKPT_RETAIN = 5                # 保留最近 5 个

CKPT_DIR = "s3://my-org-bucket/training/run_xxx/ckpts"

def save_checkpoint(model, optimizer, scheduler, step, loader):
    state_dict = {
        'model': model.state_dict(),
        'optim': optimizer.state_dict(),
        'scheduler': scheduler.state_dict(),
        'dataloader': loader.state_dict(),   # torchdata StatefulDataLoader, 见 §19.7.1
        'step': step,
        # 注意: 普通 tensor 在 DCP 里按"各 rank 相同"去重, 只存一份;
        # 各 rank 不同的 RNG 状态要放进带 rank 号的 key, 否则加载后全是同一个 rank 的值
        f'rng_states_rank{torch.distributed.get_rank()}': {
            'cpu': torch.get_rng_state(),
            'cuda': torch.cuda.get_rng_state_all(),
        },
    }

    # async + DCP + 写 S3
    # async_save 要求进程组带 CPU 后端, 如 init_process_group("cpu:gloo,cuda:nccl")
    future = dcp.async_save(
        state_dict,
        checkpoint_id=f"{CKPT_DIR}/step_{step}",
    )
    return future

def load_latest_checkpoint(model, optimizer, scheduler, loader):
    ckpts = list_s3_dir(CKPT_DIR)
    for ckpt_id in sorted(ckpts, key=lambda x: int(x.split('_')[-1]), reverse=True):
        try:
            state_dict = init_state_dict(model, optimizer, scheduler, loader)
            dcp.load(state_dict, checkpoint_id=ckpt_id)
            validate_loaded(state_dict)        # checksum + 一个 dummy forward
            return state_dict
        except Exception as e:
            print(f"corrupted: {ckpt_id}, trying older: {e}")
    raise RuntimeError("no valid ckpt found")

def cleanup_old_checkpoints():
    ckpts = sorted(list_s3_dir(CKPT_DIR), key=lambda x: int(x.split('_')[-1]))
    for old in ckpts[:-CKPT_RETAIN]:
        delete_s3_dir(old)

主训练循环：

ckpt_future = None
for step, batch in enumerate(loader):
    if step % CKPT_INTERVAL == 0:
        if ckpt_future is not None:
            ckpt_future.result()              # 等上一个 ckpt 写完
        ckpt_future = save_checkpoint(model, optim, sched, step, loader)
        cleanup_old_checkpoints()

    loss = train_step(batch)
    if torch.isnan(loss):
        # NaN 检测 + 自动 rollback
        state = load_latest_checkpoint(model, optim, sched, loader)
        # 调小 lr 重试
        for g in optim.param_groups:
            g['lr'] *= 0.5

这套配置覆盖了：

• DCP（千卡 reshard 友好）
• async（不阻塞训练）
• S3（云原生，无需本地落盘）
• rolling retention（控制存储成本）
• NaN guard（自动 rollback）
• RNG 完整保存（resume 可复现）
• checksum 验证（silent corruption 防御）

理解每条配置的 why，让你能为自家训练任务量身定制。

19.8 文件格式与底层写入

19.8.1 PyTorchFileReader / PyTorchFileWriter

torch._C.PyTorchFileReader / PyTorchFileWriter（C++ 实现在 caffe2/serialize/inline_container.cc，类名是 PyTorchStreamReader / PyTorchStreamWriter）是 zip 容器的底层，内部用 miniz。写入器技术上支持压缩（writeRecord 有个 compress 形参，走 MZ_BEST_COMPRESSION），但 torch.save 全程不用它 —— 头文件注释把 All files are stored uncompressed 直接列为容器的硬保证。为什么：

• 训练好的浮点权重字节接近随机，DEFLATE 压不动多少
• 压缩要多一遍全量数据的 CPU 处理，save / load 都变慢
• 压缩会直接杀死 mmap —— 压缩过的 entry 不能当内存映射用

配套的第二条保证是对齐：所有 entry 对齐到 64 字节边界（getPadding 塞 padding），这样 mmap 之后拿到的 tensor 数据指针本身就是对齐的。第三条是统一写 ZIP64（见 §19.8.2）。

这三条合起来才让 torch.load(..., mmap=True) 真的能零拷贝工作。

19.8.2 大文件支持：ZIP64 与 stream

标准 zip 单文件最大 4 GB（header offset 是 32-bit）。大模型 ckpt 远超这个限制。PyTorch 的做法是 无条件写 ZIP64 —— inline_container.h 的注释原话是 We universally write in ZIP64 format for consistency，而不是"文件大了才切换"：

• 文件大小用 64-bit 表示
• entry offset 也用 64-bit
• 外层结构兼容标准 zip（小文件仍可被 unzip 工具打开）

所以哪怕只存一个几 KB 的小 tensor，.pt 也是 ZIP64 格式。好处是大文件不用特判、.pt 的读写路径只有一条；代价是老旧或裁剪过的 zip 工具可能读不了。遇到命令行 zip 工具报错时别急着怀疑文件坏了，先换 Python 的 zipfile（原生支持 ZIP64）试一下。

诊断 .pt 文件时用 Python 的 zipfile 模块（强制 ZIP64 兼容）：

import zipfile
with zipfile.ZipFile('model.pt') as z:
    print(z.namelist())                    # 列文件
    print([info.file_size for info in z.infolist()])

这层细节不常用，但调 ckpt 问题时能避免"为什么 unzip 看不了"的困惑。

19.8.3 storage 共享检测：避免存重复字节

PyTorch 张量经常共享 storage（第 2 章 §2.6 view 机制）。torch.save 不会傻傻地把每个 tensor 的字节都存一份 —— 它做了storage 共享检测：

a = torch.randn(10, 10)
b = a.view(100)             # b 与 a 共享 storage
c = a[0:5]                  # c 也共享

state = {'a': a, 'b': b, 'c': c}
torch.save(state, 'shared.pt')
# 文件只含 1 份 storage 字节, 不是 3 份

实现机制是 pickle 的 persistent_id / persistent_load 钩子（torch/serialization.py:1200 的 persistent_id）。去重的那一行很直白：

storage_key = id_map.setdefault(storage._cdata, str(len(id_map)))

以 storage 的 C 层地址 _cdata 为键查表：同一份 storage 无论被多少个 tensor 引用，都只拿到一个 key，zip 里也就只写一个 data/<key> entry（:1280 起的循环）。load 时 persistent_load 按 key 取字节构造 storage，多个 tensor 再共用这一份。

这套机制对 view-heavy 的模型（比如把一个大 weight 切成 Q / K / V 三个 view）能省下成倍的空间。

但有个反直觉的坑：view 张量也存了它的"完整 base storage"，不是它自己的子集。如果你 b = a[0:5]; del a; torch.save({'b': b})，文件大小是整个 a 的，不是 5 个元素。要省空间得 b = a[0:5].clone() 让 b 有独立 storage。

19.8.4 自定义 reduce：让自家类能被 save

如果你的 state_dict 含自定义类（如 LR scheduler 的 metadata），默认 pickle 可能失败。解法：实现 __reduce__：

class MyScheduler:
    def __init__(self, base_lr, warmup_steps):
        self.base_lr = base_lr
        self.warmup_steps = warmup_steps
        self._cache = ...    # 不想 pickle 的私有状态

    def __reduce__(self):
        # pickle 时只保存关键参数, 加载时重建
        return (MyScheduler, (self.base_lr, self.warmup_steps))

__reduce__ 返回 (callable, args)。pickle 调用时调 callable(*args) 重建对象。这种"自定义重建"让你完全控制序列化字段。

但 weights_only=True（v2.6+ 默认）下 自定义 __reduce__ 不被信任 —— 出于安全考虑只允许白名单类型。要让自定义类支持加载：

import torch.serialization
torch.serialization.add_safe_globals([MyScheduler])

# 现在可以加载
state = torch.load('ckpt.pt')

或者用 safetensors 替代 torch.save —— 它根本不支持 Python 对象（只支持 tensor 字典），这种"局限性"反而让安全成为默认。

实战建议：state_dict 应该只含 tensor 与原生 Python 类型（int/float/str/list/dict），不含自定义类。如果非要保存复杂对象（如 scheduler 的 internal state），用 JSON 序列化简单 metadata、复杂逻辑在加载侧重建。这是从 pickle 安全 + 跨版本兼容角度都最好的实践。

19.8.5 序列化的内存峰值

很多人以为 torch.save(state_dict, path) 内存占用 = state_dict size。实际峰值更高：

graph LR
    SD[state_dict<br/>GPU 上 1 份] --> Cpu[逐个 storage .cpu&#40;&#41;<br/>CPU 上多出临时副本]
    Cpu --> Zip[write_record 直接写进 zip<br/>不经过 pickle]
    Zip --> File[最终文件<br/>约等于 state_dict 大小]
    style Cpu fill:#fef3c7
    style Zip fill:#dcfce7

图里要强调的一点：v2.11 的 _save（torch/serialization.py:1184 起）不把 tensor 字节塞进 pickle 流。pickle 里只有 persistent_id 返回的元组（storage 类型 / key / device / numel），tensor 字节由 :1313 的 zip_file.write_record(name, storage, num_bytes) 逐个直接写进 zip。所以并不存在"整份 state_dict 的 pickle 副本"这一格。

真正的峰值来源在别处：

• GPU → CPU 拷贝：CUDA storage 要先落到 CPU 才能写盘（:1289 的注释明说 given that we copy things around anyway）。这一项的峰值是"同时在手的那几份"，不是全部之和
• state_dict() 自身：如果构造时做了 .clone() 或 dtype cast，那是实打实的又一整份
• 写云存储：fsspec 一类实现可能把整段 buffer 攒在内存里再上传

想只写元数据、连字节都不落，用 torch.serialization.skip_data()（:385）—— 它让 :1286 的 write_record_metadata 只占位不写内容。

DCP 的 async save 是另一条降峰值的路：staging 占一份 CPU 内存，写盘后释放，且每 rank 只负责自己那 1/N。

19.9 跨格式与跨生态

19.9.1 sharded safetensors：单文件大小限制怎么处理

safetensors 的格式本身对单文件大小没有硬限制，但托管平台通常会限制单文件体积，且大文件断点续传体验差 —— 所以大模型一律切片。切分逻辑（transformers 侧实现，不在 PyTorch 里）：

1. 按 fqn 顺序遍历所有 tensor
2. 累计 size > shard_size_limit (默认 5GB) 时切到下一个文件
3. 每个 tensor 不跨文件（一个 tensor 必须完整在一个文件里）
4. 写完所有文件后生成 model.safetensors.index.json

注意点：

• shard 大小：太小 → 文件数多、每个文件的请求开销被放大；太大 → 单文件下载失败要重来、也容易撞平台限制。几 GB 到十几 GB 是常见区间
• 顺序敏感：fqn 顺序变化会让 sharding 结果不一致 → ckpt 不可复现。HF 用 sorted(state_dict.keys()) 保证一致
• 加载时按 index 找：每加载一个 weight，从 index 反查在哪个文件、再从该文件 mmap → 一次只读必要的字节

下载耗时基本就是 总字节数 / 有效带宽，并发只能把单连接的带宽上限撑开 —— 这就是 from_pretrained 首次调用要等的根本原因，与反序列化本身无关。

19.9.2 GGUF / GGML：推理优化的 weight 格式

safetensors 是训练 / 通用分发场景，端侧 / CPU 推理有更专门的格式：GGUF（GGML 的下一代）。它专为 llama.cpp 等 CPU/Apple Silicon 推理框架设计：

维度	safetensors	GGUF
量化支持	不内置	内置（int4/int5/int8 等多种）
元信息	tensor 名 + shape	加上 chat template、stop tokens、tokenizer 配置
加载方式	mmap 友好	mmap 友好 + 端侧推理优化
主要生态	HuggingFace Hub	llama.cpp / Ollama 生态
适用场景	训练 / 服务器推理	CPU / 端侧推理

PyTorch 这层不直接支持 GGUF，但 llama.cpp 仓库提供了从 HuggingFace 模型转 GGUF 的脚本与量化工具（命令以 llama.cpp 官方 tools/quantize/README.md 为准）：

python convert_hf_to_gguf.py ./hf-model-dir --outfile model-bf16.gguf --outtype bf16
./build/bin/llama-quantize model-bf16.gguf model-Q4_K_M.gguf Q4_K_M

转换后用 llama.cpp 加载、在 CPU / Mac M-series 上跑推理。这条路径让 PyTorch 训练的模型能落地到端侧 —— 训练用 PyTorch、部署给社区用 GGUF。

理解 GGUF 在 PyTorch 生态外的存在让你看到 model 序列化是个多格式共存的世界。每种格式服务一类场景：safetensors 通用、.pt2 (AOTI) 服务器推理、.pte (ExecuTorch) 移动端、GGUF 端侧 CPU。PyTorch 自身只覆盖前几个，社区补足后两个。

19.9.3 ONNX 互操作：跨框架部署的中间格式

ONNX（Open Neural Network Exchange）是跨框架部署的标准 IR。PyTorch 的 ONNX 导出：

import torch.onnx

torch.onnx.export(
    model,
    args=(example_input,),
    f="model.onnx",
    input_names=['input'],
    output_names=['output'],
    dynamic_axes={'input': {0: 'batch'}, 'output': {0: 'batch'}},
    opset_version=18,
)

v2.11 的 API 已经收敛成一个：torch.onnx.dynamo_export 已被移除（torch/onnx/__init__.py 的 __all__ 里只剩 export / is_in_onnx_export 等），两条路径合并成 torch.onnx.export 的一个开关：

torch.onnx.export(model, args, "m.onnx", dynamo=True)   # 默认

• dynamo=True（torch/onnx/__init__.py:78，v2.11 的默认值）：走 torch.export / Dynamo 抓图再转 ONNX。动态形状用 dynamic_shapes= 描述（与 torch.export 同一套语法），dynamic_axes 只作为 BC 选项保留
• dynamo=False：回落到基于 TorchScript 的老导出器（源码在 torch/onnx/_internal/torchscript_exporter/），已标注 deprecated

也就是说上面那段示例里的 dynamic_axes= 写法仍能跑（BC 选项），但新代码应该改用 dynamic_shapes=。

ONNX 的价值：

• 跨框架推理：训练用 PyTorch，推理用 TensorRT / OpenVINO / Core ML / DirectML
• 嵌入式 / 边缘：Coral / Jetson 等支持 ONNX runtime
• 优化器生态：ONNX Runtime 自带 graph optimization、量化 toolkit

代价：

• 算子覆盖度：PyTorch 几千个 op，ONNX 标准化几百个。冷门 op 导出失败
• dynamic shape 支持有限：导出时要用 dynamic_shapes（老路径是 dynamic_axes）显式声明，否则 shape 写死
• trace 不支持复杂控制流：含 if x.sum() > 0 的 model 用 trace 导出会丢分支

实战建议：首选 .pt2 (AOTI) 服务器推理、ExecuTorch 移动端、ONNX 仅作为"必须用 TensorRT/OpenVINO" 的桥梁。v2.11 的 ONNX 导出本身也建立在 torch.export 之上，抓图这一步与 AOTI / ExecuTorch 共用。

19.9.4 跨格式转换矩阵

PyTorch 生态有多种 ckpt 格式，转换关系：

源格式 → 目标	工具	备注
.pt (pickle) → .safetensors	safetensors 库	安全 + 加载快，强烈推荐
.safetensors → .pt (pickle)	safetensors 库	兼容老代码
.pt / .safetensors → .pt2 (AOTI)	torch._inductor.aoti_compile_and_package	服务器推理
.pt / .safetensors → .pte (ExecuTorch)	executorch 库	移动端 / 嵌入式
.safetensors → ONNX	先加载成模型再 torch.onnx.export	跨框架
.safetensors → GGUF	llama.cpp 的 convert_hf_to_gguf.py	CPU / Mac 推理
.bin（HF legacy）→ .safetensors	transformers 库	HF 现在默认输出后者
TF SavedModel → .safetensors	自家工具	偶尔用，质量参差

没有"通用万能格式"。每种服务一类场景：

graph TB
    Train[训练阶段] --> Pt[.pt / .safetensors]
    Pt -->|服务器推理| Aoti[.pt2 AOTI]
    Pt -->|移动端| Pte[.pte ExecuTorch]
    Pt -->|跨框架| Onnx[ONNX]
    Pt -->|CPU 推理| Gguf[GGUF]
    Pt -->|分发| Hub[HuggingFace Hub]
    style Pt fill:#fef3c7
    style Aoti fill:#dcfce7
    style Pte fill:#dbeafe
    style Gguf fill:#fce7f3

工程决策树：

• 训练 → 还要继续训练：用 DCP，支持 reshard
• 训练 → 单机推理 (GPU 服务器)：safetensors + AOTI 编译
• 训练 → 移动端：ExecuTorch
• 训练 → 别的框架：ONNX（最后选择）
• 训练 → 社区分发：safetensors + HF Hub
• 训练 → CPU / Mac 端推理：转 GGUF

理解这个矩阵让你不会盲目"转 ONNX 部署"——很多场景 .pt2 / .pte 更优。每条转换链路都有 tradeoff，用最少的转换跳转达到目标平台是工程正解。

19.9.5 HuggingFace from_pretrained：业界实际加载流程

PyTorch 用户极少手写 torch.load，更常见的是：

from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained("meta-llama/Meta-Llama-3-8B")

这一行的内部流程：

graph TB
    Call[from_pretrained] --> DL[1. 从 Hub 下载文件<br/>config.json / model.safetensors / tokenizer]
    DL --> Cache[本地 cache: ~/.cache/huggingface/hub/]
    Cache --> Cfg[2. 解析 config.json<br/>找到 model class + dtype + ...]
    Cfg --> Build[3. 实例化空 model<br/>只 build 网络结构, 权重 init 为 meta tensor]
    Build --> Sharded{4. 单文件还是 sharded?}
    Sharded -->|单文件| Single[加载 model.safetensors]
    Sharded -->|多文件| Multi[加载 model-00001-of-00007.safetensors<br/>+ model.safetensors.index.json]
    Single --> Load[5. load_state_dict to model<br/>用 mmap, 一次填一个 tensor]
    Multi --> Load
    Load --> Done[ready model]
    style Build fill:#fef3c7
    style Load fill:#dcfce7

几个工程精巧点：

• meta tensor build：第 3 步用 torch.device('meta') 创建模型骨架，无显存分配。第 5 步把 mmap 字节直接 cast 到 GPU，跳过 "CPU full copy → GPU copy" 双重内存压力
• sharded 文件命名：model-00001-of-00007.safetensors 是事实标准，每片大小由保存时的分片上限决定（与 safetensors 格式本身无关，见 §19.9.1）
• index.json：含 {"weight_map": {"model.layers.0.weight": "model-00001-of-00007.safetensors"}} 的反查表 —— 加载某个权重时知道去哪个文件读

理解这套流程让你看 vLLM / SGLang / 自家推理服务的"模型加载几分钟"开销主要在哪 —— 一是网络下载、二是 GPU memory 分配、三是 cuda 上下文 init。真正的 deserialization 反而最快（mmap + zero-copy）。

19.10 加载、校验与运维

19.10.1 跨设备加载：map_location 与 device_map

ckpt 保存时 tensor 在 GPU，加载时机器没那么多 GPU 怎么办？map_location 解决这个：

# 把所有 tensor 加载到 CPU
state = torch.load('ckpt.pt', map_location='cpu')

# 把所有 tensor 加载到 cuda:0 (即便保存时在 cuda:7)
state = torch.load('ckpt.pt', map_location='cuda:0')

# 自定义函数: 参数是 storage 和它保存时的 location 字符串 (如 "cuda:7")
def loader(storage, loc):
    return storage.cuda(0)
state = torch.load('ckpt.pt', map_location=loader)

底层：torch.load 在恢复 storage 时调 map_location 决定目标 device。

device_map 是 HuggingFace transformers 的高级版本：

model = AutoModelForCausalLM.from_pretrained(
    "Llama-70B",
    device_map="auto",      # 自动分配 layer 到多 GPU
)

device_map="auto" 算法：

1. 估算每 layer 占多少 memory
2. 按层数把 layer 平均分到所有 GPU
3. 加载 weight 时直接 map 到目标 GPU（不经过 CPU）

实战：8 卡机器加载 Llama-70B（140 GB），device_map="auto" 让每卡 ~17.5 GB。配合 offload_folder 还能让多余 layer 落到磁盘 → 单卡推理 7B、双卡 70B、单卡 + offload 70B 都能跑。

理解这层让你看到"内存不够也能 load 模型"不是魔法，是 PyTorch + HF 联合做的精细 device 管理。

19.10.2 ckpt 完整性校验

长时间训练 ckpt 损坏会让"resume 后 loss 直接 NaN"。生产建议加 checksum：

import hashlib

def save_with_checksum(state_dict, path):
    torch.save(state_dict, path)
    with open(path, 'rb') as f:
        sha256 = hashlib.sha256(f.read()).hexdigest()
    with open(path + '.sha256', 'w') as f:
        f.write(sha256)

def load_with_checksum(path):
    with open(path, 'rb') as f:
        actual = hashlib.sha256(f.read()).hexdigest()
    with open(path + '.sha256') as f:
        expected = f.read().strip()
    if actual != expected:
        raise ValueError(f"checksum mismatch: {path}")
    return torch.load(path)

顺带澄清：DCP 自身并不做内容校验。.metadata 里的 Metadata dataclass（torch/distributed/checkpoint/metadata.py:138）只有 state_dict_metadata / planner_data / storage_data / storage_meta / version 五个字段，torch/distributed/checkpoint/ 下搜不到任何 sha256 / checksum 逻辑。zip 层面倒是有 CRC32（miniz 自带）：.distcp 本身不是 zip，但默认格式下其中每个张量 chunk 都是一段 torch.save 产物（torch/distributed/checkpoint/filesystem.py 的 _write_item），CRC32 只覆盖到这一层，BYTE_IO 项和 .metadata 都不在其内；而且 torch.serialization.set_crc32_options(False) 还能把它关掉（C++ 侧即 MZ_ZIP_FLAG_DO_NOT_COMPUTE_CRC32）。想要端到端校验，得自己在上层加。

corruption 来源：

• 存储介质位翻转：硬件老化 / 宇宙射线导致的 silent data corruption。发生率极低但不为零，ckpt 越大、保留越久，撞上的概率越高
• 网络传输错误：TCP checksum 不能保证 100%，云存储传输偶发 corruption
• 进程 crash 中断写入：写一半 SIGKILL、ckpt 不完整

每种都让 ckpt 看起来"加载正常但 resume 后 loss 异常"。checksum 让这种 silent failure 立刻暴露。生产大模型训练必须有 checksum，否则浪费几小时训练在坏 ckpt 上。

19.10.3 ckpt 内容 inspector

调试 ckpt 问题时常需"打开看里面有啥"，但又不想真加载几百 GB。常用工具：

# 1. 看 zip 结构 (不读 tensor 数据)
import zipfile
with zipfile.ZipFile('model.pt') as z:
    for name in z.namelist():
        print(name)

# 2. 看 metadata 不加载 tensor (map_location='meta': 不分配内存也不读字节)
state = torch.load('model.pt', map_location='meta')
print(state.keys())                   # 看 fqn 列表
print({k: v.shape for k, v in state.items()})

# 3. safetensors header 直接解析
from safetensors import safe_open
with safe_open('model.safetensors', framework='pt') as f:
    print(f.metadata())               # header 里的 __metadata__ 字段 (可能为 None)
    print(list(f.keys()))             # tensor 名列表

# 4. DCP 的 .metadata 直接看 (是 pickle 出来的 Metadata dataclass, 不是 JSON)
from torch.distributed.checkpoint import FileSystemReader
meta = FileSystemReader('ckpt_dir').read_metadata()
print(meta.state_dict_metadata.keys())   # 每个 fqn
print(meta.storage_data)                 # fqn/offset -> 物理文件与偏移

实战场景：

• 比较两个 ckpt 结构是否一致：第 2 种方式拿 fqn list + shape 对比
• 找某个 weight 在哪个 shard 文件：safetensors 的 weight_map (model.safetensors.index.json)
• 查 ckpt 训练步数：state['step'] 或 state['epoch']
• 诊断 reshard 失败：看 DCP .metadata 里 state_dict_metadata 的 chunk 坐标是否符合预期

把这些命令做成 inspect_ckpt.py 脚本，是大模型训练运维必备的工具。

19.10.4 ckpt 版本管理：HuggingFace Hub 的 git-like 模型

大模型项目都需要"版本管理"——v1.0 的 ckpt、v1.1 的 ckpt、experimental 分支的 ckpt 共存。HuggingFace Hub 用 git LFS 提供 git-like 版本管理：

# Hub 上一个 model repo 实际是 git repo
git clone https://huggingface.co/meta-llama/Meta-Llama-3-8B
ls -la
# .git/                    <- git 历史
# config.json              <- 元信息
# model.safetensors        <- 权重 (LFS-tracked, 大文件不在 git history)
# tokenizer.json
# README.md

每次模型更新 = git commit + LFS push。用户可以：

• git checkout v1.0：拿历史版本
• git diff main..v2.0 config.json：看不同版本配置变化
• git log --oneline：看模型更新历史

对应到 PyTorch 序列化层面：ckpt 不再是单文件，而是一组协同的文件（safetensors + config + tokenizer + chat template + ...）。HF 把这组文件放进同一个 model repo（说明文档即 README.md 形式的 model card），让模型分发等价于 git repo 分发。

实战影响：

• 不要把 from_pretrained("path/to/local") 当成"加载一个文件"，它实际加载一个目录 + 解析 config + ...
• ckpt 兼容性：升级 transformers 库时旧 ckpt 可能加载失败（因为 config schema 变了）
• 自家训练的 ckpt 上传 Hub 的工程：用 model.push_to_hub() 自动打包 + LFS push

理解这套生态让你看到"序列化"在 LLM 时代不只是"save tensor"，是"模型生命周期管理"的入口。PyTorch 自身只覆盖 tensor 字节，HF 生态补足版本管理 + 元信息 + 分发协议。

19.10.5 推理服务的 ckpt 加载优化

大模型推理引擎的"启动慢"问题常被诟病 —— 加载一个几十 B 的模型动辄几分钟。生产优化方向：

1. 模型预热池

服务集群保持 N 个 worker 进程，每个预先加载好 model。新请求来时直接路由到 worker，不重新 load。代价是常驻 RAM 高（每 worker 一份 weight）。

2. 共享内存加载

多个 worker 进程共享同一份 weight：master 进程 mmap 一次，子进程 fork 后共享 page。Linux CoW 让"读"操作不复制内存，N 个 worker 仅占 1 份 weight。

# master
weight_mmap = torch.load('model.pt', mmap=True)
# fork 子进程
for _ in range(num_workers):
    pid = os.fork()
    if pid == 0:
        # 子进程: weight_mmap 与 master 共享
        serve(weight_mmap)

CUDA tensor 复杂些 —— GPU 显存不能靠 fork 共享，要走 CUDA IPC（torch.multiprocessing 的 CUDA 张量共享就是这个机制）。做对之后，N 个 worker 共用同一份权重而不是各持一份，省下的显存与 worker 数成正比。

3. layered loading

启动时只加载 embedding + 第一层，后续 layer lazy load（首请求时 page-in）。服务能更早进入 ready 状态，代价是首请求变慢。适合"长尾请求 + 低 QPS"的场景。

4. tensor parallel pre-shard

ckpt 保存时按 TP 维度预先切好（如 8 个文件、每文件 1/8）。各 GPU 直接读自己负责的文件，跳过 master broadcast。生产部署的标准模式。

理解这套让你看到推理引擎的启动时间不是个恒定数字，是"工程取舍 × 硬件配置"的输出。每种优化都有 tradeoff：内存 vs 启动时间 vs 复杂度。

19.10.6 ckpt 加密：付费模型的工程做法

商业模型（如付费版本的微调模型）需要防止权重被直接复制使用。常见方案：

1. ckpt 整体加密

# 伪代码: derive_key / AES 只示意流程, 不是某个真实库的 API
key = derive_key(license)         # 用 license 算出对称密钥
encrypted = AES.encrypt(open('model.safetensors', 'rb').read(), key)
open('model.enc', 'wb').write(encrypted)

# 解密 + 加载
data = AES.decrypt(open('model.enc', 'rb').read(), key)
state = safetensors.torch.load(data)   # 从内存 bytes 反序列化

代价：解密时整个模型在内存里，运行时机器仍能 dump RAM 拿到明文。

2. weight 分层加密

只加密关键 layer（如 attention），公开 backbone。攻击者拿到部分 weight 不够推理，提高破解成本。但密码学上仍是"obfuscation 而非真加密"。

3. TEE / SGX 内运行

模型 weight 在 trusted execution environment 内解密 + 推理，OS 都看不到明文。代价：只在特定硬件可用 + 性能损失大。

4. 服务化 (model as service)

最常用的"加密"是不分发 weight、只分发 API。客户付费调用 API、不能拿权重。OpenAI、Anthropic 都用这种。

PyTorch 序列化层不直接支持加密（设计哲学是"开放"），上层应用自己实现。safetensors 的 metadata 字段可以存 license 信息，但不强制 enforce。

理解这层让你看到"模型保护"是个产品 / 商业问题，不是技术问题 —— PyTorch 把权重存成可读字节是正确选择，分发策略另说。

19.11 几条工程经验

1. 默认 weights_only=True：从 v2.6 起这是默认（PR #137602 翻转，逻辑在 torch/serialization.py:89 的 _default_to_weights_only）。如果遇到加载报错说"unsupported global"，先想想是否真的需要那个类，能换掉就换

2. 分发模型给社区用 safetensors：兼容性、安全性、加载速度全面优于 pickle

3. 大模型加载显式传 mmap=True：torch.load 的 mmap 默认仍是 None（torch/serialization.py:1316 起的签名），None 会回落到 torch.utils.serialization.config.load.mmap（默认 False），要自己开。前提是文件用新 zip 格式存的（_use_new_zipfile_serialization=True，默认）

4. 训练 ckpt 推荐 DCP：单 rank 写大文件已经过时

5. 只看元数据用 map_location="meta"：load 时只想查 ckpt 含哪些 key、什么 shape，传 map_location="meta" 既不读字节也不分配内存；torch.serialization.skip_data 上下文虽然也跳过读字节，但仍会按原大小分配 CPU storage（torch/serialization.py:2070），不适合几百 GB 的 ckpt

6. ckpt 校验：写完 ckpt 后立刻 load 一遍验证，避免训练几小时后才发现 ckpt 坏了。可以用单独 rank 异步验证

7. 跨 PyTorch 版本加载：state_dict 通常向后兼容。整对象 pickle 跨版本可能 deserialize 失败 —— 又一个用 state_dict 的理由

8. 写云存储分清两条路：torch.save / torch.load 不认 s3:// 这种 URI，要自己用 fsspec.open() 拿 file-like 对象再传进去；DCP 则内置 FsspecWriter / FsspecReader，dcp.save(..., checkpoint_id='s3://...') 可以直接跑

19.12 跨专栏关联

• 第 9 章 §9.6 nn.Module.state_dict：序列化的"上游"，本章是它的"下游"
• 第 18 章 FSDP：FSDP 分片训练的 ckpt 首选 DCP，免得先 gather 成全量再写
• 《vLLM 推理内核深度解析》专栏第 7 章 模型加载：vLLM 推理时大量用 safetensors + mmap 加载预训练模型

19.13 设计启示

序列化设计的核心思想：

第一：容器格式 + 数据分离：zip + tensor entries 让 mmap、并行加载、增量更新成为可能。比单一 blob 序列化灵活得多

第二：weights_only 把"加载"从代码执行降级到字节解析：这条原则在所有"加载第三方数据"场景都成立 —— JSON / Protobuf / safetensors 都比 pickle 安全。能不执行就不执行，是数据格式设计的底线

第三：逻辑坐标而非物理位置：DCP 用逻辑张量坐标解耦 "保存时拓扑" 与 "加载时拓扑"，让 ckpt 可重 reshard。这条思想在分布式数据库分片管理也常见，是"让数据脱离硬件假设"的工程标准做法

第四：state_dict 让模型架构与权重解耦：架构是 Python 代码（活的、可演进），权重是数据（死的、可移植）。这种"代码 + 数据"分离是任何"长寿模型"的工程基础

第五：异步、并行、增量是大规模训练的三件套：async save 让训练不阻塞、DCP 让每 rank 并行写、PEFT/LoRA 让分发只发增量。三个机制叠加才让千卡训练 ckpt 从"工程难题"变成"标准操作"。这条"在多个维度同时下功夫"的设计思路，对任何重型 IO 系统的设计都成立 —— 不要指望一招制胜，要在每条流水线上都挖性能

下一章拆量化与混合精度训练 —— FSDP 已经给了精度策略的初步配置，这章把整个 PyTorch 量化生态摊开，看 autocast 的 CastPolicy 决策、GradScaler 的动态范围管理，以及已迁往 torchao 的 PT2E 量化框架怎么把整个量化栈现代化。

19.13.1 ckpt 设计的演进时间线

PyTorch 序列化的几个关键节点：

下面几条都能从仓库 git 历史里查证：

版本	改进	依据
v1.6	_use_new_zipfile_serialization 成为默认，.pt 变成 zip 容器	PR #39460；torch/serialization.py:950/984 仍保留 =False 的老路径
v2.0	DCP 从 torch.distributed._shard.checkpoint 移到 torch.distributed.checkpoint	PR #88698
v2.1	fsspec storage reader / writer 进入 DCP	PR #98387（_fsspec_filesystem.py 首次出现）
v2.3	dcp.async_save 公开	PR #121325
v2.6	torch.load 的 weights_only 默认翻成 True	PR #137602
v2.9	torch.onnx.dynamo_export 移除，只剩 export(..., dynamo=...)；到 v2.11（本专栏基准）dynamo=True 为默认	PR #158130；torch/onnx/__init__.py:78 的默认值

dcp.async_save 至今仍带着 This feature is experimental and subject to change 的 warning —— 它可用，但别当成冻结的 API。

整体趋势：

• v1.x：从单文件 pickle 走向 zip 容器，让 mmap / 大文件可行
• v2.0-2.1：从单机走向分布式（DCP）+ 云存储后端
• v2.3 / v2.6：不阻塞训练（async_save）与从可用走向安全（weights_only）

ckpt 值得这么持续投入的理由很直接：ckpt 写入时间是直接从训练吞吐里扣的。ckpt 间隔越密、单步越快，写盘那段停顿占的比例就越大 —— 这也是 async save 存在的全部理由。