使用檢查點(diǎn)支持容錯(cuò)訓(xùn)練

在整個(gè)RLHF訓(xùn)練過(guò)程中，可能會(huì)出現(xiàn)訓(xùn)練錯(cuò)誤或機(jī)器故障，
因此建議啟用檢查點(diǎn)功能以最小化損失。

API接口已在 :ref:config-explain-page 中列出，
此處不再贅述。但仍有一些技術(shù)細(xì)節(jié)需要說(shuō)明。

.. note::

注意：對(duì)于FSDP檢查點(diǎn)，``checkpoint.contents`` 字段中除 ``hf_model`` 外均無(wú)效，
其余三個(gè)字段（model/optimizer/extra）綁定保存和加載。建議同時(shí)包含 ``model``, ``optimizer`` 和 ``extra``。

檢查點(diǎn)保存目錄結(jié)構(gòu)

通常，我們使用 ppo_trainer.yaml 或 ppo_megatron_trainer.yml 中聲明的 default_local_dir
作為檢查點(diǎn)保存路徑的前綴，即 checkpoints/${trainer.project_name}/${trainer.experiment_name}。

FSDP 的檢查點(diǎn)內(nèi)部結(jié)構(gòu)如下：

.. 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

所有模型分片、優(yōu)化器和附加狀態(tài)均以分片分布式方式存儲(chǔ)。

當(dāng)前 Megatron 的檢查點(diǎn)結(jié)構(gòu)為：

.. code::

checkpoints/${trainer.project_name}/${trainer.experiment_name}
├── global_steps_${i}
│   ├── actor
│   │   ├── huggingface     # 默認(rèn)保存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_cp{c}_dpu0z1t8os.pt
│   │   └── rng_states
│   └── critic
│   │   ├── huggingface
│   │   ├── model
│   │   ├── optim
│   │   └── rng_states
└── latest_checkpointed_iteration.txt

將FSDP和Megatron檢查點(diǎn)轉(zhuǎn)換為HuggingFace格式模型

我們提供了將FSDP和Megatron檢查點(diǎn)轉(zhuǎn)換為HuggingFace格式模型的工具，
位于 scripts/model_merger.py。

該腳本支持兩個(gè)子命令：merge（轉(zhuǎn)換保存檢查點(diǎn)）和 test（驗(yàn)證合并后檢查點(diǎn)與參考模型的一致性）。
merge 子命令的參數(shù)如下：

.. 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
                       模型檢查點(diǎn)保存路徑
--hf_model_path HF_MODEL_PATH
                       (已棄用) 原始Hugging Face模型配置路徑
--tie-word-embedding  是否綁定詞嵌入權(quán)重（當(dāng)前僅Megatron支持）
--is-value-model      是否為價(jià)值模型（當(dāng)前僅Megatron支持）
--target_dir TARGET_DIR
                       合并后HuggingFace模型保存目錄
--hf_upload_path HF_UPLOAD_PATH
                       上傳模型的Hugging Face倉(cāng)庫(kù)ID
--private             是否上傳至私有Hugging Face倉(cāng)庫(kù)

合并Megatron檢查點(diǎn)的示例：

.. 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檢查點(diǎn)的示例：

.. 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合并器實(shí)現(xiàn)細(xì)節(jié)

當(dāng)前解碼器層使用 nn.ModuleList 存儲(chǔ)層級(jí)，
導(dǎo)致每個(gè)PP rank和VPP rank上的模型層索引均從0開始。

有三種修正方案：

1. 修改解碼器層的state_dict，為每層索引添加偏移量（offset），重寫 nn.ModuleList 實(shí)現(xiàn)。
2. 保存檢查點(diǎn)時(shí)修正層級(jí)索引，加載時(shí)恢復(fù)原始索引。
3. 檢查點(diǎn)合并器動(dòng)態(tài)計(jì)算偏移量（僅通過(guò)state_dict），實(shí)現(xiàn)較復(fù)雜。

當(dāng)前采用方案2。

HuggingFace模型轉(zhuǎn)Megatron分布式檢查點(diǎn)細(xì)節(jié)

對(duì)于超大模型，推薦使用Megatron分布式檢查點(diǎn)（dist-checkpoint）加載。
該方式支持不同模型并行配置，加載速度遠(yuǎn)優(yōu)于原始檢查點(diǎn)。

如需將原始HuggingFace模型轉(zhuǎn)為Megatron分布式檢查點(diǎn)，
可使用 scripts/converter_hf_to_mcore.py 腳本。大型MoE模型暫支持CPU初始化（速度較慢），
我們正在優(yōu)化大型模型支持方案。

轉(zhuǎn)換命令示例：

.. 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模型需要

原始檢查點(diǎn)工具集

原始檢查點(diǎn)工具集指 verl/models/[model]/megatron/checkpoint_utils 中的實(shí)現(xiàn)。

目前僅需 [model]_loader.py，因我們移除了每次保存 hf_model 的設(shè)計(jì)（大型模型訓(xùn)練中不推薦，建議僅保存分片模型）。

.. note::

注意：``[model]_loader`` 僅適用于**存儲(chǔ)集群能與所有計(jì)算節(jié)點(diǎn)直連**的環(huán)境。
因其采用**分片加載機(jī)制最小化開銷**——各rank直接從所有節(jié)點(diǎn)可訪問(wèn)的 ``state_dict`` 加載自身數(shù)據(jù)。
由于保存的state_dict僅由DP rank 0生成，無(wú)需在DP rank間廣播。

若**僅能將HuggingFace模型置于單設(shè)備**，可使用舊版實(shí)現(xiàn) ``[model]_loader_deprecated``。
該實(shí)現(xiàn)中：rank 0向所有TP/PP rank廣播權(quán)重，DP rank 0再向所有DP rank廣播，存在內(nèi)存溢出(OOM)風(fēng)險(xiǎn)。

如需使用舊版加載器，請(qǐng)修改 ``load_state_dict_to_megatron_llama`` 的導(dǎo)入包。