openrlhf-training

OpenRLHF - 高性能 RLHF トレーニング

クイックスタート

OpenRLHF は、vLLM による推論高速化を備えた分散トレーニング向けに最適化された、Ray ベースの RLHF フレームワークです。

インストール:

# Launch Docker container
docker run --runtime=nvidia -it --rm --shm-size="10g" --cap-add=SYS_ADMIN \
  -v $PWD:/openrlhf nvcr.io/nvidia/pytorch:25.02-py3 bash

# Uninstall conflicts
sudo pip uninstall xgboost transformer_engine flash_attn pynvml -y

# Install OpenRLHF with vLLM
pip install openrlhf[vllm]

PPO トレーニング（ハイブリッドエンジン）:

ray start --head --node-ip-address 0.0.0.0 --num-gpus 8

ray job submit --address="http://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 4 --vllm_tensor_parallel_size 2 \
  --colocate_all_models \
  --vllm_gpu_memory_utilization 0.5 \
  --pretrain OpenRLHF/Llama-3-8b-sft-mixture \
  --reward_pretrain OpenRLHF/Llama-3-8b-rm-700k \
  --save_path ./output/llama3-8b-rlhf \
  --micro_train_batch_size 8 --train_batch_size 128 \
  --micro_rollout_batch_size 16 --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 --packing_samples \
  --vllm_enable_sleep --deepspeed_enable_sleep

GRPO トレーニング（Group Normalized Policy Optimization）:

# Same command as PPO, but add:
--advantage_estimator group_norm

一般的なワークフロー

ワークフロー 1: 完全な RLHF パイプライン（SFT → 報酬モデル → PPO）

ステップ 1: 報酬モデルのトレーニング（DPO）:

deepspeed --module openrlhf.cli.train_rm \
  --save_path ./output/llama3-8b-rm \
  --save_steps -1 --logging_steps 1 \
  --eval_steps -1 --train_batch_size 256 \
  --micro_train_batch_size 1 --pretrain meta-llama/Meta-Llama-3-8B \
  --bf16 --max_epochs 1 --max_len 8192 \
  --zero_stage 3 --learning_rate 9e-6 \
  --dataset OpenRLHF/preference_dataset_mixture2_and_safe_pku \
  --apply_chat_template --chosen_key chosen \
  --rejected_key rejected --flash_attn --gradient_checkpointing

ステップ 2: PPO トレーニング:

ray start --head --node-ip-address 0.0.0.0 --num-gpus 8

ray job submit --address="http://127.0.0.1:8265" \
  -- 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 4 --vllm_tensor_parallel_size 2 \
  --colocate_all_models \
  --pretrain OpenRLHF/Llama-3-8b-sft-mixture \
  --reward_pretrain ./output/llama3-8b-rm \
  --save_path ./output/llama3-8b-ppo \
  --micro_train_batch_size 8 --train_batch_size 128 \
  --micro_rollout_batch_size 16 --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 \
  --vllm_enable_sleep --deepspeed_enable_sleep

ワークフロー 2: GRPO トレーニング（クリティックモデルが不要）

PPO よりメモリ効率の良い代替手法です:

ray job submit --address="http://127.0.0.1:8265" \
  -- python3 -m openrlhf.cli.train_ppo_ray \
  --advantage_estimator group_norm \
  --ref_num_nodes 1 --ref_num_gpus_per_node 8 \
  --reward_num_nodes 1 --reward_num_gpus_per_node 8 \
  --actor_num_nodes 1 --actor_num_gpus_per_node 8 \
  --vllm_num_engines 4 --vllm_tensor_parallel_size 2 \
  --colocate_all_models \
  --pretrain OpenRLHF/Llama-3-8b-sft-mixture \
  --reward_pretrain OpenRLHF/Llama-3-8b-rm-700k \
  --save_path ./output/llama3-8b-grpo \
  --micro_train_batch_size 8 --train_batch_size 128 \
  --micro_rollout_batch_size 16 --rollout_batch_size 1024 \
  --max_epochs 1 --bf16 \
  --actor_learning_rate 5e-7 \
  --init_kl_coef 0.01 --use_kl_loss --kl_estimator k3 \
  --normalize_reward --no_advantage_std_norm

GRPO の主なパラメーター:

