_checkpoint-page:

使用 Checkpoints 支持容错训练

最后更新: 06/25/2025.

在整个 RLHF 训练过程中可能会出现训练错误或机器故障，因此建议启用 checkpoints 以最大程度地减少损失。

API 接口已在配置说明中列出，此处不再重复。但仍有一些技术细节希望在此澄清。

Note

请注意，除了 hf_model 之外，checkpoint.contents 字段对 FSDP checkpoint 没有影响，其他三个字段会绑定在一起进行保存和加载。我们建议包含 model、optimizer 和 extra。

Checkpoint 保存目录结构

通常，我们在保存 checkpoints 时使用 ppo_trainer.yaml 或 ppo_megatron_trainer.yml 中声明的 default_local_dir 作为前缀，其格式为 checkpoints/${trainer.project_name}/${trainer.experiment_name}。

因此，FSDP 的内部 checkpoint 结构如下：

checkpoints/${trainer.project_name}/${trainer.experiment_name}
├── global_steps_${i}
│   ├── actor
│   │   ├── huggingface      # 默认保存配置和 tokenizer，如果 checkpoint.contents 中包含 ``hf_model``，则保存 huggingface 模型
│   │   └── fsdp_config.json # FSDP 配置文件，包含 world_size 和 fsdp 版本
│   │   ├── model_world_size_{self.world_size}_rank_{self.rank}.pt
│   │   ├── optim_world_size_{self.world_size}_rank_{self.rank}.pt
│   │   └── extra_state_world_size_{self.world_size}_rank_{self.rank}.pt
│   ├── critic
│   │   ├── huggingface
│   │   └── fsdp_config.json
│   │   ├── model_world_size_{self.world_size}_rank_{self.rank}.pt
│   │   ├── optim_world_size_{self.world_size}_rank_{self.rank}.pt
│   │   └── extra_state_world_size_{self.world_size}_rank_{self.rank}.pt
└── latest_checkpointed_iteration.txt

所有模型分片、优化器和额外状态会以分片（sharded）和分布式（distributed）的方式一起存储。

而 Megatron 当前的 checkpoint 结构如下：

checkpoints/${trainer.project_name}/${trainer.experiment_name}
├── global_steps_${i}
│   ├── actor
│   │   ├── huggingface     # 默认保存配置和 tokenizer，如果 checkpoint.contents 中包含 ``hf_mode``，则保存 huggingface 模型
│   │   └── dist_ckpt       # 保存分片模型/优化器/rng_states，命名与 Megatron 相同
│   └── critic
│   │   ├── huggingface
│   │   └── dist_ckpt
└── latest_checkpointed_iteration.txt

将 FSDP 和 Megatron Checkpoints 转换为 HuggingFace 格式模型

我们提供了一个工具来将 FSDP 和 Megatron checkpoints 转换为 HuggingFace 格式模型。该工具位于 verl/model_merger。对于不包含 fsdp_config.json 的旧版本 verl checkpoint，可以使用位于 verl/scripts/legacy_model_merger.py 的旧版模型合并工具。

该脚本支持两个主要子命令：merge`（用于转换和保存 checkpoints）和 `test`（用于针对参考模型验证合并后的 checkpoints）。 `merge 子命令的参数如下：

usage: python -m verl.model_merger merge [-h] --backend {fsdp,megatron} [--local_dir LOCAL_DIR] [--tie-word-embedding] [--is-value-model] [--use_cpu_initialization] [--target_dir TARGET_DIR]
                     [--hf_upload_path HF_UPLOAD_PATH] [--private]

options:
-h, --help            显示此帮助消息并退出
--backend {fsdp,megatron}
                        模型的后端
--local_dir LOCAL_DIR
                        已保存模型 checkpoints 的路径
--tie-word-embedding  是否绑定词嵌入权重（目前仅 Megatron 支持）
--is-value-model      模型是否为价值模型（目前仅 Megatron 支持）
--use_cpu_initialization
                        是否对模型使用 CPU 初始化。这对于在初始化过程中无法放入 GPU 内存的大型模型很有用。
--target_dir TARGET_DIR
                        保存合并后的 huggingface 模型的目录
--hf_upload_path HF_UPLOAD_PATH
                        上传模型的 Hugging Face 仓库 ID
