報酬モデル訓練パイプライン(ホウシュウモデルクンレンパイプライン)
LLMの選好データ収集からモデル訓練・評価・デプロイまでの一連のワークフロー。SFT済みモデルをバックボーンとし、ペアワイズ比較データで報酬ヘッドを訓練する標準的なMLパイプラインである。
報酬モデル訓練パイプラインとは
報酬モデル訓練パイプラインは、LLMの出力品質を評価する報酬モデル(Reward Model)を構築するための体系的なワークフローである。データ収集・前処理・モデル訓練・評価・デプロイの各段階を経て、RLHF や Best-of-N サンプリングで実用される RM を完成させる。
2026年現在、HuggingFace TRL・OpenRLHF・DeepSpeed-Chat・Axolotl など複数のフレームワークが RM 訓練パイプラインを提供しており、A100/H100 クラスの GPU があれば 8B〜70B パラメータの RM を数時間〜数日で訓練可能である。
パイプラインの全体構成
報酬モデル訓練パイプラインは以下の5段階で構成される。
Stage 1: SFT(Supervised Fine-Tuning)
RM のバックボーンとなるモデルを instruction-following データで事前微調整する。
- データ: OpenAssistant、ShareGPT、UltraChat(約140万会話)等
- 訓練量: 1〜3エポック、学習率 2e-5、コサインスケジューラ
- 目的: 応答品質の基準線を確立し、RM が「良い応答」と「悪い応答」を識別できる基盤を作る
Stage 2: 選好データ収集
SFT モデルから複数の応答を生成し、人間または AI が比較ランク付けを行う。
- サンプリング: temperature 0.7〜1.0 で同一プロンプトから 4〜8 個の応答を生成
- 比較形式: ペアワイズ(2択)が最も一般的。K個のランキングからは C(K,2) 個のペアを生成可能
- 品質管理: Inter-Annotator Agreement(IAA)を測定し、κ > 0.6 を基準とする
Stage 3: RM 訓練
ペアワイズ比較データを用いて報酬ヘッドを訓練する。
| パラメータ | 推奨値 | 備考 |
|---|---|---|
| 学習率 | 1e-5 〜 5e-6 | SFT より低く設定 |
| バッチサイズ | 64〜128(ペア数) | GPU メモリに依存 |
| エポック数 | 1〜2 | 過学習防止のため少なめ |
| 重み減衰 | 0.01 | 正則化 |
| 最大系列長 | 2048〜4096 | 長い応答も評価するため |
| 損失関数 | Bradley-Terry | ペアワイズランキング損失 |
Stage 4: 評価
RewardBench・MT-Bench・人間評価の3軸で RM の品質を検証する。
Stage 5: デプロイ
訓練済み RM を推論サーバーにデプロイし、RL 訓練ループまたは Best-of-N リランキングに組み込む。vLLM や TensorRT-LLM でバッチ推論を最適化することで、RM 推論のスループットを 10 倍以上向上できる。
フレームワーク比較
| フレームワーク | 言語 | RM 訓練 | RL 統合 | マルチGPU | ライセンス |
|---|---|---|---|---|---|
| TRL (HuggingFace) | Python | RewardTrainer | PPOTrainer/DPOTrainer | DeepSpeed/FSDP | Apache 2.0 |
| OpenRLHF | Python | 専用 RM Trainer | PPO/DPO/KTO | Ray + vLLM | Apache 2.0 |
| DeepSpeed-Chat | Python | Step 2: RM | Step 3: RLHF | DeepSpeed ZeRO-3 | Apache 2.0 |
| Axolotl | Python | YAML設定 | DPO/ORPO | FSDP | Apache 2.0 |
| LLaMA-Factory | Python | GUI/CLI | PPO/DPO/KTO/ORPO | DeepSpeed | Apache 2.0 |
実装例(TRL)
from trl import RewardTrainer, RewardConfig
from transformers import AutoModelForSequenceClassification, AutoTokenizer
model = AutoModelForSequenceClassification.from_pretrained(
"meta-llama/Llama-3.1-8B-Instruct", num_labels=1
)
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.1-8B-Instruct")
config = RewardConfig(
output_dir="./reward_model",
per_device_train_batch_size=4,
gradient_accumulation_steps=16,
learning_rate=1e-5,
num_train_epochs=1,
bf16=True,
max_length=2048,
)
trainer = RewardTrainer(
model=model,
tokenizer=tokenizer,
args=config,
train_dataset=preference_dataset,
)
trainer.train()
コスト見積もり
| 構成 | GPU | 8B RM 訓練時間 | 70B RM 訓練時間 | 概算コスト(クラウド) |
|---|---|---|---|---|
| シングル GPU | A100 80GB | 約4時間 | 不可(VRAM不足) | $8〜$12 |
| 4x GPU | A100 80GB×4 | 約1.5時間 | 約12時間 | $12〜$100 |
| 8x GPU | H100 80GB×8 | 約30分 | 約4時間 | $30〜$120 |
よくある質問(FAQ)
Q1: RM のバックボーンはポリシーモデルと同じサイズにすべきですか? A: 必ずしも同じサイズである必要はない。一般的にポリシーモデルの1/4〜1/2サイズで十分な精度が得られる。ただし、Nemotron-4 340B RM のようにポリシーモデルと同等サイズの RM が最高精度を達成した例もある。
Q2: 訓練データはどのくらいの量が必要ですか? A: 最低10,000ペア、推奨50,000〜100,000ペアである。UltraFeedback(約64,000件)で訓練した 8B RM が RewardBench で90%以上のスコアを達成している実績がある。
Q3: RM 訓練で過学習を防ぐにはどうすればよいですか? A: エポック数を1〜2に制限し、学習率を SFT より低く(1e-5以下)設定する。また、訓練データの多様性確保(複数ドメイン・複数アノテーター)と、バリデーションセットでの早期停止が有効である。
まとめ
- RM 訓練パイプラインは SFT → 選好データ収集 → RM 訓練 → 評価 → デプロイの5段階
- TRL・OpenRLHF 等のフレームワークで数行のコードから実装可能
- 8B RM は A100 1枚で約4時間、概算$10以下で訓練可能
- 過学習防止にはエポック数制限・低学習率・データ多様性確保が重要