训练中断不用慌：OpenRLHF断点续训全攻略

【免费下载链接】OpenRLHF A Ray-based High-performance RLHF framework (for large models) 项目地址: https://gitcode.com/gh_mirrors/op/OpenRLHF

在大模型训练过程中，意外断电、显存溢出或网络中断等问题时常导致训练被迫中止。重新开始不仅浪费算力，更可能错过最佳训练窗口。OpenRLHF作为基于Ray和vLLM的高性能RLHF框架，提供了完善的断点续训机制。本文将从核心原理、实操步骤到高级技巧，全面解析如何高效恢复训练进程，让70B模型训练也能如丝般顺滑。

断点续训核心机制

OpenRLHF通过三级保障实现断点续训：

1. 自动检查点系统
   框架会定期保存模型权重、优化器状态和训练元数据到指定路径。关键参数包括：

   • --save_steps：每隔N步保存一次检查点（默认-1表示不自动保存）
   • --ckpt_path：指定检查点存储目录，支持绝对路径和相对路径
   • --load_checkpoint：启动时加载最新检查点

2. 分布式状态同步
   在Ray分布式环境下，所有worker节点的状态通过Ray的对象存储保持一致。即使部分节点故障，主节点也能从检查点恢复完整训练状态。

3. 混合引擎资源管理
   通过--colocate_all_models和--vllm_enable_sleep参数，框架可动态调整GPU资源分配，避免恢复训练时的资源冲突。

OpenRLHF的分布式架构支持跨节点检查点同步，图源docs/openrlhf_architecture.svg

基础续训流程（以PPO为例）

1. 配置检查点参数

在训练脚本中添加检查点配置：

# 示例脚本：[examples/scripts/train_ppo_llama_ray.sh](https://link.gitcode.com/i/b360a1f4915ca9a339d91f79f187b339)
ray job submit --address="http://127.0.0.1:8265" \
  --runtime-env-json='{"working_dir": "/openrlhf"}' \
  -- python3 -m openrlhf.cli.train_ppo_ray \
  --pretrain OpenRLHF/Llama-3-8b-sft-mixture \
  --reward_pretrain OpenRLHF/Llama-3-8b-rm-700k \
  --save_path /openrlhf/final_checkpoint \
  --ckpt_path /openrlhf/training_checkpoints \  # 检查点存储路径
  --save_steps 100 \                           # 每100步保存一次
  --load_checkpoint \                          # 启动时加载检查点
  --micro_train_batch_size 8 \
  --train_batch_size 128 \
  # 其他参数...

2. 恢复中断训练

当训练中断后，只需重新提交相同脚本。框架会自动检测--ckpt_path下的最新检查点并恢复：

# 重新提交训练任务
ray job submit --address="http://127.0.0.1:8265" \
  --runtime-env-json='{"working_dir": "/openrlhf"}' \
  -- python3 -m openrlhf.cli.train_ppo_ray \
  # 保持与之前相同的参数...

注意：若修改了关键超参数（如--train_batch_size或--learning_rate），可能导致恢复失败。建议使用examples/scripts/train_reinforce_baseline_ray_agent_multiturn.sh中演示的条件加载方式：

# 条件加载示例
if [ -d "/openrlhf/training_checkpoints" ]; then
  load_args="--load_checkpoint"
else
  load_args=""
fi
ray job submit ... $load_args

高级场景处理

跨设备恢复训练

当需要更换硬件环境（如从8卡A100迁移到4卡H100），需调整张量并行参数：

# 从原训练脚本修改
--vllm_tensor_parallel_size 2 \        # 原参数：4
--ds_tensor_parallel_size 2 \          # 原参数：8
--ckpt_path /openrlhf/training_checkpoints \
--load_checkpoint \

参考examples/scripts/train_dpo_ring_llama.sh中的RingAttention配置，可实现跨节点训练状态迁移。

检查点文件结构

--ckpt_path目录下会生成以下文件：

training_checkpoints/
├── latest  # 指向最新检查点的符号链接
├── step_100/
│   ├── actor_model/          # Actor模型权重
│   ├── critic_model/         # Critic模型权重
│   ├── optimizer_states/     # 优化器状态
│   └── training_metadata.json  # 训练步数、学习率等元数据
├── step_200/
│   └── ...

定期清理旧检查点可使用--max_ckpt_num参数限制保留数量（默认5个）。

处理损坏的检查点

若检查点文件损坏，可从最近的完整步数恢复：

# 指定从特定步数恢复
ray job submit ... \
  --load_checkpoint \
  --ckpt_path /openrlhf/training_checkpoints/step_100 \  # 直接指定步数目录

建议配合examples/scripts/train_prm_mistral.sh中使用的--save_steps 500参数，设置合理的保存间隔。

避坑指南与最佳实践

1. 检查点路径规划

   • 使用共享存储（如NFS）存储检查点，避免节点故障导致数据丢失
   • 生产环境推荐设置--ckpt_path为独立磁盘分区，防止IO竞争

2. 训练稳定性保障

   • 启用--deepspeed_enable_sleep和--vllm_enable_sleep减少资源冲突
   • 参考性能调优指南中的GPU分配策略：vLLM:Actor:Critic=1:1:1

3. LoRA模型特殊处理
   若使用LoRA适配器，恢复后需合并权重：

   python -m openrlhf.cli.lora_combiner \
     --model_path meta-llama/Meta-Llama-3-8B \
     --lora_path /openrlhf/training_checkpoints/latest \
     --output_path /openrlhf/merged_model \
     --bf16
   

4. 监控与告警
   结合--use_wandb参数监控检查点状态，设置以下告警指标：

   • 连续3个检查点的奖励分数无提升
   • 检查点文件大小异常（通常70B模型约140GB）

常见问题解答

Q: 为什么设置了--save_steps 100却没有生成检查点？
A: 可能是因为训练未达到指定步数就中断，或--save_path权限不足。检查examples/scripts/train_dpo_ring_llama.sh中的路径设置，确保容器内有写入权限。

Q: 恢复训练后损失值突然飙升怎么办？
A: 这通常是因为优化器状态未正确加载。尝试删除--ckpt_path下的optimizer_states目录，仅恢复模型权重。

Q: 能否将检查点转换为Hugging Face格式？
A: 可以，添加--save_hf_ckpt参数会在--save_path同时生成HF格式模型。

通过本文介绍的方法，即使面对复杂的70B模型训练，也能实现99%以上的算力利用率。掌握断点续训技巧，是大模型工业化训练的必备技能。建议结合docs/ppo_examples.md中的案例进行实操，遇到问题可查阅README_zh.md的故障排除章节或提交Issue获取社区支持。

下一步行动：

1. 收藏本文以备训练中断时参考
2. 尝试在examples/scripts/train_ppo_llama_ray_70b.sh中添加断点续训参数
3. 关注项目更新，获取动态检查点压缩等新功能

OpenRLHF框架持续迭代中，所有脚本和文档以最新版本为准。

【免费下载链接】OpenRLHF A Ray-based High-performance RLHF framework (for large models) 项目地址: https://gitcode.com/gh_mirrors/op/OpenRLHF