_checkpoint-page: 使用 Checkpoints 支持容错训练 ================================= 最后更新: 06/25/2025. 在整个 RLHF 训练过程中可能会出现训练错误或机器故障,因此建议启用 checkpoints 以最大程度地减少损失。 API 接口已在 :ref:`config-explain-page` 中列出,此处不再重复。但仍有一些技术细节希望在此澄清。 .. 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 结构如下: .. code:: 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 结构如下: .. code:: 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` 子命令的参数如下: .. code:: bash 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 的示例用法: .. code:: bash 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 的示例用法: .. code:: bash 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 的示例用法: .. code:: bash 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 种方法可以纠正此行为: 1. 修改解码器层的 state_dict,为每个层的索引添加 ``offset``,从而重写 ``nn.ModuleList`` 实现。 2. 在保存 checkpoint 时修改层索引,并在加载 checkpoint 时恢复它们。 3. 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 初始化, 这会稍慢一些。我们正在努力提供更好的解决方案来支持大型模型。 转换模型的命令示例如下: .. code:: bash 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 这样的大型模型的命令示例如下: .. code:: bash 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`` 的导入包。