知识花园

模型权重保存

14分钟阅读

深度学习模型的 Checkpoint（ckpt，检查点）分为训练态格式和推理态格式两大类：

  • 训练态：为断点续训设计，除模型权重外，还保存优化器状态、学习率调度器、随机种子、并行分片信息等，通常与训练框架、并行策略强绑定。
  • 推理态：仅保留模型权重与结构配置，面向部署、分发和推理，格式通用度高。

下面按主流框架和场景，逐一介绍常见格式的保存机制、文件结构与加载注意事项。

一、PyTorch 原生基础格式（.pt / .pth）

这是所有PyTorch生态的基础格式，本质是通过torch.save()基于Pickle序列化的Python对象。

1. 两种保存模式

  • 完整模型保存（不推荐）：直接序列化整个nn.Module对象，包含模型结构+权重。 缺点：与模型代码强耦合，跨版本、跨项目兼容性极差，仅适合简单小模型临时保存。
  • state_dict 保存（行业标准）：仅保存权重张量字典（model.state_dict()），可额外搭配优化器、调度器状态。 典型保存内容： 
    {
    "model": model.state_dict(),
    "optimizer": optimizer.state_dict(),
    "scheduler": scheduler.state_dict(),
    "epoch": epoch,
    "global_step": step
}

2. 加载注意事项

  1. 结构强匹配：加载时模型类的结构、层名称必须和保存时完全一致，key不匹配会直接报错。
  2. 设备映射：GPU保存的权重加载到CPU/其他GPU时，必须指定map_location参数，否则会默认加载到原设备导致显存不足。
  3. 安全风险：Pickle格式支持执行任意Python代码，来路不明的.pt文件严禁加载，存在后门风险。
  4. 版本兼容：高版本PyTorch保存的ckpt可能无法在低版本加载，建议保存和加载的主版本一致。
  5. 适用场景：单卡训练、简单数据并行（DP/DDP）、小模型推理。

二、分布式大模型训练专用格式

大模型训练普遍采用张量并行、流水线并行、分片优化器等策略，权重无法存储在单文件中，因此形成了各框架专属的分片ckpt格式。

2.1 Megatron-LM 系列检查点

Megatron-LM、Megatron-DeepSpeed、NeMo-Megatron均采用此设计，也是你当前使用的格式。

核心机制

基于张量并行(MP) + 流水线并行(PP) + 数据并行(DP) 三级并行，每个并行rank独立保存自己的分片权重与状态，不做全局合并。

标准目录结构

iter_0068000/                # 按迭代步数命名的检查点目录
├── latest_checkpointed_iteration.txt  # 最新迭代标记文件
├── mp_rank_00/              # 张量并行第0号分片
│   ├── pp_rank_00/          # 流水线并行第0阶段
│   │   ├── model_optim_rng.pt  # 分片权重+优化器超参+RNG随机状态
│   │   └── distrib_optim.pt    # 分布式优化器状态（Adam动量/方差）
│   └── pp_rank_01/          # 流水线并行第1阶段
│       └── ...
└── mp_rank_01/              # 张量并行第1号分片
    └── ...

你当前目录没有pp_rank层级，说明流水线并行度 PP=1；mp_rank有2个，说明张量并行度 MP=2。

关键文件说明

  • model_optim_rng.pt：三合一文件，包含分片模型权重、优化器超参数配置、训练随机数(RNG)状态，保证续训时dropout、数据打乱等随机性完全对齐。
  • distrib_optim.pt：分布式优化器的状态分片，存储Adam的一阶动量m、二阶动量v，是断点续训的核心。

加载注意事项

  1. 并行度必须严格匹配：加载时的MP、PP、DP数量必须和保存时完全一致，否则会出现维度不匹配、权重缺失。
  2. 分片完整性：所有mp_rank、pp_rank目录必须完整，缺少任意一个分片都无法正常加载。
  3. 不能直接通用：单个mp_rank下的.pt文件只是权重切片，不能直接当作完整模型加载，必须通过Megatron框架初始化后加载。
  4. 推理导出：用于推理时，必须使用官方convert_checkpoint工具，合并所有MP/PP分片，导出为HuggingFace格式或单文件权重。
  5. 分支差异：原生Megatron-LM、Megatron-DeepSpeed、NeMo的ckpt结构有细微差异，不能跨分支直接加载。

2.2 DeepSpeed ZeRO 检查点

DeepSpeed ZeRO是大模型训练最常用的优化器分片方案，分ZeRO-1/2/3三个阶段，保存逻辑差异显著。

核心机制

  • ZeRO-1/2：模型权重每个GPU保留完整副本，仅优化器状态分片；
  • ZeRO-3：模型权重、梯度、优化器状态全部分片，每个GPU仅持有一部分。

标准目录结构

checkpoint/
├── latest                         # 最新step软链接
└── global_step_1000/              # 迭代步目录
    ├── zero_pp_rank_0_model_states.pt   # 第0号DP分片的模型权重
    ├── zero_pp_rank_1_model_states.pt   # 第1号DP分片的模型权重
    ├── zero_pp_rank_0_optim_states.pt   # 第0号DP分片的优化器状态
    ├── zero_pp_rank_1_optim_states.pt
    └── lr_scheduler.pt            # 学习率调度器状态（全局一份）

如果同时开启张量并行，会增加mp_rank_xx层级。

两种保存模式

  1. 分片模式（默认）：每个rank独立保存自己的分片，保存速度快，无额外显存开销，用于断点续训。
  2. Universal Checkpoint（统一格式）：导出合并后的全量权重，跨框架兼容，适合模型迁移。