• --advantage_estimator group_norm - GRPO を有効化
• --use_kl_loss - GRPO 論文の KL 損失
• --kl_estimator k3 - 損失関数（k2 ≈ k1）
• --no_advantage_std_norm - 標準偏差による正規化を無効化

ワークフロー 3: DPO トレーニング（選好最適化）

報酬モデルを使わない、よりシンプルな代替手法です:

deepspeed --module openrlhf.cli.train_dpo \
  --save_path ./output/llama3-8b-dpo \
  --save_steps -1 --logging_steps 1 \
  --eval_steps -1 --train_batch_size 256 \
  --micro_train_batch_size 2 --pretrain meta-llama/Meta-Llama-3-8B \
  --bf16 --max_epochs 1 --max_len 8192 \
  --zero_stage 3 --learning_rate 5e-7 --beta 0.1 \
  --dataset OpenRLHF/preference_dataset_mixture2_and_safe_pku \
  --apply_chat_template --chosen_key chosen \
  --rejected_key rejected --flash_attn --gradient_checkpointing

使い分けと代替手法

OpenRLHF を使うべき場合:

• RL で大規模モデル（7B〜70B 以上）をトレーニングする場合
• vLLM による推論高速化が必要な場合
• Ray による分散アーキテクチャを利用したい場合
• マルチノードの GPU クラスターがある場合
• 1 つのフレームワークで PPO/GRPO/RLOO/DPO を使いたい場合

アルゴリズムの選択:

• PPO: 最大限の制御が可能で、複雑な報酬に最適
• GRPO: メモリ効率が高く、クリティックが不要
• RLOO: トークンごとの KL を用いる修正版 PPO
• REINFORCE++: GRPO より安定し、PPO より高速
• DPO: 最もシンプルで、報酬モデルが不要

代わりに使うべきツール:

• TRL: シングルノードでのトレーニングに適した、よりシンプルな API
• veRL: ByteDance による 671B モデル向けフレームワーク
• DeepSpeedChat: DeepSpeed エコシステムと統合されている

よくある問題

問題: 大規模モデルで GPU の OOM が発生する

モデルのコロケーションを無効化します:

# Remove --colocate_all_models flag
# Allocate separate GPUs for each model
--actor_num_gpus_per_node 8 \
--critic_num_gpus_per_node 8 \
--reward_num_gpus_per_node 8 \
--ref_num_gpus_per_node 8

問題: DeepSpeed で GPU インデックスが範囲外になる

環境変数を設定します:

export RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES=1

問題: トレーニングが不安定になる

非同期ではなくハイブリッドエンジンを使用します:

--colocate_all_models \
--vllm_enable_sleep \
--deepspeed_enable_sleep

KL 係数を調整します:

--init_kl_coef 0.05  # Increase from 0.01

問題: PPO 中の生成が遅い

vLLM による高速化を有効にします:

--vllm_num_engines 4 \
--vllm_tensor_parallel_size 2 \
--vllm_gpu_memory_utilization 0.5

応用トピック

ハイブリッドエンジンによる GPU 共有: vLLM のスリープモード、DeepSpeed のスリープモード、最適なノード割り当てについては、references/hybrid-engine.md を参照してください。

アルゴリズムの比較: PPO、GRPO、RLOO、REINFORCE++ のベンチマークとハイパーパラメーターについては、references/algorithm-comparison.md を参照してください。

マルチノードのセットアップ: Ray クラスターの設定とフォールトトレランスについては、references/multi-node-training.md を参照してください。

カスタム報酬関数: 強化ファインチューニングとエージェント RLHF については、references/custom-rewards.md を参照してください。

ハードウェア要件

• GPU: NVIDIA A100/H100 を推奨
• VRAM:
  • 7B モデル: 8× A100 40GB（ハイブリッドエンジン）
  • 70B モデル: 48× A100 80GB（vLLM:Actor:Critic = 1:1:1）
• マルチノード: InfiniBand を備えた Ray クラスターを推奨
• Docker: NVIDIA PyTorch コンテナ 25.02 以降

パフォーマンス:

• DeepSpeedChat の 2 倍高速
• vLLM による推論高速化
• ハイブリッドエンジンにより GPU のアイドル時間を最小化

リソース

• ドキュメント: https://github.com/OpenRLHF/OpenRLHF
• 論文: https://arxiv.org/abs/2405.11143
• サンプル: https://github.com/OpenRLHF/OpenRLHF/tree/main/examples
• Discord: コミュニティサポート