PiD Checkpoint管理实战:DCP分布式格式详解与.pth单文件转换技巧
【免费下载链接】PiD[NeurIPS 2026] PiD: Fast and High-Resolution Latent Decoding with Pixel Diffusion项目地址https://gitcode.com/gh_mirrors/pid18/PiD点击查看免费下载PiD 训练时默认采用 PyTorch Distributed CheckpointDCP分片格式保存权重但推理和新阶段初始化却需要单文件.pth。本文带你完整搞懂PiD Checkpoint 管理DCP 目录结构长什么样、Teacher 与 Distillation 两种训练的差异以及如何用官方脚本把 DCP 一步转换为可直接torch.load()的.pth文件并附常见踩坑清单。一、为什么 PiD 训练要用 DCP 格式PiD 的分布式训练由两套 Checkpointer 负责通过 Hydra 的ckpt_type选项切换训练类型ckpt_typeCheckpointer源码位置Teacher 训练dcpDistributedCheckpointerpid/_ext/imaginaire/checkpointer/dcp.py蒸馏训练dcp_distillDistillationCheckpointerpid/_src/checkpointer/dcp_distill.pyDCP 的核心特点是目录式分片一个逻辑 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/两个关键差异要记牢目录名取决于配置optim_key/、scheduler_key/的后缀来自模型optimizer_dict的 key关闭 GAN 后 discriminator 相关目录就不存在不要硬编码。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后一次生成三个文件输出文件内容model.pth完整 state dict含net.*、net_ema.*等model_ema_fp32.pth仅 EMA 权重key 由net_ema.*重命名为net.*model_ema_bf16.pth同上float32 tensor 转为 bfloat16mkdir -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 EMAfake_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只能用于推理或新阶段权重初始化。五、不同场景该用哪个文件场景推荐文件说明同一 job 中断后续训完整 DCP 目录 latest_checkpoint.txt严格恢复 model/optim/scheduler/trainer从外部 DCP 显式恢复外部iter_xxxxxxxxx/checkpoint.load_path指向它load_training_stateFalse时仅加载model/推理 / 新阶段初始化转换后的model_ema_*.pth从 iteration 0 开始只含权重浏览/迁移权重转换后的model.pth可直接torch.load()查看六、延伸阅读完整官方文档含 metadata 字段级细节docs/checkpointing_CN.md发布模型的路径索引docs/checkpoints.mdHydra 配置项注册pid/_src/configs/common/defaults/ckpt_type.py掌握以上要点你就能在 PiD 的分布式训练、续训和推理部署之间自如切换 Checkpoint 形态不再被.distcp分片卡住。赞分享【免费下载链接】PiD[NeurIPS 2026] PiD: Fast and High-Resolution Latent Decoding with Pixel Diffusion项目地址https://gitcode.com/gh_mirrors/pid18/PiD点击查看免费下载相关推荐洛雪音乐桌面版免费开源跨平台音乐聚合播放器的终极指南洛雪音乐桌面版免费开源跨平台音乐聚合播放器的终极指南 你是否厌倦了在各个音乐平台之间来回切换只为找到一首想听的歌曲是否对各大平台的会员制度和曲库限制感到困桌面应用音视频前端LLaMA-Adapter模型转换工具checkpoint提取与格式转换详解LLaMA Adapter模型转换工具checkpoint提取与格式转换详解 LLaMA Adapter是一个革命性的轻量级适配器方法能够在短短1小时内使用人工智能大模型微调多模态PEFT Checkpoint 格式完全指南文件结构、格式转换与模型存储PEFT Checkpoint 格式完全指南文件结构、格式转换与模型存储 本指南以 PEFT 官方开发者文档《PEFT checkpoint format》为人工智能大模型微调LoRA创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考