オープンソース / 包括的 / 軽量 / 使いやすい

[ English | 中文 | 日本語 ]

OpenRLHFは、Ray、DeepSpeed、およびHF Transformersを基盤とした高性能なRLHFフレームワークです：

シンプルで使いやすい: OpenRLHFは現在利用可能な最もシンプルな高性能RLHFライブラリの一つであり、Huggingfaceのモデルとデータセットとシームレスに互換性があります。
高性能: RLHFトレーニングの80％の時間はサンプル生成段階に費やされます。RayとPacking SamplesおよびvLLM生成加速の能力を活用することで、OpenRLHFのパフォーマンスはOptimized DeepSpeedChat with Hybrid Engineの3〜4倍以上です。
分散RLHF: OpenRLHFは、Actor、Reward、Reference、およびCriticモデルをRayを使用して別々のGPUに分散し、AdamオプティマイザをCPUに配置します。これにより、複数のA100 80G GPUとvLLMを使用して70B+モデルのフルスケールの微調整が可能になり、複数の24GB RTX 4090 GPUで7Bモデルを微調整できます。
PPO実装の最適化: トレーニングの安定性を向上させるために、PPOの実装トリックを統合しました。詳細はZhihuおよびNotionブログを参照してください。

詳細はスライド | 技術報告 | ドキュメントをご覧ください。

ニュース

[2024/12] 私たちは😊 REINFORCE++: A Simple and Efficient Approach for Aligning Large Language Modelsを「提案」しました。
[2024/12] NotionブログでPPO、REINFORCE++、GRPO、およびRLOOを分析しました。

特徴

Rayに基づく分散 PPOおよびREINFORCE++/RLOOの実装。
70億以上のパラメータを持つモデルの完全なRLHF微調整のサポート。
RLHFタスクでの生成を加速するためのvLLMの統合（--vllm_num_engines）。
複数の報酬モデル（--reward_pretrain model1,model2...）およびリモート報酬モデル（--remote_rm_url）のサポート。
DPO（直接選好最適化）/IPO/cDPOおよびKahneman-Tversky Optimization（KTO）の実装。
反復DPO（GitHub: Online-RLHF）のサポート。
拒否サンプリングのサポート。
条件付きSFT（arXiv:2308.12050）の実装。
知識蒸留（Microsoft: minillm）のサポート。
プロセス報酬モデル（PRM）の統合。
SFT、DPO、RM、PRM、およびPPOのトレーニングサンプルのパッキング（--packing_samples）。
RingAttentionの実装（--ring_attn_size、--ring_head_stride）。
専門家の混合モデル（MoE）のサポート（--aux_loss_coef）。
FlashAttention2の統合（--flash_attn）。
QLoRA（--load_in_4bit）およびLoRA（--lora_rank、--target_modules）のサポート。
HuggingFaceのtokenizer.apply_chat_templateとの互換性（--apply_chat_templateおよび--input_key）。
Wandb（--use_wandb）およびTensorBoard（--use_tensorboard）によるログ記録のサポート。
チェックポイントの回復機能（--load_checkpointおよび--save_steps）。
DPOおよびRay PPOなどのマルチノードトレーニングスクリプトを提供。

PPOサポートマトリックス

特徴	OpenRLHF	DSChat	CAIChat	TRL
16 A100-80GBで70B+のフルチューニング	✅	❌	❌	❌
4 RTX4090で7Bのフルチューニング	✅	❌	❌	❌
8 A100-80GBで34B DPOのフルチューニング	✅	❌	❌	❌
PPOでの推論エンジンのサポート	✅	✅	❌	❌
PPO実装のトリック	✅	❌	❌	✅
QLoRAのサポート	✅	❌	❌	✅
Mixtral 8*7bのサポート	✅	❌	❌	❌
未結合のActor-Criticのサポート	✅	✅	✅	❌
複数の報酬モデルのサポート	✅	❌	❌	❌
Huggingfaceモデルのサポート	✅	✅	✅	✅
使いやすさ	✅	❌ (HybridEngineのバグ)	✅	✅

クイックスタート

インストール

OpenRLHFを使用するには、まずDockerコンテナを起動し（推奨）、Dockerコンテナ内でpip installを実行してopenrlhfをインストールします：

# Dockerコンテナを起動
docker run --runtime=nvidia -it --rm --shm-size="10g" --cap-add=SYS_ADMIN -v $PWD:/openrlhf nvcr.io/nvidia/pytorch:24.07-py3 bash
sudo pip uninstall xgboost transformer_engine flash_attn -y

# pip install
pip install openrlhf

# vLLM加速を使用する場合（vLLM 0.6.5をインストール）
pip install openrlhf[vllm]
# 最新のvLLMもサポートされています
pip install openrlhf[vllm_latest]

