使用检查点支持容错训练
在整个RLHF训练过程中,可能会出现训练错误或机器故障,
因此建议启用检查点功能以最小化损失。
API接口已在 :ref:config-explain-page 中列出,
此处不再赘述。但仍有一些技术细节需要说明。
.. note::
注意:对于FSDP检查点,``checkpoint.contents`` 字段中除 ``hf_model`` 外均无效,
其余三个字段(model/optimizer/extra)绑定保存和加载。建议同时包含 ``model``, ``optimizer`` 和 ``extra``。
检查点保存目录结构
通常,我们使用 ppo_trainer.yaml 或 ppo_megatron_trainer.yml 中声明的 default_local_dir
作为检查点保存路径的前缀,即 checkpoints/${trainer.project_name}/${trainer.experiment_name}。
FSDP 的检查点内部结构如下:
.. code::
checkpoints/${trainer.project_name}/${trainer.experiment_name}
├── global_steps_${i}
│ ├── actor
│ │ ├── 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
│ ├── actor_huggingface
│ ├── critic
│ │ ├── 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
└── latest_checkpointed_iteration.txt
所有模型分片、优化器和附加状态均以分片分布式方式存储。
当前 Megatron 的检查点结构为:
.. code::
checkpoints/${trainer.project_name}/${trainer.experiment_name}
├── global_steps_${i}
│ ├── actor
│ │ ├── huggingface # 默认保存tokenizer,若checkpoint.contents包含hf_model则保存HuggingFace模型
│ │ ├── model # 保存分片模型,命名与Megatron一致
│ │ │ ├── mp_rank_xx_yyy # xx为2位TP rank,yyy为3位PP rank
│ │ │ │ └── model_states.pt
│ │ │ └── mp_rank_xx_xxx
│ │ ├── optim
│ │ │ └── distrib_optim_pp{a}_tp{b}_cp{c}_dp{d}.pt
│ │ └── rng_states
│ └── critic
│ │ ├── huggingface
│ │ ├── model
│ │ ├── optim
│ │ └── rng_states
└── latest_checkpointed_iteration.txt
将FSDP和Megatron检查点转换为HuggingFace格式模型
我们提供了将FSDP和Megatron检查点转换为HuggingFace格式模型的工具,
位于 scripts/model_merger.py。
该脚本支持两个子命令:merge(转换保存检查点)和 test(验证合并后检查点与参考模型的一致性)。
merge 子命令的参数如下:
.. code:: bash
usage: model_merger.py merge [-h] --backend {fsdp,megatron} --local_dir LOCAL_DIR [--hf_model_path HF_MODEL_PATH]
[--tie-word-embedding] [--is-value-model] [--target_dir TARGET_DIR]
[--hf_upload_path HF_UPLOAD_PATH] [--private]
options:
-h, --help 显示帮助信息
--backend {fsdp,megatron}
模型后端类型
--local_dir LOCAL_DIR
模型检查点保存路径
--hf_model_path HF_MODEL_PATH
(已弃用) 原始Hugging Face模型配置路径
--tie-word-embedding 是否绑定词嵌入权重(当前仅Megatron支持)
--is-value-model 是否为价值模型(当前仅Megatron支持)
--target_dir TARGET_DIR
合并后HuggingFace模型保存目录
--hf_upload_path HF_UPLOAD_PATH
上传模型的Hugging Face仓库ID
--private 是否上传至私有Hugging Face仓库
合并Megatron检查点的示例:
.. code:: bash
python scripts/model_merger.py 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检查点的示例:
.. code:: bash
python scripts/model_merger.py 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合并器实现细节
当前解码器层使用 nn.ModuleList 存储层级,
导致每个PP rank和VPP rank上的模型层索引均从0开始。
有三种修正方案:
- 修改解码器层的state_dict,为每层索引添加偏移量(offset),重写
nn.ModuleList实现。 - 保存检查点时修正层级索引,加载时恢复原始索引。
- 检查点合并器动态计算偏移量(仅通过state_dict),实现较复杂。
当前采用方案2。
HuggingFace模型转Megatron分布式检查点细节
对于超大模型,推荐使用Megatron分布式检查点(dist-checkpoint)加载。
该方式支持不同模型并行配置,加载速度远优于原始检查点。
如需将原始HuggingFace模型转为Megatron分布式检查点,
可使用 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模型需要
原始检查点工具集
原始检查点工具集指 verl/models/[model]/megatron/checkpoint_utils 中的实现。
目前仅需 [model]_loader.py,因我们移除了每次保存 hf_model 的设计(大型模型训练中不推荐,建议仅保存分片模型)。
.. note::
注意:``[model]_loader`` 仅适用于**存储集群能与所有计算节点直连**的环境。
因其采用**分片加载机制最小化开销**——各rank直接从所有节点可访问的 ``state_dict`` 加载自身数据。
由于保存的state_dict仅由DP rank 0生成,无需在DP rank间广播。
若**仅能将HuggingFace模型置于单设备**,可使用旧版实现 ``[model]_loader_deprecated``。
该实现中:rank 0向所有TP/PP rank广播权重,DP rank 0再向所有DP rank广播,存在内存溢出(OOM)风险。
如需使用旧版加载器,请修改 ``load_state_dict_to_megatron_llama`` 的导入包。