2025-06-09

使用检查点支持容错训练

在整个RLHF训练过程中,可能会出现训练错误或机器故障,
因此建议启用检查点功能以最小化损失。

API接口已在 :ref:config-explain-page 中列出,
此处不再赘述。但仍有一些技术细节需要说明。

.. note::

注意:对于FSDP检查点,``checkpoint.contents`` 字段中除 ``hf_model`` 外均无效,
其余三个字段(model/optimizer/extra)绑定保存和加载。建议同时包含 ``model``, ``optimizer`` 和 ``extra``。

检查点保存目录结构

通常,我们使用 ppo_trainer.yamlppo_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开始。

有三种修正方案:

  1. 修改解码器层的state_dict,为每层索引添加偏移量(offset),重写 nn.ModuleList 实现。
  2. 保存检查点时修正层级索引,加载时恢复原始索引。
  3. 检查点合并器动态计算偏移量(仅通过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`` 的导入包。
最后编辑于
©著作权归作者所有,转载或内容合作请联系作者
平台声明:文章内容(如有图片或视频亦包括在内)由作者上传并发布,文章内容仅代表作者本人观点,简书系信息发布平台,仅提供信息存储服务。