# 最新バージョンをpip install
pip install git+https://github.com/OpenRLHF/OpenRLHF.git

# またはgit clone
git clone https://github.com/OpenRLHF/OpenRLHF.git
cd OpenRLHF
pip install -e .

Note

vLLM 0.6.4以降の使用をお勧めします。他のバージョン（vLLM >= 0.4.2）は、Glooを介して重みの同期が必要な場合があります（--vllm_sync_backend gloo）。また、vLLM用のDockerfileおよびNvidia-Dockerのワンクリックインストールスクリプトも提供しています。

データセットの準備

OpenRLHFは、データセットクラス内で複数のデータ処理方法を提供しています。例えば、Prompt Datasetでは：

def preprocess_data(data, input_template=None, input_key="input", apply_chat_template=None) -> str:
    if apply_chat_template:
        chat = data[input_key]
        if isinstance(chat, str):
            chat = [{"role": "user", "content": chat}]
        prompt = apply_chat_template(chat, tokenize=False, add_generation_prompt=True)
    else:
        prompt = data[input_key]
        if input_template:
            prompt = input_template.format(prompt)
    return prompt

--input_keyを使用して、入力データセットのJSON key nameを指定し、--prompt_data {name or path}（PPO）または--dataset {name or path}を使用し、--apply_chat_templateを使用してHuggingface Tokenizerのchat_templateを利用できます。
--apply_chat_templateを使用したくない場合は、代わりに--input_templateを使用するか、事前にデータセットをオフラインで前処理することができます。
OpenRLHFは、--prompt_data_probs 0.1,0.4,0.5（PPO）または--dataset_probs 0.1,0.4,0.5を使用して複数のデータセットを混合することもサポートしています。

Chat Templatingの動作方法：

dataset = [{"input_key": [
  {"role": "user", "content": "Hello, how are you?"},
  {"role": "assistant", "content": "I'm doing great. How can I help you today?"},
  {"role": "user", "content": "I'd like to show off how chat templating works!"},
]}]

tokenizer.apply_chat_template(dataset[0]["input_key"], tokenize=False)

"<s>[INST] Hello, how are you? [/INST]I'm doing great. How can I help you today?</s> [INST] I'd like to show off how chat templating works! [/INST]"

トレーニングおよびテストデータセットの指定方法：

data_type@data_dir形式を使用して指定できます。例えば、データセットは--dataset json@./dataとして設定できます。

data
├── test.jsonl
└── train.jsonl

Note

デフォルトでは、trainおよびtestを使用してHuggingfaceのトレーニングおよびテストデータセットを区別します。 JSON keyオプションは特定のデータセットに依存します。詳細はReward DatasetおよびSFT Datasetを参照してください。

教師あり微調整

OpenRLHFのモデルチェックポイントはHuggingFaceモデルと完全に互換性があります。--pretrain {name or path}、--reward_pretrain {name or path}、および--critic_pretrain {name or path}を使用してモデル名またはパスを指定できます。いくつかの事前トレーニング済みチェックポイントとデータセットをHuggingFace OpenRLHFで提供しています。

次に、examples/scriptsディレクトリに提供されている起動スクリプトを使用するか、以下のコマンドを使用してトレーニングを開始できます。

deepspeed --module openrlhf.cli.train_sft \
   --max_len 4096 \
   --dataset Open-Orca/OpenOrca \
   --input_key question \
   --output_key response \
   --input_template $'User: {}\nAssistant: ' \
   --train_batch_size 256 \
   --micro_train_batch_size 2 \
   --max_samples 500000 \
   --pretrain meta-llama/Meta-Llama-3-8B \
   --save_path ./checkpoint/llama3-8b-sft \
   --save_steps -1 \
   --logging_steps 1 \
   --eval_steps -1 \
   --zero_stage 2 \
   --max_epochs 1 \
   --packing_samples \
   --bf16 \
   --flash_attn \
   --learning_rate 5e-6 \
   --gradient_checkpointing \
   --use_wandb {wandb_token}

# HF tokenizer.apply_chat_templateのサポート
# --apply_chat_template 
# --tokenizer_chat_template {HF Chat Template}

# RingAttentionのサポート
# pip install ring_flash_attn
#   --ring_attn_size 2 \
#   --ring_head_stride 2 \

# 継続的な事前トレーニングにも使用できます
# --pretrain_mode

Note

OpenRLHF SFT/DPO/RewardModel/PPOトレーナーは--packing_samplesをサポートしています --flash_attnに基づく

報酬モデルのトレーニング

deepspeed --module openrlhf.cli.train_rm \
   --save_path ./checkpoint/llama3-8b-rm \
   --save_steps -1 \
   --logging_steps 1 \
   --eval_steps -1 \
   --train_batch_size 256 \
   --micro_train_batch_size 1 \
   --pretrain OpenRLHF/Llama-3-8b-sft-mixture \
   --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 \
   --packing_samples \
   --gradient_checkpointing \
   --use_wandb {wandb_token}

