Skip to content

Latest commit

 

History

History
529 lines (434 loc) · 29.3 KB

README_ja.md

File metadata and controls

529 lines (434 loc) · 29.3 KB
OpenRLHF logo

GitHub Contributors Issues Issues GitHub pull requests GitHub stars
オープンソース / 包括的 / 軽量 / 使いやすい

[ English | 中文 | 日本語 ]

OpenRLHFは、Ray、vLLM、ZeRO-3、およびHuggingFace Transformersを基盤とした最初の高性能RLHFフレームワークです：

  • Rayベースの分散アーキテクチャ
    OpenRLHFはRayを活用して効率的な分散スケジューリングを実現します。Actor、Reward、Reference、およびCriticモデルを異なるGPUに分散し、70Bパラメータまでのモデルのトレーニングをサポートします。
    また、Hybrid Engineスケジューリングもサポートしており、すべてのモデルとvLLMエンジンがGPUリソースを共有し、アイドル時間を最小限に抑え、GPU利用率を最大化します。
  • vLLM 推論加速 + AutoTP
    RLHF トレーニングの 80% の時間はサンプル生成段階に費やされます。vLLM と Auto Tensor Parallelism (AutoTP) を活用し、OpenRLHF は高スループットでメモリ効率の良いサンプル生成を実現します。HuggingFace Transformers とのネイティブ統合により、シームレスで高速な生成を保証し、現在最も高速な RLHF フレームワークとなっています。
  • ZeRO-3ベースのメモリ効率の良いトレーニング
    DeepSpeedのZeRO-3とdeepcompileを基盤とし、OpenRLHFは重量級フレームワークなしで大規模モデルのトレーニングを可能にします。HuggingFaceと直接連携し、事前学習済みモデルの簡単なロードと微調整を実現します。
  • 最適化されたPPO実装
    実践ガイドとコミュニティのベストプラクティスに基づいた高度なPPOテクニックを統合し、RLHFワークフローのトレーニング安定性と報酬品質を向上させます。ZhihuとAdvanced Tricks for Training Large Language Models with Proximal Policy Optimizationを参照。

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

ニュース

特徴

クイックスタート

インストール

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:25.02-py3 bash
sudo pip uninstall xgboost transformer_engine flash_attn pynvml -y

# pip install
pip install openrlhf

# vLLM加速を使用する場合（vLLM 0.9.1をインストール）
pip install openrlhf[vllm]
# 最新のvLLMもサポートされています
pip install openrlhf[vllm_latest]
# vLLM、ring-flash-attention、およびLiger-Kernelをインストール
pip install openrlhf[vllm,ring,liger]

# 最新バージョンを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.9.1以降の使用をお勧めします。 また、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]"

テストデータセットの指定方法は？

テストデータセットのパスは --eval_dataset {name or path} を使用して設定してください。

Note

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とvLLMを使用したPPO/REINFORCE++

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

# コンテナ内で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 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 /openrlhf/examples/test_scripts/final/llama3-8b-rlhf \
  --ckpt_path /openrlhf/examples/test_scripts/ckpt/llama3-8b-rlhf \
  --save_hf_ckpt \
  --micro_train_batch_size 8 \
  --train_batch_size 128 \
  --micro_rollout_batch_size 16 \
  --rollout_batch_size 1024 \
  --n_samples_per_prompt 1 \
  --max_epochs 1 \
  --prompt_max_len 1024 \
  --max_samples 100000 \
  --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 \
  --gradient_checkpointing \
  --packing_samples \
  --vllm_sync_backend nccl \
  --enforce_eager \
  --vllm_enable_sleep \
  --deepspeed_enable_sleep \
  --use_wandb {wandb_token}

# REINFORCE++  | RLOO | REINFORCE++-baseline | GRPO | Dr. GRPO をサポート
# --advantage_estimator reinforce | rloo | reinforce_baseline | group_norm | dr_grpo

# --init_kl_coef を 0 に設定すると参照モデルが起動しません

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

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

Note

また、setup_commandsを使用してRayに環境を自動的にデプロイさせることもできます。例：--runtime-env-json='{"setup_commands": ["pip install openrlhf[vllm]"]}'

Note

OpenRLHFのRLOOとREINFORCE++-baselineはREINFORCE++に基づく修正版です：

  • REINFORCE++は、PPOの主要な最適化技術（アドバンテージ正規化やPPO-clipロスなど）を統合し、criticネットワークの必要性を排除します。
  • REINFORCE++-baselineは、同じプロンプトから生成された複数のサンプルの平均報酬をベースラインとして報酬を再形成します（グローバルバッチ正規化 /std を使用）。
  • OpenRLHFのRLOOは、トークンごとのKL報酬を導入し、PPO-clipロスを使用することで元のバージョンを修正しています。
  • Dr. GRPOは、GRPOのグループ正規化 /std を削除します。

Note

deepspeedがGPUデバイスをセットアップする際にインデックス範囲外のエラーが発生した場合は、環境変数 RAY_EXPERIMENTAL_NOSET_*_VISIBLE_DEVICES を設定することで一時的な解決策として対応できます。

# NVIDIA GPUの場合：
export RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES=1

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

強化学習によるファインチューニング (RFT)

OpenRLHFは、便利で効率的な強化学習によるファインチューニングをサポートしています。カスタムreward_func関数を含むファイルを実装し、そのパスをremote_rm_urlパラメータに渡すだけで済みます。例えば：

# reward_func.py
import torch