加载注意事项

  1. ZeRO阶段匹配：ZeRO-3的分片ckpt不能在ZeRO-2配置下加载，阶段必须对齐。
  2. 续训要求：断点续训必须保留所有optim_states.pt分片，缺少任意一个都会导致优化器状态不完整。
  3. 权重提取：仅需要模型权重推理时，使用DeepSpeed官方脚本zero_to_fp32.py合并分片，导出单文件全量权重。
  4. 禁止直接加载：单个分片文件不能用torch.load直接当作完整模型使用，必须通过DeepSpeed引擎加载。

2.3 PyTorch FSDP 检查点

PyTorch官方的完全分片数据并行（Fully Sharded Data Parallel），是原生分布式训练的主流方案。

三种保存模式（核心区别）

模式说明适用场景
FULL_STATE_DICT所有rank通信合并，保存完整的单文件权重导出推理、跨框架迁移
LOCAL_STATE_DICT每个rank仅保存自身分片，无通信开销同配置断点续训
SHARDED_STATE_DICT基于PyTorch DCP（Distributed Checkpoint）的分布式保存格式，支持异步保存、断点续传超大规模模型训练

文件结构

  • FULL模式：单文件.pt，和PyTorch原生格式完全兼容；
  • SHARDED（DCP）模式：目录形式，包含大量.ckpt分片文件、元数据文件，不能直接用torch.load读取。

加载注意事项

  1. LOCAL_STATE_DICT仅支持同并行度的FSDP续训，并行数修改后无法加载。
  2. DCP格式必须通过torch.distributed.checkpoint.load加载，不支持原生torch.load。
  3. FSDP与DeepSpeed、Megatron的ckpt互不通用，需要通过全量权重中转转换。

三、推理与分发类格式

这类格式仅保留推理必需的权重与结构信息，通用性强，适合模型分享、部署。

3.1 HuggingFace Transformers 标准格式

目前开源模型的事实标准，几乎所有训练框架都支持导出此格式。

目录结构

model_dir/
├── config.json          # 模型结构配置（层数、隐层维度、头数等）
├── pytorch_model.bin    # 模型权重（小模型单文件）
├── pytorch_model-00001-of-00003.bin  # 大模型分片权重
├── pytorch_model.bin.index.json      # 分片映射表
├── tokenizer.json       # 分词器配置
└── special_tokens_map.json

加载注意事项

  1. 配置文件必需：config.json是加载的核心，缺少则无法初始化模型结构。
  2. 分片完整性：大模型分片必须完整，且index.json不能损坏，否则会加载失败。
  3. 自动设备映射：大模型加载建议使用device_map="auto"，自动拆分权重到多卡/CPU。
  4. 版本兼容：不同Transformers版本可能存在层命名差异，升级后加载可能需要调整key映射。

3.2 Safetensors 格式（.safetensors）

HuggingFace推出的新一代张量格式，专门替代Pickle格式的.pt/.bin。

核心优势

  1. 安全：不执行任意代码，彻底杜绝Pickle的后门风险，适合公开分发的模型。
  2. 高速：零拷贝加载，速度比torch.load快2~10倍。
  3. 跨语言：支持Python、C++、Rust等多语言，跨框架兼容。

加载注意事项

  1. 需要安装safetensors依赖库，HF的from_pretrained可直接加载。
  2. 仅能存储张量数据，不能保存Python对象，因此无法保存优化器、RNG状态，仅用于模型权重。
  3. 和HF格式完全兼容，可直接替换.bin文件使用。

3.3 GGUF / GGML 格式

面向端侧、CPU轻量推理的格式，是llama.cpp生态的标准，主打量化压缩与低资源推理。

特点

  • 单文件包含模型结构、权重、量化信息，无需额外配置文件；
  • 支持多种量化等级（Q4_K_M、Q8_0、FP16等），体积远小于原生格式。

加载注意事项

  1. 只能通过llama.cpp或支持GGUF的推理框架加载，PyTorch无法直接读取。
  2. 不同量化等级的文件不通用，不能互相转换精度。
  3. GGML是旧版格式，已被GGUF全面取代，新模型均使用GGUF。

3.4 ONNX / TensorRT 推理引擎格式

  • ONNX：开放神经网络交换格式，跨框架、跨硬件通用，适合部署到不同推理后端。
  • TensorRT：NVIDIA专属推理引擎格式，针对GPU硬件做深度优化，推理速度最快。

注意事项

  1. 均为静态图格式，导出后无法修改模型结构，仅用于推理，不能续训。
  2. TensorRT格式与显卡架构、CUDA版本强绑定，不同硬件、不同驱动版本不能通用。

四、通用避坑与最佳实践

  1. 续训前核对并行配置：所有分布式分片ckpt，续训时的并行度（MP/PP/DP/FSDP分片数）必须和保存时完全一致，这是最常见的加载失败原因。
  2. 区分训练与推理用途：
    • 仅推理：只保留模型权重，删除优化器、RNG状态文件，可节省70%以上存储空间。
    • 需续训：完整保留所有分片、状态文件，严禁删除任何rank目录。
  3. 优先使用安全格式：公开分发、第三方模型优先选择safetensors，避免加载不明来源的.pt文件。
  4. 格式转换用官方工具：不同格式之间不要手动修改文件名或拼接权重，必须使用对应框架的官方转换脚本，避免权重错乱。
  5. 完整性校验：大模型ckpt传输、下载后，建议校验MD5/SHA，避免文件损坏导致加载异常。

需要我补充一份不同格式之间互转的常用工具和命令清单吗？

关系图谱