報酬モデルの--value_prefix_headオプションをscoreに設定することをお勧めします。これにより、AutoModelForSequenceClassificationを使用してモデルをロードできます：

reward_model = AutoModelForSequenceClassification.from_pretrained(
              reward_model_path,
              num_labels=1,
              torch_dtype=torch.bfloat16,
              attn_implementation="flash_attention_2",
              use_cache=False,
          )
inputs = xxxx (Left Padding Input Tokens)
reward = reward_model.model(*inputs).last_hidden_state
reward = reward_model.score(reward)[:, -1]

Rayを使用しないPPO

deepspeed --module openrlhf.cli.train_ppo \
  --pretrain OpenRLHF/Llama-3-8b-sft-mixture \
  --reward_pretrain OpenRLHF/Llama-3-8b-rm-mixture \
  --save_path ./checkpoint/llama-3-8b-rlhf \
  --save_steps -1 \
  --logging_steps 1 \
  --eval_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 2 \
  --bf16 \
  --actor_learning_rate 5e-7 \
  --critic_learning_rate 9e-6 \
  --init_kl_coef 0.01 \
  --prompt_data OpenRLHF/prompt-collection-v0.1 \
  --input_key context_messages \
  --apply_chat_template \
  --max_samples 100000 \
  --normalize_reward \
  --adam_offload \
  --flash_attn \
  --gradient_checkpointing \
  --use_wandb {wandb_token}

# リモート報酬モデルのサポート（HTTP）
# --remote_rm_url http://localhost:5000/get_reward

RayとvLLMを使用したPPO/REINFORCE++

RLHFトレーニング速度を向上させるか、70Bモデルをサポートするために、RayとvLLM加速を使用したPPOを使用できます

# コンテナ内でRayのマスターノードを起動
ray start --head --node-ip-address 0.0.0.0 --num-gpus 8

# さらに多くのノードでRayを起動する場合は
ray start --address {MASTER-NODE-ADDRESS}:6379  --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 2 \
  --reward_num_nodes 1 \
  --reward_num_gpus_per_node 2 \
  --critic_num_nodes 1 \
  --critic_num_gpus_per_node 2 \
  --actor_num_nodes 1 \
  --actor_num_gpus_per_node 2 \
  --vllm_num_engines 2 \
  --vllm_tensor_parallel_size 2 \
  --colocate_critic_reward \
  --colocate_actor_ref \
  --pretrain OpenRLHF/Llama-3-8b-sft-mixture \
  --reward_pretrain OpenRLHF/Llama-3-8b-rm-mixture \
  --save_path /openrlhf/examples/checkpoint/llama3-8b-rlhf \
  --micro_train_batch_size 8 \
  --train_batch_size 128 \
  --micro_rollout_batch_size 16 \
  --rollout_batch_size 1024 \
  --max_samples 100000 \
  --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 \
  --prompt_data OpenRLHF/prompt-collection-v0.1 \
  --input_key context_messages \
  --apply_chat_template \
  --normalize_reward \
  --packing_samples \
  --adam_offload \
  --flash_attn \
  --gradient_checkpointing \
  --use_wandb {wandb_token}

# REINFORCE++ | RLOOのサポート
# --advantage_estimator reinforce | rloo

# リモート報酬モデルのサポート（HTTP）
# --remote_rm_url http://localhost:5000/get_reward


# Nサンプルのサポート
# --n_samples_per_prompt 4

Note

--vllm_num_enginesを設定しない場合は、vLLMエンジンを使用しないことを意味します。 setup_commandsを使用してRayが自動的に環境をデプロイすることもできます。例えば、--runtime-env-json='{"setup_commands": ["pip install openrlhf[vllm]"]}'。

[!NOTE] OPENRLHFのRLOOは、REINFORCE++を基に改良されたものであり、オリジナル版とは異なります。

Note

deepspeedがGPUデバイスを設定する際にインデックスが範囲外に関連するエラーが発生した場合、環境変数RAY_EXPERIMENTAL_NOSET_*_VISIBLE_DEVICESを設定して回避策を試すことができます。

# NVIDIA GPUの場合:
export RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES=1

サポートされているアルゴリズムの起動スクリプトとドキュメントはexample/scriptsおよびDocuments - Usageにあります。

パフォーマンス

Adamオフロードの有効化、報酬モデル（RM）および参照モデル（Ref）オフロードなどの技術を使用して、DSChatのパフォーマンスを最大限に最適化し、推論段階でのマイクロバッチサイズを増やし、メモリ不足の問題を回避しました。LLaMA2のハイブリッドエンジン（HE）を有効にするために、DSChatのいくつかのバグも修正しました。Optimized DSChatとOpenRLHFを使用して1024のプロンプトを1つのPPOエポックでトレーニングするのにかかる平均時間（秒）は次のとおりです：

