PiD Checkpoint管理实战：DCP分布式格式详解与.pth单文件转换技巧

PiD 训练时默认采用 PyTorch Distributed Checkpoint（DCP）分片格式保存权重，但推理和新阶段初始化却需要单文件.pth。本文带你完整搞懂PiD Checkpoint 管理：DCP 目录结构长什么样、Teacher 与 Distillation 两种训练的差异，以及如何用官方脚本把 DCP 一步转换为可直接torch.load()的.pth文件，并附常见踩坑清单。

一、为什么 PiD 训练要用 DCP 格式？

PiD 的分布式训练由两套 Checkpointer 负责，通过 Hydra 的ckpt_type选项切换：

DCP 的核心特点是目录式分片：一个逻辑 state dict 不再是一个.pth文件，而是一个目录：

model/ ├── .metadata # 记录每个 key 的 shape/dtype/chunk 及在 shard 中的位置 ├── __0_0.distcp # rank 0 的 tensor payload ├── __1_0.distcp # rank 1 的 tensor payload └── ...

这样设计带来三大好处：

• ⚡多进程并行 I/O，保存/加载吞吐更高；
• 🧠降低单进程内存峰值，各 rank 只写自己的 shard；
• 🔀按当前分布式拓扑按需读取，续训时灵活匹配。

⚠️ 注意：DDP 下各 rank 完全复制的状态可能只写入最低 rank，其他 rank 的.distcp可能是 0 字节——这是正常现象，不代表 checkpoint 损坏。

二、两种训练的真实 Checkpoint 目录结构

每次保存以 9 位 iteration 编号建目录，根目录下的latest_checkpoint.txt记录最新保存的目录名（不是软链接），用于同一 job 自动续训：

checkpoints/ ├── latest_checkpoint.txt └── iter_000000025/ ├── model/ # 模型权重（net.* + net_ema.*） ├── optim/ # optimizer 状态 ├── scheduler/ # 学习率调度器状态 └── trainer/ # iteration、GradScaler 状态

Teacher 训练每个逻辑目录就是一个独立 DCP；而蒸馏训练会把 optimizer/scheduler 按组件拆成独立目录，因为 PiD 蒸馏同时训练 student、fake-score 和 discriminator 三套参数：

iter_000000025/ ├── model/ # 含 net.* / net_ema.* / fake_score.* / discriminator.* ├── optim_net/ # student optimizer ├── optim_fake_score/ # fake-score optimizer ├── optim_discriminator/ # discriminator optimizer ├── scheduler_net/ ├── scheduler_fake_score/ ├── scheduler_discriminator/ └── trainer/

两个关键差异要记牢：

1. 目录名取决于配置：optim_<key>/、scheduler_<key>/的后缀来自模型optimizer_dict的 key，关闭 GAN 后 discriminator 相关目录就不存在，不要硬编码。
2. Frozen teacher 不在蒸馏 DCP 里：teacher 由配置中的pretrained_teacher_path单独加载，迁移蒸馏环境时必须同时保证 teacher checkpoint 可访问。

三、DCP 转 .pth：官方脚本一步到位

DCP 目录无法直接torch.load()，PiD 提供了转换脚本 scripts/convert_distcp_to_pt.py，内部封装了torch.distributed.checkpoint.format_utils.dcp_to_torch_save，把分片合并为单文件。

基本用法：只合并权重

输入必须精确指向某个 iteration 下的model/目录（不要传checkpoints/根目录，也不接受s3://等对象存储地址）：

mkdir -p /path/to/output_dir PYTHONPATH=. python scripts/convert_distcp_to_pt.py \ /path/to/checkpoints/iter_000000025/model \ /path/to/output_dir

默认输出model.pth——即model/中完整 state dict 的合并结果（不含 optimizer/scheduler/trainer）。

进阶用法：导出 EMA 推理权重

推理场景通常只想用 EMA 权重。加--ema参数（见 scripts/convert_distcp_to_pt.py）后一次生成三个文件：

mkdir -p /path/to/output_dir PYTHONPATH=. python scripts/convert_distcp_to_pt.py \ /path/to/checkpoints/iter_000000025/model \ /path/to/output_dir \ --ema

对蒸馏 checkpoint，--ema导出的是student EMA；fake_score.*和discriminator.*不会进入 EMA-only 文件。

四、避坑清单：这些细节能省你半天

• 📥用普通python单进程运行，别用torchrun：转换会把整个逻辑 state dict 合并进 CPU 内存，样例规模下峰值内存可能明显高于文件大小（teacher 约 11 GB、蒸馏约 16 GB），先确认 RAM 和磁盘余量。
• 🗑️脚本会先删除输出目录中的旧产物：启动时即unlink已有的model.pth/model_ema_fp32.pth/model_ema_bf16.pth，即使后续因输入错误或 OOM 失败，旧文件也没了——转换前务必确认输出路径。
• 📁脚本不会创建输出目录，必须先mkdir -p。
• 🔁checkpoint 没有net_ema.*时：脚本先写出model.pth，随后在 EMA 提取阶段抛ValueError。
• 🏷️扩展名是.pth不是.pt：脚本输出固定为.pth，别按.pt找文件。
• 🧩--keep-original优先于--ema：传了它会在写出model.pth后直接退出，不做 EMA 提取。
• ⛔.pth不能替代 DCP 无损续训：转换结果只含模型权重，不含 optimizer、scheduler 和 iteration，只能用于推理或新阶段权重初始化。

五、不同场景该用哪个文件？

六、延伸阅读

• 完整官方文档（含 metadata 字段级细节）：docs/checkpointing_CN.md
• 发布模型的路径索引：docs/checkpoints.md
• Hydra 配置项注册：pid/_src/configs/common/defaults/ckpt_type.py

掌握以上要点，你就能在 PiD 的分布式训练、续训和推理部署之间自如切换 Checkpoint 形态，不再被.distcp分片"卡住"。