--private             是否将模型上传到私有的 Hugging Face 仓库

合并 Megatron checkpoints 的示例用法：

python -m verl.model_merger merge \
    --backend megatron \
    --tie-word-embedding \
    --local_dir checkpoints/verl_megatron_gsm8k_examples/qwen2_5_0b5_megatron_saveload/global_step_1/actor \
    --target_dir /path/to/merged_hf_model

分布式合并 Megatron checkpoints 的示例用法：

torchrun --nproc_per_node 1 --nnodes 8 --node_rank ${RANK} -m verl.model_merger merge \
    --backend megatron \
    --tie-word-embedding \
    --local_dir checkpoints/verl_megatron_gsm8k_examples/qwen2_5_0b5_megatron_saveload/global_step_1/actor \
    --target_dir /path/to/merged_hf_model

合并 FSDP checkpoints 的示例用法：

python -m verl.model_merger merge \
    --backend fsdp \
    --local_dir checkpoints/verl_fsdp_gsm8k_examples/qwen2_5_0b5_fsdp_saveload/global_step_1/actor \
    --target_dir /path/to/merged_hf_model

Megatron Merger 详细信息

当前解码器层（decoder layers）的实现使用 nn.ModuleList 来存储层，因此，每个 PP rank 和 VPP rank 上的模型层索引都从 0 开始。

有 3 种方法可以纠正此行为：

修改解码器层的 state_dict，为每个层的索引添加 offset，从而重写 nn.ModuleList 实现。
在保存 checkpoint 时修改层索引，并在加载 checkpoint 时恢复它们。
Checkpoint merger 完成此工作，仅从 state_dict 计算实际的 offset，这有点复杂。

当前实现使用的是方法 2。

HuggingFace 到 Megatron DistCheckpoint 详细信息

如果您的模型非常巨大，我们建议您使用 Megatron dist-checkpoint 来加载模型。 Megatron dist-checkpoint 支持使用不同类型的模型并行（model parallelism）进行加载，并且比原始的 checkpoint 加载速度快得多。

要将原始 HuggingFace 模型转换为 Megatron dist-checkpoint，您可以使用 scripts/converter_hf_to_mcore.py 脚本。大型 MoE 模型暂时支持 CPU 初始化，这会稍慢一些。我们正在努力提供更好的解决方案来支持大型模型。

转换模型的命令示例如下：

python scripts/converter_hf_to_mcore.py \
    --hf_model_path Qwen/Qwen1.5-MoE-A2.7B-Chat \
    --output_path /mnt/disk/Qwen/Qwen1.5-MoE-A2.7B-Chat \
    --use_cpu_initialization    # 仅适用于 MoE 模型

转换为像 deepseekv3 671B 这样的大型模型的命令示例如下：

torchrun --nproc_per_node 1 --nnodes 8 --node_rank ${RANK} scripts/converter_hf_to_mcore.py \
    --hf_model_path deepseek-ai/DeepSeek-V3 \
    --output_path /mnt/disk/deepseek-ai/DeepSeek-V3 \
    --use_cpu_initialization    # 仅适用于 MoE 模型

原始 Checkpoint 工具

原始 Checkpoint 工具（Original Checkpoint Utils）指的是 verl/models/[model]/megatron/checkpoint_utils 中的原始 checkpoint 实现。

现在只需要原始 checkpoint 工具中的 [model]_loader.py，因为我们不再每次都存储 ``hf_model``（不推荐用于大型模型训练，如果可能，请尝试仅保存分片模型）。

Note

请注意，[model]_loader 仅支持**存储集群可以连接到每个计算节点的**环境。因为它利用**分片加载方式来最小化加载 checkpoint 的开销**。每个 rank 从可以被所有 rank 访问的 state_dict 加载自己的数据。同时，也没有必要在 DP rank 之间进行广播，因为保存的 state_dict 仅由 DP rank 0 生成。

对于**只能将 huggingface 模型放在一个设备上的用户**，我们保留了 [model]_loader_deprecated 中原来的耗时实现。在此实现中，rank 0 将所有权重广播给每个 tp 和 pp rank，然后 dp rank 0 再广播给所有 dp rank。存在 OOM（内存溢出）的风险。

要使用已弃用的加载器，请更改 load_state_dict_to_megatron_llama 的导入包。