サイズ	NVIDIA A800-80GB GPU	Optimized DSChat（ハイブリッドエンジン付き）	OpenRLHF	スピードアップ
7B	16	855.09	471.11	1.82x
13B	32	1528.93	608.93	2.5x
34B	32	3634.98	1526.4	2.4x
70B	32	10407.0	4488.53	2.3x

Note

データは古くなっています。再テストのためにパフォーマンスチューニングセクションを参照してください。

パフォーマンスチューニングガイド

最適なパフォーマンスを達成するために、vLLMエンジンにより多くのノードを割り当てることをお勧めします。例えば、32個のA100 GPUを持つ70Bモデルの場合、16個のA100 GPUをvLLMエンジンに割り当て、8個のGPUをActorモデルに、残りの8個のGPUをCriticモデルに割り当てることをお勧めします。さらに、--colocate_critic_reward、--colocate_actor_refオプションを有効にしてノードをマージします。最後に、rollout_micro_batch_size（およびvLLMエンジンのTPサイズを最小化）を可能な限り増やすべきです。トレーニングフェーズでは、より大きな--micro_train_batch_sizeが望ましく、--packing_samplesを有効にします。十分なGPUがある場合、--adam_offloadを無効にし、--overlap_commを有効にします。マルチノードRLHFの場合、vLLM 0.6.4+で--vllm_sync_backend ncclを使用してください。

OpenRLHFを使用している企業と組織

Google
ByteDance
Tencent
Alibaba
Baidu
China Telecom
Vivo
Allen AI
NexusFlow
Jülich Supercomputing Centre (JSC)
Berkeley Starling Team
M-A-P
...

参加方法

参加方法

janhu9527@gmail.comにメールを送るか、GitHub Organizationに参加してください。以下の詳細を含めてください：
- あなたの名前
- あなたのGitHubユーザー名
- あなたの興味のある分野
- NLPおよび/またはAIに関連するスキルと経験
公式GitHubOpenRLHF ↗プロジェクトページを通じて参加することもできます。貢献したい興味についてのissueを作成するだけで、私たちが連絡します。

何ができるか

チームに参加し、OpenRLHFプロジェクトの開発に参加します。
プロジェクトに貢献するためにプルリクエストを提出します。
ドキュメントの改善、バグの修正、新機能の作成を手伝います。
プロジェクトを共有し、コミュニティの成長を支援します。

スポンサー

スポンサーシップは、OpenRLHFの維持と改善に役立ちます。このプロジェクトが役立つと感じた場合は、スポンサーを検討してください。Open Collective ↗でスポンサーになることができます。

スター履歴

貢献者

すべての貢献者に感謝します！貢献したい場合は、プルリクエストを作成するか、issueを作成してください。

参考文献と謝辞

AIおよびNLP分野への貢献に対して、以下のプロジェクトおよび組織に感謝します：

Hugging Face Transformers ↗
OpenAI GPT ↗
LLaMA ↗
DeepSpeed ↗
Ray ↗

私たちのプロジェクトは、ColossalChatおよびDeepSpeedChatにも感謝します。プロジェクトの初期段階で、彼らのコードデザインを参考にしました。

(2024/7) 私たちのGitHub組織はOpenLLMAIからOpenRLHFに変更されました。

引用

@article{hu2024openrlhf,
  title={OpenRLHF: An Easy-to-use, Scalable and High-performance RLHF Framework},
  author={Jian Hu and Xibin Wu and Zilin Zhu and Xianyu and Weixun Wang and Dehao Zhang and Yu Cao},
  journal={arXiv preprint arXiv:2405.11143},
  year={2024}
}

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

README_ja.md

README_ja.md

ニュース

特徴

PPOサポートマトリックス

クイックスタート

インストール

データセットの準備

教師あり微調整

報酬モデルのトレーニング

Rayを使用しないPPO

RayとvLLMを使用したPPO/REINFORCE++

パフォーマンス

パフォーマンスチューニングガイド

OpenRLHFを使用している企業と組織

参加方法

スポンサー

スター履歴

貢献者

参考文献と謝辞

引用

Files

README_ja.md

Latest commit

History

README_ja.md

File metadata and controls

ニュース

特徴

PPOサポートマトリックス

クイックスタート

インストール

データセットの準備

教師あり微調整

報酬モデルのトレーニング

Rayを使用しないPPO

RayとvLLMを使用したPPO/REINFORCE++

パフォーマンス

パフォーマンスチューニングガイド

OpenRLHFを使用している企業と組織

参加方法

スポンサー

スター履歴

貢献者

参考文献と謝辞

引用