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 トレーニング（グループ正規化ポリシー最適化）：

# 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 トレーニング（Critic モデル不要）

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 クラスターを利用できる場合
• PPO／GRPO／RLOO／DPO を単一のフレームワークで使いたい場合

アルゴリズムの選択：

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

代わりに次の選択肢を使用してください：

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

よくある問題

問題：大規模モデルで GPU のメモリ不足が発生する

モデルの共存を無効にします：

# 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 共有：references/hybrid-engine.md で、vLLM のスリープモード、DeepSpeed のスリープモード、最適なノード割り当てについて確認できます。

アルゴリズムの比較：references/algorithm-comparison.md で、PPO、GRPO、RLOO、REINFORCE++ のベンチマークとハイパーパラメーターを比較できます。

マルチノードのセットアップ：references/multi-node-training.md で、Ray クラスターの構成とフォールトトレランスについて確認できます。

カスタム報酬関数：references/custom-rewards.md で、強化ファインチューニングとエージェント向け RLHF について確認できます。

ハードウェア要件

• 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：コミュニティサポート

エージェントの作業手順

1. 環境を確認します。 GPU の種類、メモリ、ドライバー／CUDA のバージョン（nvidia-smi）、フレームワークのバージョン、データセットの場所とサイズを確認します。
2. 入力を確定します。 データまたはユーザーに確認し、形式、識別子、パラメーターを確定します。結果に影響する値は推測せず、確認してください。
3. まず小規模なバージョンを実行します。 スモークテストとして、小さなモデルまたはサブセットを使い、少ないステップ数で実行します。損失が減少し、チェックポイントが保存されることを確認します。
4. 上記の手順と参考情報を使って、タスク全体を実行します。
5. 結果を検証します。 保留データでメトリクスを追跡し、ベースラインと比較して、シード、設定、ハードウェアを記録します。
6. 報告します。 実行した内容（バージョン、コマンド、パラメーター）、確認した内容、未解決の不確実な点を述べます。

次の事象が起きた場合	対応
CUDA のメモリ不足	バッチサイズを減らし、勾配累積／チェックポイントまたは混合精度を有効にするか、モデルをシャーディングします。
損失が NaN になる、または発散する	学習率を下げ、データに無効な値がないか確認し、勾配クリッピングを有効にします。
この手順に記載された関数、フラグ、エンドポイントがインストール済みのバージョンに存在しない	インストール済みバージョンのドキュメント（help()、--help、公式ドキュメント）を確認して適応し、ユーザーに伝えます。API をでっち上げてはいけません。
必須の入力、識別子、パラメーターが曖昧	ユーザーに確認するか、実行前に仮定を明示します。

誠実性に関するルール

• 結果、パラメーター、識別子、引用、統計情報を捏造してはいけません。実行または検証できないことは、明確に伝えます。
• ログがない状態で学習結果を主張してはいけません。大規模なジョブを開始する前に、計算コストを見積もります。
• ここに記載されたバージョン固有の情報は古くなっている可能性があるものとして扱い、インストール済みバージョンの公式ドキュメントで確認します。
• 費用が発生する、共有 GPU やクラウドのクォータを消費する、個人データまたは患者データに触れる、あるいは取り消せない操作を行う前に確認します。

関連スキル

• fine-tuning-with-trl：TRL を使った LLM の強化学習によるファインチューニング。指示チューニングの SFT、選好アラインメントの DPO、報酬最適化の PPO／GRPO…
• verl-rl-training：verl（Volcano Engine RL）を使った LLM の強化学習に関するガイダンスを提供します。
• simpo-training：LLM のアラインメントのためのシンプルな選好最適化。