Chapter 10
Checkpoint：dcp 与 HF 转换
TorchTitan 默认用 PyTorch 的 distributed checkpoint（dcp）格式，每张卡保自己那份；要给 HF 生态用还得转格式。本章把两条路径走通。
10.1dcp 格式
dcp（torch.distributed.checkpoint）是 PyTorch 原生的"每 rank 一份"格式：
./checkpoint/step-1000/
├── .metadata # 张量元信息（shape / dtype / device_mesh）
├── __0_0.distcp # rank 0 的数据
├── __1_0.distcp # rank 1
├── __2_0.distcp # ...
└── ...优势：
- save / load 并行，时间不随 GPU 数线性增加；
- 支持不同 world_size 加载（reshard）；
- 跟 DTensor / FSDP2 / TP 完美兼容。
10.2save / load
TorchTitan 封装在 torchtitan/components/checkpoint.py::CheckpointManager：
ckpt_mgr = CheckpointManager(model, optimizer, scheduler, ...)
ckpt_mgr.save(step=1000)
ckpt_mgr.load() # 从 latest 加载内部调 torch.distributed.checkpoint.async_save() 或 save()。
10.3async checkpoint
save 是 I/O 密集，会阻塞训练。TorchTitan 支持 async：
[checkpoint]
async_mode = "async" # disabled / async / async_with_pinned_mem| mode | 说明 |
|---|---|
| disabled | 同步保存，训练阻塞 |
| async | fork 后台进程保存 |
| async_with_pinned_mem | + pinned memory 加速 |
async 后训练不阻塞，但要预留 CPU 内存放 staging tensor。
10.4dcp → 单文件（推理用）
dcp 不能直接给 HF 加载。要先合并：
from torch.distributed.checkpoint.format_utils import dcp_to_torch_save
dcp_to_torch_save(
dcp_checkpoint_dir="./checkpoint/step-1000",
torch_save_path="./model.pth",
)产物 ./model.pth 是单 torch.save 文件。
10.5dcp → HF safetensors
TorchTitan 提供脚本 scripts/convert_to_hf.py（具体路径以仓库为准）：
python scripts/convert_to_hf.py \
--ckpt_path ./checkpoint/step-1000 \
--output_path ./hf_out \
--tokenizer_path ./tokenizer/tokenizer.model \
--model_name llama3 \
--flavor 8B核心步骤：
- 用 dcp API 把分布式 ckpt 加载到单进程；
- 按 HF 命名约定重命名 state_dict 的 key；
- 保存为 safetensors 分片。
10.6HF → TorchTitan
反向也类似——从 HF safetensors 加载到 model，按 TorchTitan 期望的 name 命名，然后 dcp save。代码在 scripts/convert_from_hf.py。
10.7resume 包含什么
| 子项 | 是否默认保 |
|---|---|
| model state_dict | ✅ |
| optimizer state | ✅ |
| lr_scheduler state | ✅ |
| step / epoch | ✅ |
| RNG state | ✅ |
| data_loader state (iterator 位置) | ✅ |
所以 resume 是完全确定的，loss 曲线接得上。
10.8跨 world_size resume
dcp 支持变 world_size加载：训时 16 卡，resume 时换 8 卡，dcp 自动 reshard：
# 训时
torchrun --nproc_per_node=8 --nnodes=2 ...
# resume 时变 1 卡
torchrun --nproc_per_node=1 ... # 自动 reshard10.9这章你需要带走的
- dcp = 每张卡一份 + metadata 总文件；
- save / load 并行，支持 async；
- HF 互转用
convert_to_hf.py/convert_from_hf.py； - resume 默认包含 model + optim + scheduler + RNG + data_loader 位置；
- dcp 支持 reshard，跨 world_size 加载。