def reward_func(queries, prompts, labels):
    # queriesはprompts + responses
    # labelsはanswers
    print(queries)

    # 例としてランダムな報酬を生成
    # 実際のアプリケーションでは、これを実際の報酬計算ロジックに置き換える必要があります
    reward = torch.randint(0, 2, (len(queries),)).float()

    return {
        "rewards": reward,  # アドバンテージ計算用の報酬
        "scores": reward,  # 動的フィルタリング用のスコア（0-1報酬）
        "extra_logs": {"dummy_scores": reward},  # wandb用の追加ログ情報
    }

そして、以下のように設定するだけです：

ray job submit --address="http://127.0.0.1:8265" \
  --runtime-env-json='{"working_dir": "/openrlhf"}' \
  -- python3 -m openrlhf.cli.train_ppo_ray \
  ...
  --remote_rm_url /path/to/reward_func.py \
  --label_key answer

ここで、label_keyパラメータは、答えなどの追加のサンプル情報を報酬関数に渡すために使用されます。

非同期RLHFとAgent RLHF

OpenRLHFは、非同期RLHFとエージェントベースのRLHF実装の両方を包括的にサポートしています。これらの機能を使用するには、トレーニング設定に--async_trainと--agent_func_pathパラメータを含めるだけです。

# agent_func.py
step_idx = 0
max_steps = 2

async def step(state, action, label, **kwargs) -> Tuple[float, Dict[str, Any], bool]:
    global step_idx, max_steps
    # 検証後に終了
    if step_idx >= max_steps:
        done = True
        # torch.randを使用してランダムな報酬を生成
        reward = torch.rand(1)
        next_state = state + action + " The answer is correct. <|endoftext|>"
    else:
        done = False
        reward = torch.tensor(0)
        # 状態を更新
        next_state = state + action + " The answer is not correct, please try again: "
    step_idx += 1

    return {
        "rewards": reward,  # アドバンテージ計算用の報酬
        "scores": reward,  # 動的フィルタリング用のスコア（0-1報酬）
        "next_state": next_state,  # 次のステップのvLLMの更新状態
        "done": done,  # エピソードが完了したかどうかを示すブール値
        "sampling_params": kwargs.get("sampling_params", None),  # 次のステップのvLLMサンプリングのパラメータ
        "extra_logs": {"dummy_scores": reward},  # 追加のログ情報
    }

また、export OPENRLHF_ASYNC_NUM_TASKS=128を設定することで、vLLMエンジンごとの最大同時エージェント数を設定できます。 さらに、環境でexport OPENRLHF_ASYNC_QUEUE_SIZE=1（このパラメータはバッファに保存できるデータのバッチ数を制御します）を設定することで、オフポリシーサンプリングの程度を制御できます。

Note

OpenRLHFのAgent RLHFはハイブリッドエンジントレーニングもサポートしています。この機能を有効にするには、--async_trainフラグを削除し、--colocate_all_modelsを有効にしてください。

Warning

非同期トレーニングはトレーニングの安定性に影響を与える可能性があります。ハイブリッドエンジンまたは同期トレーニングモードを優先することをお勧めします。

LoRA

LoRA (Low-Rank Adaptation)を使用する場合、OpenRLHFはデフォルトで完全な重みを保存せず、代わりにLoRA Adapterを保存します。タスクを正常に続行するには、Adapterをベースモデルの重みと結合する必要があります

python -m openrlhf.cli.lora_combiner \
    --model_path meta-llama/Meta-Llama-3-8B \
    --lora_path ./checkpoint/llama3-8b-rm \
    --output_path ./checkpoint/llama-3-8b-rm-combined \
    --is_rm \
    --bf16

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

最適なパフォーマンスを得るために、ノードを vLLM:Actor:Critic = 1:1:1 の比率で割り当てることをお勧めします。

  • 例えば、70Bモデルと48個のA100 GPUの場合、16個のA100 GPUをvLLMエンジンに、16個のGPUをActorモデルに、残りの16個のGPUをCriticモデルに割り当てることをお勧めします。
  • RLアルゴリズムの収束性が要求を満たす場合は、非同期トレーニング --async_train を有効にしてください。
  • GPUメモリが十分にある場合は、分散RLHFではなく、hybrid engine --colocate_all_models と --vllm_enable_sleep および --deepspeed_enable_sleep を使用してください。
  • --colocate_critic_reward、--colocate_actor_ref オプションを有効にしてノードを統合します。
  • rollout_micro_batch_size を可能な限り増やし（vLLMエンジンのTPサイズを最小化）、トレーニングフェーズでは --micro_train_batch_size を大きくし、--packing_samples を有効にしてください。
  • GPUメモリが十分にある場合は、--adam_offload を無効にし、--overlap_comm を有効にしてください。また、--deepcompile を有効にしてトレーニングを高速化してください。
  • vLLMには --vllm_sync_backend nccl を使用してください。
  • n_samples_per_prompts > 1 の場合は、vLLM生成で enable_prefix_caching を有効にしてください。
  • 大規模なベースモデルの場合、OOMが発生した場合は、--colocate_xxxx オプションを使用しないでください。

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

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

参加方法

参加方法は？

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

何ができますか？

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

スポンサー

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

スター履歴

Star History Chart

貢献者

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

参考文献と謝辞

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

私たちのプロジェクトは ColossalChat と DeepSpeedChat にも感謝したいと思います。プロジェクトの初期段階で、彼らのコード設計を参考にしました。 私たちのプロジェクトは、リングアテンション開発のGPUサポートを提供してくれた Netmind.AI にも感謝したいと思います。

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

引用

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}
}

REINFORCE++-baseline

@article{hu2025reinforce++,
  title={Reinforce++: A simple and efficient approach for aligning large language models},
  author={Hu, Jian},
  journal={arXiv preprint arXiv:2501.03262},
  year={2025}
}

OpenRLHF © 2025 OpenRLHF. All Rights Reserved.