OpenRLHF 多节点训练实战：基于 Ray 集群的跨机分布式 RLHF 完整指南

OpenRLHF 多节点训练实战基于 Ray 集群的跨机分布式 RLHF 完整指南【免费下载链接】AI-Research-SKILLsComprehensive open-source library of AI research and engineering skills for any AI model. Package the skills and your claude code/codex/gemini agent will be an AI research agent with full horsepower. Maintained by Orchestra Research.项目地址: https://gitcode.com/gh_mirrors/ai/AI-Research-SKILLs本文是 AI-Research-SKILLs 仓库中 OpenRLHF 多节点训练文档 的完整展开版。OpenRLHF 是一个基于 Ray vLLM DeepSpeed ZeRO-3 的高性能 RLHF 框架参见 SKILL 说明本文聚焦其多机多卡部署这一核心场景从零搭建 Ray 跨节点集群、为 Actor/Critic/Reward/Reference 四类模型做 GPU 拓扑分配、启用 Hybrid Engine 资源共用到断点续训与故障恢复。读完本文你将具备在一台头节点master 若干工作节点worker的 GPU 集群上完整跑通 70B 级模型 PPO/GRPO 多节点训练的实操能力并理解 checkpoint 机制与容错设计背后的工作原理。总体架构为什么多节点训练依赖 RayOpenRLHF 用 Ray 作为分布式调度层这是它与单机框架如 TRL最本质的区别。在 RLHF 训练中一个训练任务往往同时承载四类模型Actor策略模型被训练的策略负责输出动作分布Critic价值模型PPO 中的价值函数用于计算优势值Reward奖励模型为生成结果打分Reference参考模型作为 KL 惩罚的基线约束策略不偏离 SFT 模型过远。在多节点场景下这四类模型可以各自跨多台机器、多张 GPU 部署由 Ray 统一调度同时支持基于 checkpoint 的容错与任务自动重新调度。从 Hybrid Engine 文档 可知OpenRLHF 还允许这些模型与 vLLM 推理引擎共享 GPU 资源通过 sleep/wake 周期最大化 GPU 利用率——这也是多节点配置中最重要的调优维度之一。第一步搭建 Ray 集群多节点训练的第一步不是提交训练任务而是先把 Ray 集群拉起来。OpenRLHF 的官方运行环境是 NVIDIA PyTorch 官方 Docker 容器25.02集群分为 1 个 head头节点与 N 个 worker工作节点。1. 启动头节点Master 机器在所有节点上首先以相同参数启动容器注意--shm-size10gvLLM 与 NCCL 通信对共享内存敏感与--cap-addSYS_ADMIN缺一不可# 在 master 节点上启动容器 docker run --runtimenvidia -it --rm --shm-size10g \ --cap-addSYS_ADMIN -v $PWD:/openrlhf \ nvcr.io/nvidia/pytorch:25.02-py3 bash # 在容器内启动 Ray head 节点 ray start --head --node-ip-address 0.0.0.0 --num-gpus 8启动成功后会有类似输出Ray runtime started. Dashboard: http://0.0.0.0:8265--num-gpus 8显式声明该节点的 GPU 资源Ray 调度器将据此分配任务8265是 Ray Dashboard 端口后续监控与提交作业都靠它。2. 连接工作节点每台 Worker 机器每台 worker 机器重复同样的容器启动命令然后用 head 节点的内网 IP 加入集群docker run --runtimenvidia -it --rm --shm-size10g \ --cap-addSYS_ADMIN -v $PWD:/openrlhf \ nvcr.io/nvidia/pytorch:25.02-py3 bash # 连接 head 节点6379 是 Ray GCS 端口 ray start --address {MASTER-NODE-IP}:6379 --num-gpus 8其中{MASTER-NODE-IP}必须替换为 head 节点的真实内网 IP。多节点通信对网络极为敏感若节点间存在防火墙需要放行 6379GCS、8265Dashboard、10001–10100worker 通信等端口。3. 验证集群状态在 head 节点上执行ray status典型输出4 节点、单机 8 卡为例Nodes: 4 - 1 head node (8 GPUs) - 3 worker nodes (8 GPUs each) Total GPUs: 32ray status确认集群规模与 GPU 总量后才可以进入训练提交环节。集群搭建的最小化验证可参考仓库中 Ray Train 多节点说明其给出了ray status、ray.cluster_resources()等配套诊断手段。分布式训练配置多节点 PPO4 节点集群32 GPU训练 70B 模型集群就绪后通过ray job submit提交训练任务。以下命令是 4 节点、32 卡规模下跑 Llama-2-70B PPO 的完整示例ray job submit --addresshttp://127.0.0.1:8265 \ --runtime-env-json{working_dir: /openrlhf} \ -- python3 -m openrlhf.cli.train_ppo_ray \ --ref_num_nodes 1 --ref_num_gpus_per_node 8 \ --reward_num_nodes 1 --reward_num_gpus_per_node 8 \ --critic_num_nodes 1 --critic_num_gpus_per_node 8 \ --actor_num_nodes 1 --actor_num_gpus_per_node 8 \ --vllm_num_engines 2 --vllm_tensor_parallel_size 4 \ --pretrain meta-llama/Llama-2-70b-hf \ --reward_pretrain ./reward-model-70b \ --save_path ./output/llama-70b-ppo \ --ckpt_path ./checkpoints/llama-70b-ppo \ --save_steps 100 --logging_steps 1 \ --micro_train_batch_size 2 --train_batch_size 128 \ --micro_rollout_batch_size 4 --rollout_batch_size 1024 \ --max_epochs 1 --prompt_max_len 1024 --generate_max_len 1024 \ --zero_stage 3 --bf16 \ --actor_learning_rate 5e-7 --critic_learning_rate 9e-6 \ --init_kl_coef 0.01 --normalize_reward \ --gradient_checkpointing --flash_attn对应的 GPU 拓扑分配一目了然——每个模型独占一个节点Node 1Reference 模型8 GPUsNode 2Reward 模型8 GPUsNode 3Critic 模型8 GPUsNode 4Actor 模型8 GPUs这里--runtime-env-json指定工作目录为/openrlhf与容器挂载一致确保所有节点能访问同一份代码与数据集。模型分布参数细粒度控制每个模型占多少节点、多少卡OpenRLHF 为每类模型提供一对参数*_num_nodes与*_num_gpus_per_node两者相乘即该模型占用的总 GPU 数# Actor 模型 --actor_num_nodes 2 # 2 个节点 --actor_num_gpus_per_node 8 # 每节点 8 卡 共 16 卡 # Critic 模型 --critic_num_nodes 1 --critic_num_gpus_per_node 8 # Reward 模型 --reward_num_nodes 1 --reward_num_gpus_per_node 8 # Reference 模型 --ref_num_nodes 1 --ref_num_gpus_per_node 8理解这套参数是掌握多节点调度的钥匙--actor_num_nodes 2 --actor_num_gpus_per_node 8表示 Actor 需要 2 个节点、每节点 8 卡Ray 会据此在集群中寻找满足条件的资源组合。当你打算把某一类模型铺得更开例如 70B 模型需要 2 节点 × 8 卡 16 卡做 ZeRO-3 分片时只需修改对应模型的这对参数。Hybrid Engine多模型共享 GPU2 节点 / 16 卡示例如果不希望四类模型各占一批 GPU可以开启模型共置colocation让多类模型在同一批 GPU 上通过 sleep/wake 周期轮流使用显存。相关参数如下--colocate_all_models # 所有模型共享同一批 GPU --colocate_actor_ref # 仅 Actor Reference 共置 --colocate_critic_reward # 仅 Critic Reward 共置2 节点、16 卡场景下的完整示例ray job submit --addresshttp://127.0.0.1:8265 \ -- python3 -m openrlhf.cli.train_ppo_ray \ --colocate_all_models \ --vllm_enable_sleep --deepspeed_enable_sleep \ --actor_num_nodes 2 --actor_num_gpus_per_node 8 \ --critic_num_nodes 0 --critic_num_gpus_per_node 0 \ --reward_num_nodes 0 --reward_num_gpus_per_node 0 \ --ref_num_nodes 0 --ref_num_gpus_per_node 0 \ --vllm_num_engines 4 --vllm_tensor_parallel_size 4 \ # ... other args注意这里 Critic/Reward/Reference 的节点数全部置 0表示它们不再单独申请资源而是全部塞进 Actor 的 16 卡上同时需要配合--vllm_enable_sleep与--deepspeed_enable_sleep详见下文容错章节实现 GPU 内存的轮流释放。最终所有模型共享 16 卡。关于共置的底层实现Hybrid Engine 文档 给出了关键约束与原理等量约束actor_num_nodes × actor_num_gpus_per_node vllm_num_engines × vllm_tensor_parallel_size必须成立vLLM 引擎才能被调度到 Actor 节点的 GPU 上放置组Placement Group开启--colocate_all_models后框架会创建{bundle: [{GPU: actor_num_gpus_per_node}], strategy: PACK}形式的放置组把 vLLM 引擎与 DeepSpeed 模型压到同一批节点收益对比同一 70B 模型Hybrid Engine 用 48 卡即可vLLM 16 卡与 Actor 16 卡共享、Critic 16 卡、Reward 8 卡、Reference 8 卡而完全隔离需要 64 卡共置可节省约 25% 的 GPU 资源。vLLM 配置采样引擎的并行与显存多节点 PPO 的 rollout 采样由 vLLM 承担其并行配置直接影响吞吐与显存占用。张量并行Tensor Parallelism--vllm_num_engines 4 # 4 个引擎实例 --vllm_tensor_parallel_size 4 # 每个引擎跨 4 卡 共 16 卡vllm_num_enginesvLLM 引擎数量决定采样并发度vllm_tensor_parallel_size单个引擎的张量并行度决定单引擎内部用几张卡切分权重。GPU 显存控制--vllm_gpu_memory_utilization 0.5 # 单卡 50% 显存预留给 vLLM显存计算示例A100 80GB80GB × 0.5 40GB 分配给 vLLM剩余 40GB 留给同卡上的其他模型共置场景下。0.5是保守取值0.7是激进取值若共置后 OOM优先降到0.4。这一点在 Hybrid Engine 文档 的排障章节中有对应说明。检查点Checkpoint机制多节点训练动辄数天checkpoint 是唯一的后悔药。OpenRLHF 的 checkpoint 参数分为三类。开启与格式--save_path ./output/model # 最终模型保存路径 --ckpt_path ./checkpoints/model # 中间 checkpoint 目录 --save_steps 100 # 每 100 步保存一次 --save_value_network # 同时保存 critic价值网络格式相关开关--save_hf_ckpt # 额外导出 HuggingFace 格式便于加载 --use_ds_universal_ckpt # 使用 DeepSpeed universal checkpoint跨 ZeRO 阶段兼容其中--use_ds_universal_ckpt的跨 ZeRO 阶段兼容特性在 DeepSpeed ZeRO-3 资料 中有原理层面的佐证ZeRO-3 会把参数、梯度、优化器状态三类模型状态全部切分universal checkpoint 正是为了在不同切分配置间无损迁移而设计。保存内容每次保存的 checkpoint 包含以下状态示意结构{ global_step: 1000, episode: 10, data_loader_state_dict: {...}, actor_model: {...}, # DeepSpeed checkpoint critic_model: {...} # 仅当开启 --save_value_network }磁盘上的文件布局以--ckpt_path ./checkpoints/llama-70b-ppo为例checkpoints/llama-70b-ppo/ ├── global_step_1000/ │ ├── actor/ │ │ ├── mp_rank_00_model_states.pt │ │ ├── zero_pp_rank_0_mp_rank_00optim_states.pt │ │ └── ... │ └── critic/ (仅当开启 --save_value_network) │ └── ... └── hf_ckpt/ (仅当开启 --save_hf_ckpt) ├── config.json ├── pytorch_model.bin └── ...可以看到 actor 目录下同时存在模型权重文件mp_rank_*_model_states.pt与优化器状态文件zero_pp_rank_*_optim_states.pt这正是 DeepSpeed ZeRO 分片的典型产物。断点续训Resume恢复训练只需在提交命令中加入--load_checkpoint并保持--ckpt_path不变ray job submit --addresshttp://127.0.0.1:8265 \ -- python3 -m openrlhf.cli.train_ppo_ray \ --load_checkpoint # 开启恢复 --ckpt_path ./checkpoints/llama-70b-ppo # checkpoint 目录 # ... other args (必须与原训练一致)恢复逻辑PPOTrainer.fit()内部流程检查ckpt_path下是否存在历史 checkpoint加载最新的一步global_step_*目录恢复global_step、episode、dataloader 状态从该步继续训练。容错机制多节点训练怎么扛住故障多节点训练中节点宕机、OOM 几乎不可避免OpenRLHF 提供了三层防护。1. Ray 内置的任务自动重调度若某 worker 节点失败Ray 会尝试在可用节点上重新调度任务前提是剩余节点有足够资源承接被重调度的任务部分组件可能需要重新初始化Ray 会尽力保持状态一致。2. DeepSpeed Sleep 模式防 OOM--deepspeed_enable_sleep # 不训练时把模型权重卸载到 CPUsleep/wake 周期训练结束后模型状态从 GPU 卸载offload到 CPU释放的 GPU 显存可供其他组件如 vLLM使用下一次训练步骤前再从 CPU 重新加载reload加载过程通过 Ray barrier 同步避免多个模型同时抢占显存。从 Hybrid Engine 文档 可以看到其实现位置在PPOTrainer.ppo_train()中对 actor 调用self.actor.reload_states()/self.actor.offload_states()对 critic 同理初始化完成后若开启--deepspeed_enable_sleep会立即offload_states()让模型先睡在 CPU。OOM 预防的核心在于vLLM 先 wake → generate → sleep之后 DeepSpeed 再 reload → train → offload两者不会同时占据显存。3. 基于 Checkpoint 的灾难恢复训练中断节点崩溃、OOM 等重启 Ray 集群head workers以--load_checkpoint重新提交任务训练从最近一次保存的 step 继续。最佳实践--save_steps 100即每 100 步保存一次 checkpoint将灾难恢复的损失控制在 100 步以内。监控与实验跟踪Ray Dashboardhttp://{HEAD-NODE-IP}:8265Dashboard 可监控节点状态active / idle / failed每节点的 GPU 利用率任务调度情况哪个模型跑在哪个节点上资源使用内存、CPU、GPU。Weights Biases 集成--use_wandb {your-wandb-token} --wandb_org your-org --wandb_project llama-70b-ppoWB 中会记录每步的训练损失training loss奖励分数reward scoresKL 散度KL divergence每节点的 GPU 利用率。性能优化多节点通信与显存效率InfiniBand多节点通信加速对配有 InfiniBand 的节点需在启动 Ray 前设置 NCCL 环境变量让梯度通信走 RDMA 而非以太网export NCCL_IB_HCAmlx5_0 # InfiniBand 设备 export NCCL_SOCKET_IFNAMEib0 export NCCL_IB_DISABLE0 ray start --head --node-ip-address 0.0.0.0 --num-gpus 8原文档记载开启 InfiniBand 后多节点通信可提升 2–3 倍这一做法与仓库 Ray Train 多节点说明 中高性能网络配置一节的思路一致同样设置NCCL_SOCKET_IFNAMEib0、NCCL_IB_DISABLE0。多节点训练中跨机 AllReduce 通信是最大瓶颈之一IB 几乎属于标配。显存与计算优化开关--gradient_checkpointing # 用计算换显存可训练更大模型 --flash_attn # 使用 FlashAttention 2需预先安装 --packing_samples # 多个样本打包进同一 batch提升 GPU 利用率其中--packing_samples通过减少短序列的 padding 浪费来提升吞吐是长 prompt/长生成场景下的常用手段。故障排查Troubleshooting问题 1Worker 节点连不上 Head症状worker 节点启动后无法加入集群。排查检查防火墙与网络head 节点需放行默认端口 6379Redis/GCS、8265Dashboard、10001–10100worker 通信# 在 worker 节点上测试与 head 的连通性 telnet {HEAD-NODE-IP} 6379问题 2训练中节点失败症状Ray 报告 node failure。方案一——从 checkpoint 恢复# 修复故障节点或将其从集群中移除 ray stop # 在故障节点上执行 # 然后用 --load_checkpoint 继续训练方案二——缩减资源规模# 部分节点故障后降低模型占用的节点数 --actor_num_nodes 1 # 原来是 2问题 3多节点仍然 OOM症状明明是多节点集群依然出现显存溢出。方案一——缩小 batch--micro_train_batch_size 1 # 从 2 降下来 --micro_rollout_batch_size 2 # 从 4 降下来方案二——开启 sleep 模式--vllm_enable_sleep --deepspeed_enable_sleep方案三——提高 ZeRO 阶段--zero_stage 3 # 最大化参数/梯度/优化器状态切分问题 4Checkpoint 加载失败症状resume 时报FileNotFoundError。检查 checkpoint 目录ls -la ./checkpoints/llama-70b-ppo/ # 确认 global_step_* 目录存在解决确保--ckpt_path与保存时完全一致--ckpt_path ./checkpoints/llama-70b-ppo # 必须与保存路径相同完整示例8 节点 / 64 GPU 训练 70B 模型最后给出一个覆盖本文全部要点的端到端示例——8 节点、64 卡跑 Llama-2-70B四类模型各占 2 节点 16 卡同时开启 vLLM 4 引擎 × TP4、WB 记录与高频 checkpoint。头节点Node 1ray start --head --node-ip-address 10.0.0.1 --num-gpus 8工作节点Node 2–8ray start --address 10.0.0.1:6379 --num-gpus 8提交训练作业ray job submit --addresshttp://10.0.0.1:8265 \ --runtime-env-json{working_dir: /openrlhf} \ -- python3 -m openrlhf.cli.train_ppo_ray \ --ref_num_nodes 2 --ref_num_gpus_per_node 8 \ --reward_num_nodes 2 --reward_num_gpus_per_node 8 \ --critic_num_nodes 2 --critic_num_gpus_per_node 8 \ --actor_num_nodes 2 --actor_num_gpus_per_node 8 \ --vllm_num_engines 4 --vllm_tensor_parallel_size 4 \ --pretrain meta-llama/Llama-2-70b-hf \ --reward_pretrain ./reward-70b \ --save_path ./output/llama-70b-ppo \ --ckpt_path ./checkpoints/llama-70b-ppo \ --save_steps 100 --save_hf_ckpt \ --micro_train_batch_size 1 --train_batch_size 128 \ --micro_rollout_batch_size 2 --rollout_batch_size 1024 \ --max_epochs 1 --bf16 --zero_stage 3 \ --actor_learning_rate 5e-7 --critic_learning_rate 9e-6 \ --gradient_checkpointing --flash_attn --packing_samples \ --use_wandb {token} --wandb_project llama-70b-ppoGPU 分配Reference16 GPUs2 节点 × 8 卡Reward16 GPUs2 节点 × 8 卡Critic16 GPUs2 节点 × 8 卡Actor16 GPUs2 节点 × 8 卡合计64 GPUs。与其他训练形态的衔接多节点 PPO 是 OpenRLHF 最完整的训练形态实际工作中常有几种变体GRPO / RLOO / REINFORCE这些算法不需要 critic 网络详见 算法对比文档提交命令只需把--critic_num_nodes/--critic_num_gpus_per_node置 0并设置对应的--advantage_estimatorgroup_norm/rloo/reinforce即可省下整批 GPU 用于扩大 Actor 或减小集群规模自定义奖励RFT / Agent RLHF若不用预训练奖励模型可改用--remote_rm_url指定自定义奖励函数或--async_train --agent_func_path走多步智能体训练详见 自定义奖励文档节点分配逻辑与本文一致单机快速验证先在单机上用ray start --head跑通小模型参见 SKILL 快速开始再按本文扩展为多节点是推荐的排错路径。参考资料仓库内延伸阅读OpenRLHF 技能总览含单机快速开始与 GRPO/DPO 工作流Hybrid Engine 架构GPU 共享、sleep/wake 周期与放置组算法对比PPO / GRPO / RLOO / REINFORCE 与超参数自定义奖励函数与 Agent RLHFRay Train 多节点集群搭建与网络配置DeepSpeed ZeRO-3 与 Offload 原理【免费下载链接】AI-Research-SKILLsComprehensive open-source library of AI research and engineering skills for any AI model. Package the skills and your claude code/codex/gemini agent will be an AI research agent with full horsepower. Maintained by Orchestra Research.项目地址: https://gitcode.com/gh_mirrors/ai/AI-Research-SKILLs创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考