Checkpoint 使用指南
本文参考 torchtitan 上游 docs/checkpoint.md，说明torchtitan-npu 训练中常见的 checkpoint 保存、恢复和
Hugging Face 权重加载方式。
Checkpoint 格式
torchtitan-npu 主要涉及两类权重格式：
| 格式 | 典型文件 | 适用场景 | 是否包含优化器等训练状态 |
|---|---|---|---|
| DCP | .metadata、*.distcp |
训练断点恢复、不同并行切分间重分片加载 | 可以包含 |
| Hugging Face | model.safetensors、model.safetensors.index.json |
从开源权重初始化、导出给推理或其他生态使用 | 仅模型权重 |
DCP 是 PyTorch Distributed Checkpoint 格式。它可以保存分布式训练下的模型、
优化器、学习率调度器、dataloader 进度和 train_state，因此最适合断点续训。
保存目录通常形如：
outputs/<job_name>/checkpoint/step-1000/
其中 outputs/<job_name> 来自 --dump_folder，checkpoint 来自
checkpoint.folder，step-1000 是训练步数。
Hugging Face safetensors 是模型权重格式。它不保存优化器、学习率调度器和训练步数， 通常用于从预训练权重启动训练，或在训练结束后导出模型权重。
常用配置项
checkpoint 配置由 torchtitan.components.checkpoint.CheckpointManager.Config
提供，在模型的 config_registry.py 中设置，也可以通过 CLI 覆盖。
| 配置项 | 说明 |
|---|---|
enable |
是否启用 checkpoint。加载和保存都需要打开。 |
folder |
checkpoint 子目录名，最终路径为 {dump_folder}/{folder}。 |
interval |
每隔多少 step 保存一次 DCP checkpoint。 |
load_step |
指定加载某一步。为 -1 时自动加载最新 checkpoint。 |
initial_load_path |
当前输出目录没有 checkpoint 时，从该路径初始化。 |
initial_load_model_only |
从 initial_load_path 加载时是否只加载模型权重，默认 True。 |
initial_load_in_hf |
将 initial_load_path 或 hf_assets_path 当作 HF safetensors 权重加载。 |
last_save_model_only |
训练最后一步是否只保存模型权重，默认 True。 |
last_save_in_hf |
训练最后一步是否保存为 HF safetensors。 |
export_dtype |
last_save_model_only=True 时最终模型权重导出的 dtype。 |
load_only |
只加载不保存，适合验证或调试。 |
keep_latest_k |
只保留最近 k 个 checkpoint。设置为 0 表示全部保留。 |
exclude_from_loading |
从 DCP 加载时排除部分状态，例如 optimizer,lr_scheduler,dataloader。CLI 传值必须用英文逗号分隔，不能空格分隔。 |
async_mode |
DCP 异步保存方式，常用值为 disabled、async。 |
create_seed_checkpoint |
创建 seed checkpoint 模式，详见下文「创建 seed checkpoint」。 |
enable_first_step_checkpoint |
是否在 step 1 强制保存一次 checkpoint，常用于校验保存路径连通性。 |
initial_load_in_hf_quantized |
从 HF 量化权重加载，使用前需先开启 initial_load_in_hf。 |
enable_ft_dataloader_checkpoints |
TorchFT 场景下是否独立保存 per-replica dataloader 状态。 |
注意：如果
checkpoint.enable=False，即使配置了initial_load_path或initial_load_in_hf=True，也不会执行 checkpoint 加载逻辑。
CLI 使用约定
- 布尔字段反转：tyro 用
--checkpoint.no-<field>把 bool 置 False，不是--no-checkpoint.<field>也不是--checkpoint.<field> false。例如要把load_only关掉，写--checkpoint.no-load-only。 - 列表字段：
exclude_from_loading这类列表用英文逗号串成一个参数传入 （如--checkpoint.exclude_from_loading optimizer,lr_scheduler,dataloader）。 以空格分隔多个 token 会被 tyro 当成子命令报 Unrecognized options。 - 覆盖 registry 默认值：模型的
config_registry.py已经设过的字段（如 dsv32 4 层 debug config 把load_only=True、initial_load_path写死成某个路径）， CLI 必须显式覆盖才能生效；只追加--checkpoint.enable不会把这些字段清掉。 使用前可以先grep checkpoint=看一眼模型 registry 默认值。
保存 DCP checkpoint
CLI 启动
以 DeepSeek-V3.2 4 层配置为例，启用 DCP 保存：
NGPU=16 \
MODULE=torchtitan_npu.models.deepseek_v32 \
CONFIG=deepseek_v32_671b_4layers_debug \
bash scripts/run_train.sh \
--dump_folder ./outputs/dsv32_4layers \
--checkpoint.enable \
--checkpoint.folder checkpoint \
--checkpoint.interval 100 \
训练过程中会在 ./outputs/dsv32_4layers/checkpoint/step-100、
step-200 等目录保存 DCP checkpoint。中间 checkpoint 包含完整训练状态，
可用于断点续训。
在 config_registry.py 中设置
from torchtitan.components.checkpoint import CheckpointManager
checkpoint=CheckpointManager.Config(
enable=True,
folder="checkpoint",
interval=100,
keep_latest_k=5,
async_mode="disabled",
)
如果希望最后一步也保留完整训练状态，而不是只保存模型权重，将
last_save_model_only 设为 False：
checkpoint=CheckpointManager.Config(
enable=True,
interval=100,
last_save_model_only=False,
)
如果只关心最终模型权重，可以保留默认的 last_save_model_only=True，并指定导出精度：
checkpoint=CheckpointManager.Config(
enable=True,
interval=500,
last_save_model_only=True,
export_dtype="bfloat16",
)
读取 DCP checkpoint
自动续训
使用相同的 --dump_folder 和 --checkpoint.folder 重新启动训练时，框架会自动从
最新的 step-* DCP checkpoint 恢复：
NGPU=16 \
MODULE=torchtitan_npu.models.deepseek_v32 \
CONFIG=deepseek_v32_671b_4layers_debug \
bash scripts/run_train.sh \
--dump_folder ./outputs/dsv32_4layers \
--checkpoint.enable \
--checkpoint.folder checkpoint
如需指定某一步：
bash scripts/run_train.sh \
--dump_folder ./outputs/dsv32_4layers \
--checkpoint.enable \
--checkpoint.load_step 500
从另一个 DCP 目录初始化
当希望从旧任务的 checkpoint 启动一个新任务时，使用新的 dump_folder，并将
checkpoint.initial_load_path 指向旧 checkpoint 的 step 目录：
NGPU=16 \
MODULE=torchtitan_npu.models.deepseek_v32 \
CONFIG=deepseek_v32_671b_4layers_debug \
bash scripts/run_train.sh \
--dump_folder ./outputs/dsv32_new_job \
--checkpoint.enable \
--checkpoint.initial_load_path ./outputs/dsv32_4layers/checkpoint/step-500
默认只加载模型权重。如果要从这个 DCP 完整恢复优化器、学习率调度器和训练步数，
在 config_registry.py 中设置：
checkpoint=CheckpointManager.Config(
enable=True,
initial_load_path="./outputs/dsv32_4layers/checkpoint/step-500",
initial_load_model_only=False,
)
如果只想加载模型和部分状态，可以排除不需要的键：
bash scripts/run_train.sh \
--checkpoint.enable \
--checkpoint.initial_load_path ./outputs/dsv32_4layers/checkpoint/step-500 \
--checkpoint.exclude_from_loading optimizer,lr_scheduler,dataloader
可排除的常用键包括 optimizer、lr_scheduler、dataloader 和 train_state。
注意：如果
{dump_folder}/{checkpoint.folder}已经存在可用 checkpoint， 框架会优先从该目录续训，并忽略initial_load_path。从新权重启动实验时， 请使用新的dump_folder或清理旧的 checkpoint 目录。
读取 Hugging Face 权重
从 HF safetensors 权重初始化时，需要启用 checkpoint，并设置
initial_load_in_hf=True。路径应包含 model.safetensors 或
model.safetensors.index.json，同时建议包含 tokenizer、config.json 等 HF 资产。
NGPU=16 \
MODULE=torchtitan_npu.models.deepseek_v32 \
CONFIG=deepseek_v32_671b_61layers_4k_128die \
bash scripts/run_train.sh \
--checkpoint.enable \
--checkpoint.initial_load_in_hf \
--checkpoint.initial_load_path ./checkpoint/DeepSeek-V3.2
如果不显式传 checkpoint.initial_load_path，框架会尝试使用顶层
hf_assets_path：
bash scripts/run_train.sh \
--checkpoint.enable \
--checkpoint.initial_load_in_hf \
--hf_assets_path ./checkpoint/DeepSeek-V3.2
checkpoint.initial_load_path 的优先级高于 hf_assets_path。
HF 权重只能作为模型权重加载，不能恢复优化器或训练步数。
保存 Hugging Face 权重
如果希望训练最后一步直接导出 HF safetensors 权重，可以同时启用
last_save_in_hf 和 last_save_model_only：
NGPU=16 \
MODULE=torchtitan_npu.models.deepseek_v32 \
CONFIG=deepseek_v32_671b_61layers_4k_128die \
bash scripts/run_train.sh \
--dump_folder ./outputs/dsv32_hf_export \
--checkpoint.enable \
--checkpoint.interval 500 \
--checkpoint.last_save_in_hf \
--checkpoint.export_dtype bfloat16
最终权重会保存在类似下面的目录中：
./outputs/dsv32_hf_export/checkpoint/step-1000/
该目录用于模型权重导出，不适合作为完整训练断点恢复。保存 HF 权重时需要经过 state_dict 转换和 safetensors 合并，耗时通常高于 DCP 保存。
DeepSeek-V3.2 和 DeepSeek-V4 的 state dict adapter 还包含模型侧保存 patch，
可用于专家权重格式转换等高级场景。日常保存和加载优先使用
CheckpointManager.Config 的 DCP/HF 配置；只有模型配置明确需要时，再在
model.Config 中设置 save_patch_enabled、save_format、hf_save_dir 和
save_expert_format。
创建 seed checkpoint
seed checkpoint 用于先在单卡 CPU 初始化模型，再让多卡任务通过 DCP 重分片加载。
创建时需要关闭并行切分，并使用 NGPU=1：
NGPU=1 \
MODULE=torchtitan_npu.models.deepseek_v32 \
CONFIG=deepseek_v32_671b_4layers_debug \
bash scripts/run_train.sh \
--checkpoint.enable \
--checkpoint.create_seed_checkpoint \
--parallelism.data_parallel_replicate_degree 1 \
--parallelism.data_parallel_shard_degree 1 \
--parallelism.tensor_parallel_degree 1 \
--parallelism.pipeline_parallel_degree 1 \
--parallelism.context_parallel_degree 1 \
--parallelism.expert_parallel_degree 1
生成的 step-0 checkpoint 后续可作为 checkpoint.initial_load_path 使用。
格式转换
将 DCP checkpoint 合并为单个 PyTorch .pt 文件：
python -m torch.distributed.checkpoint.format_utils dcp_to_torch \
./outputs/dsv32_4layers/checkpoint/step-1000 \
checkpoint.pt
HF safetensors 到 DCP 的离线转换依赖模型的 state dict adapter 和对应模型入口。
当前仓库中的 scripts/checkpoint_conversion/convert_from_hf.py 仍保留旧版
TrainSpec 调用方式，对已经迁移到 config_registry.py 的模型不一定可直接使用。
对于 DeepSeek-V3.2、DeepSeek-V4、Qwen3 等当前模型，优先使用
checkpoint.initial_load_in_hf=True 在训练启动时直接读取 HF 权重。
常见问题
为什么配置了 initial_load_path 但没有加载？
通常有两个原因：一是没有设置 checkpoint.enable=True；二是当前
{dump_folder}/{checkpoint.folder} 已经存在可用 checkpoint，框架优先从当前
输出目录续训。
为什么最后一步 checkpoint 不能完整续训？
last_save_model_only=True 时，最后一步只保存模型权重，不包含优化器、
学习率调度器和训练状态。需要完整续训时，请使用中间 DCP checkpoint，或在配置中将
last_save_model_only=False。
DCP 和 HF 权重应该怎么选？
断点续训、故障恢复和分布式重分片优先用 DCP；从开源模型初始化或导出给推理生态时用 HF safetensors。
SwapOptimizer 下中间 DCP 保存或续训会报 NotImplementedError？
deepseek_v32、deepseek_v4 等模型默认启用 optimizer.swap_optimizer=True，其 state 中含
None 项，目前不能被 torch.distributed.checkpoint 的优化器扁平化逻辑序列化，典型报错：
NotImplementedError: Flattening optimizer state_dict only supports tensor, int, float, dict states now. Type is <class 'NoneType'>.
触发位置有两处：
interval触发的中间 step 保存（含完整优化器状态）。- 任何
dcp.load：因为 plan 阶段也会调用优化器的state_dict()。
短期 workaround：
- 只用
last_save_model_only=True（默认就是 True）做模型权重落盘，保存时不要触发中间 step（拉大interval或减少training.steps）。 - Resume /
initial_load_path加载时显式排除优化器和训练状态：--checkpoint.exclude_from_loading optimizer,lr_scheduler,dataloader,train_state。
这是当前框架与 SwapOptimizer 的兼容性限制，待 SwapOptimizer 把 None state 改成空张量 或上游 DCP 支持 NoneType 后可恢复完整状态保存。