grpo-rl-training

TRL を使った GRPO/RL トレーニング

Transformer Reinforcement Learning (TRL) ライブラリを使って Group Relative Policy Optimization (GRPO) を実装するための専門的なガイダンスです。このスキルは、カスタム報酬関数を用いて言語モデルをファインチューニングするための、実績のあるパターン、重要な知見、本番環境向けのワークフローを提供します。

このスキルを使うタイミング

次のような場合に GRPO トレーニングを使用してください:

• 特定の出力フォーマットを強制する（例: XML タグ、JSON、構造化された推論）
• 客観的な正しさの指標がある検証可能なタスクを教える（数学、コーディング、ファクトチェック）
• Chain-of-thought パターンを報酬として与え、推論能力を向上させる
• ラベル付きの選好データなしで、ドメイン固有の振る舞いにモデルを合わせる
• 複数の目的を同時に最適化する（フォーマット + 正しさ + スタイル）

次の用途には GRPO を使用しないでください:

• 単純な教師ありファインチューニングのタスク（代わりに SFT を使用）
• 明確な報酬シグナルがないタスク
• すでに高品質な選好ペアがある場合（代わりに DPO/PPO を使用）

---

中核となる概念

1. GRPO アルゴリズムの基礎

主要なメカニズム:

• プロンプトごとに複数の completion を生成します（グループサイズ: 4〜16）
• 報酬関数を使って、各グループ内の completion を比較します
• グループ全体に対して相対的に高い報酬を得た応答を優先するよう、ポリシーを更新します

PPO との重要な違い:

• 別個の報酬モデルが不要です
• サンプル効率が高いです（グループ内比較から学習します）
• 実装とデバッグが簡単です

数学的な直感:

For each prompt p:
  1. Generate N completions: {c₁, c₂, ..., cₙ}
  2. Compute rewards: {r₁, r₂, ..., rₙ}
  3. Learn to increase probability of high-reward completions
     relative to low-reward ones in the same group

2. 報酬関数設計の基本方針

重要なルール:

1. 複数の報酬関数を組み合わせる - 各関数が1つの側面（フォーマット、正しさ、スタイル）を担当します
2. 報酬を適切にスケールする - 重みが大きいほど強いシグナルになります
3. 段階的な報酬を使う - 部分的な遵守に対して部分点を与えます
4. 報酬を個別にテストする - 各報酬関数を単独でデバッグします

報酬関数の種類:

種類	用途	重みの例
正しさ	検証可能なタスク（数学、コード）	2.0（最も高い）
フォーマット	厳密な構造の強制	0.5-1.0
長さ	冗長性または簡潔さの促進	0.1-0.5
スタイル	望ましくないパターンへのペナルティ	-0.5 〜 0.5

---

実装ワークフロー

ステップ 1: データセットの準備

重要な要件:

• プロンプトはチャット形式にします（'role' と 'content' を持つ辞書のリスト）
• 期待する振る舞いを設定するためにシステムプロンプトを含めます
• 検証可能なタスクでは、正解を追加の列として含めます

データ構造の例:

from datasets import load_dataset, Dataset

SYSTEM_PROMPT = """
Respond in the following format:
<reasoning>
[Your step-by-step thinking]
</reasoning>
<answer>
[Final answer]
</answer>
"""

def prepare_dataset(raw_data):
    """
    Transform raw data into GRPO-compatible format.

    Returns: Dataset with columns:
    - 'prompt': List[Dict] with role/content (system + user messages)
    - 'answer': str (ground truth, optional but recommended)
    """
    return raw_data.map(lambda x: {
        'prompt': [
            {'role': 'system', 'content': SYSTEM_PROMPT},
            {'role': 'user', 'content': x['question']}
        ],
        'answer': extract_answer(x['raw_answer'])
    })

実践的なヒント:

• 複雑なフォーマットの場合は、システムプロンプトに one-shot または few-shot の例を使用します
• プロンプトは簡潔に保ちます（max_prompt_length: 256-512 トークン）
• トレーニング前にデータの品質を検証します（ゴミを入れればゴミが出てきます）

ステップ 2: 報酬関数の実装

テンプレート構造:

def reward_function_name(
    prompts,        # List[List[Dict]]: Original prompts
    completions,    # List[List[Dict]]: Model generations
    answer=None,    # Optional: Ground truth from dataset
    **kwargs        # Additional dataset columns
) -> list[float]:
    """
    Evaluate completions and return rewards.

    Returns: List of floats (one per completion)
    """
    # Extract completion text
    responses = [comp[0]['content'] for comp in completions]

    # Compute rewards
    rewards = []
    for response in responses:
        score = compute_score(response)
        rewards.append(score)

    return rewards

例 1: 正しさの報酬（数学/コーディング）

def correctness_reward(prompts, completions, answer, **kwargs):
    """Reward correct answers with high score."""
    responses = [comp[0]['content'] for comp in completions]
    extracted = [extract_final_answer(r) for r in responses]
    return [2.0 if ans == gt else 0.0
            for ans, gt in zip(extracted, answer)]

例 2: フォーマットの報酬（構造化された出力）

import re

def format_reward(completions, **kwargs):
    """Reward XML-like structured format."""
    pattern = r'<reasoning>.*?</reasoning>\s*<answer>.*?</answer>'
    responses = [comp[0]['content'] for comp in completions]
    return [1.0 if re.search(pattern, r, re.DOTALL) else 0.0
            for r in responses]

例 3: 段階的なフォーマットの報酬（部分点）

def incremental_format_reward(completions, **kwargs):
    """Award partial credit for format compliance."""
    responses = [comp[0]['content'] for comp in completions]
    rewards = []

    for r in responses:
        score = 0.0
        if '<reasoning>' in r:
            score += 0.25
        if '</reasoning>' in r:
            score += 0.25
        if '<answer>' in r:
            score += 0.25
        if '</answer>' in r:
            score += 0.25
        # Penalize extra text after closing tag
        if r.count('</answer>') == 1:
            extra_text = r.split('</answer>')[-1].strip()
            score -= len(extra_text) * 0.001
        rewards.append(score)

    return rewards

重要な知見: 堅牢なトレーニングのために 3〜5 個の報酬関数を組み合わせます。順序よりもシグナルの多様性が重要です。

ステップ 3: トレーニング設定

メモリ最適化設定（小規模 GPU）

from trl import GRPOConfig

training_args = GRPOConfig(
    output_dir="outputs/grpo-model",

    # Learning rate
    learning_rate=5e-6,          # Lower = more stable
    adam_beta1=0.9,
    adam_beta2=0.99,
    weight_decay=0.1,
    warmup_ratio=0.1,
    lr_scheduler_type='cosine',

    # Batch settings
    per_device_train_batch_size=1,
    gradient_accumulation_steps=4,  # Effective batch = 4

    # GRPO-specific
    num_generations=8,            # Group size: 8-16 recommended
    max_prompt_length=256,
    max_completion_length=512,

    # Training duration
    num_train_epochs=1,
    max_steps=None,               # Or set fixed steps (e.g., 500)

    # Optimization
    bf16=True,                    # Faster on A100/H100
    optim="adamw_8bit",          # Memory-efficient optimizer
    max_grad_norm=0.1,

    # Logging
    logging_steps=1,
    save_steps=100,
    report_to="wandb",            # Or "none" for no logging
)

高性能設定（大規模 GPU）

training_args = GRPOConfig(
    output_dir="outputs/grpo-model",
    learning_rate=1e-5,
    per_device_train_batch_size=4,
    gradient_accumulation_steps=2,
    num_generations=16,           # Larger groups = better signal
    max_prompt_length=512,
    max_completion_length=1024,
    num_train_epochs=1,
    bf16=True,
    use_vllm=True,                # Fast generation with vLLM
    logging_steps=10,
)

重要なハイパーパラメータ:

パラメータ	影響	チューニングのアドバイス
num_generations	比較のためのグループサイズ	8 から始め、GPU に余裕があれば 16 まで増やします
learning_rate	収束速度と安定性	5e-6（安全）、1e-5（高速だがリスクあり）
max_completion_length	出力の冗長さ	タスクに合わせます（推論には 512、短い回答には 256）
gradient_accumulation_steps	実効バッチサイズ	GPU メモリが限られている場合は増やします

ステップ 4: モデルのセットアップとトレーニング

標準セットアップ（Transformers）

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import LoraConfig
from trl import GRPOTrainer

# Load model
model_name = "Qwen/Qwen2.5-1.5B-Instruct"
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.bfloat16,
    attn_implementation="flash_attention_2",  # 2-3x faster
    device_map="auto"
)

tokenizer = AutoTokenizer.from_pretrained(model_name)
tokenizer.pad_token = tokenizer.eos_token

# Optional: LoRA for parameter-efficient training
peft_config = LoraConfig(
    r=16,                         # Rank (higher = more capacity)
    lora_alpha=32,               # Scaling factor (typically 2*r)
    target_modules=[
        "q_proj", "k_proj", "v_proj", "o_proj",
        "gate_proj", "up_proj", "down_proj"
    ],
    task_type="CAUSAL_LM",
    lora_dropout=0.05,
)

# Initialize trainer
trainer = GRPOTrainer(
    model=model,
    processing_class=tokenizer,
    reward_funcs=[
        incremental_format_reward,
        format_reward,
        correctness_reward,
    ],
    args=training_args,
    train_dataset=dataset,
    peft_config=peft_config,      # Remove for full fine-tuning
)

# Train
trainer.train()

# Save
trainer.save_model("final_model")

Unsloth セットアップ（2-3 倍高速）

from unsloth import FastLanguageModel

model, tokenizer = FastLanguageModel.from_pretrained(
    model_name="google/gemma-3-1b-it",
    max_seq_length=1024,
    load_in_4bit=True,
    fast_inference=True,
    max_lora_rank=32,
)

model = FastLanguageModel.get_peft_model(
    model,
    r=32,
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
                    "gate_proj", "up_proj", "down_proj"],
    lora_alpha=32,
    use_gradient_checkpointing="unsloth",
)

# Rest is identical to standard setup
trainer = GRPOTrainer(model=model, ...)
trainer.train()

---

重要なトレーニングの知見

1. 損失の挙動（想定されるパターン）

• 損失は0付近から始まり、学習中に増加します
• これは正常な動作です。損失は初期ポリシーからのKLダイバージェンスを測定しており、モデルは報酬を最適化するために元の挙動から離れていく（学習している）ためです
• 進捗の確認には、損失ではなく報酬の指標を監視してください

2. 報酬の追跡

監視すべき主な指標:

• reward: すべての補完にわたる平均
• reward_std: グループ内の多様性（0より大きい値を維持すべきです）
• kl: 参照モデルからのKLダイバージェンス（緩やかに増加するはずです）

健全な学習パターン:

Step   Reward    Reward_Std   KL
100    0.5       0.3          0.02
200    0.8       0.25         0.05
300    1.2       0.2          0.08  ← Good progression
400    1.5       0.15         0.12

警告サイン:

• 報酬の標準偏差が0に近づく（モデルが単一の応答に収束している）
• KLが急増する（> 0.5）（乖離しすぎているため、学習率を下げてください）
• 報酬が停滞する（報酬関数が厳しすぎるか、モデルの容量の問題）

3. よくある落とし穴と解決策

問題	症状	解決策
モード崩壊	すべての補完が同一	num_generationsを増やす、多様性ペナルティを追加する
学習が進まない	報酬が横ばい	報酬関数のロジックを確認する、学習率を上げる
OOMエラー	GPUメモリ超過	num_generationsを減らす、勾配チェックポイントを有効にする
学習が遅い	< 1 it/s	use_vllm=Trueを有効にする、Unslothを使う、シーケンス長を減らす
フォーマットが無視される	モデルが構造に従わない	フォーマット報酬の重みを上げる、段階的な報酬を追加する

---

応用パターン

1. 多段階学習

複雑なタスクでは、段階的に学習させてください:

# Stage 1: Format compliance (epochs=1)
trainer_stage1 = GRPOTrainer(
    model=model,
    reward_funcs=[incremental_format_reward, format_reward],
    ...
)
trainer_stage1.train()

# Stage 2: Correctness (epochs=1)
trainer_stage2 = GRPOTrainer(
    model=model,
    reward_funcs=[format_reward, correctness_reward],
    ...
)
trainer_stage2.train()

2. 適応的な報酬スケーリング

class AdaptiveReward:
    def __init__(self, base_reward_func, initial_weight=1.0):
        self.func = base_reward_func
        self.weight = initial_weight

    def __call__(self, *args, **kwargs):
        rewards = self.func(*args, **kwargs)
        return [r * self.weight for r in rewards]

    def adjust_weight(self, success_rate):
        """Increase weight if model struggling, decrease if succeeding."""
        if success_rate < 0.3:
            self.weight *= 1.2
        elif success_rate > 0.8:
            self.weight *= 0.9

3. カスタムデータセットの統合

def load_custom_knowledge_base(csv_path):
    """Example: School communication platform docs."""
    import pandas as pd
    df = pd.read_csv(csv_path)

    dataset = Dataset.from_pandas(df).map(lambda x: {
        'prompt': [
            {'role': 'system', 'content': CUSTOM_SYSTEM_PROMPT},
            {'role': 'user', 'content': x['question']}
        ],
        'answer': x['expert_answer']
    })
    return dataset

---

デプロイと推論

LoRAの保存とマージ

# Merge LoRA adapters into base model
if hasattr(trainer.model, 'merge_and_unload'):
    merged_model = trainer.model.merge_and_unload()
    merged_model.save_pretrained("production_model")
    tokenizer.save_pretrained("production_model")

推論の例

from transformers import pipeline

generator = pipeline(
    "text-generation",
    model="production_model",
    tokenizer=tokenizer
)

result = generator(
    [
        {'role': 'system', 'content': SYSTEM_PROMPT},
        {'role': 'user', 'content': "What is 15 + 27?"}
    ],
    max_new_tokens=256,
    do_sample=True,
    temperature=0.7,
    top_p=0.9
)
print(result[0]['generated_text'])

---

ベストプラクティスのチェックリスト

学習前:

• データセットの形式を検証する（promptsはList[Dict]）
• サンプルデータで報酬関数をテストする
• データから想定される最大のprompt_lengthを計算する
• GPUメモリに基づいて適切なnum_generationsを選ぶ
• ロギングを設定する（wandbを推奨）

学習中:

• 報酬の推移を監視する（増加するはずです）
• reward_stdを確認する（0.1より大きい値を維持すべきです）
• OOMエラーに注意する（必要に応じてバッチサイズを減らす）
• 50〜100ステップごとに生成結果をサンプリングする
• ホールドアウトセットでフォーマットの準拠を検証する

学習後:

• PEFTを使用している場合はLoRAの重みをマージする
• 多様なプロンプトでテストする
• ベースラインモデルと比較する
• 報酬の重みとハイパーパラメータを文書化する
• 再現性のための設定を保存する

---

トラブルシューティングガイド

デバッグのワークフロー

1. 報酬関数を切り分ける - 各関数を個別にテストする
2. データの分布を確認する - プロンプトの多様性を確保する
3. 複雑さを減らす - 単一の報酬から始め、徐々に追加する
4. 生成結果を監視する - N ステップごとにサンプルを出力する
5. 抽出ロジックを検証する - 回答のパースが機能することを確認する

クイックフィックス

# Debug reward function
def debug_reward(completions, **kwargs):
    responses = [comp[0]['content'] for comp in completions]
    for i, r in enumerate(responses[:2]):  # Print first 2
        print(f"Response {i}: {r[:200]}...")
    return [1.0] * len(responses)  # Dummy rewards

# Test without training
trainer = GRPOTrainer(..., reward_funcs=[debug_reward])
trainer.generate_completions(dataset[:1])  # Generate without updating

---

参考資料とリソース

公式ドキュメント:

• TRL GRPO Trainer: https://huggingface.co/docs/trl/grpo_trainer
• DeepSeek R1 Paper: https://arxiv.org/abs/2501.12948
• Unsloth Docs: https://docs.unsloth.ai/

サンプルリポジトリ:

• Open R1 Implementation: https://github.com/huggingface/open-r1
• TRL Examples: https://github.com/huggingface/trl/tree/main/examples

推奨文献:

• エージェント向け指示のためのProgressive Disclosure Pattern
• Reward shaping in RL (Ng et al.)
• LoRA paper (Hu et al., 2021)

---

エージェント向けの使用手順

このスキルが読み込まれたら、次の手順に従ってください。

1. GRPO トレーニングを実装する前に、このファイル全体を読んでください
2. セットアップを検証するため、最もシンプルな報酬関数（例：長さベース）から始めてください
3. templates/ ディレクトリのテンプレートを出発点として使用してください
4. タスク固有の実装については、examples/ の例を参照してください
5. ワークフローに順番に従ってください（ステップを飛ばさないでください）
6. 報酬関数は一度に 1 つずつ追加して、段階的にデバッグしてください

重要な注意事項：

• 常に複数の報酬関数を使用してください（3〜5 個が最適です）
• 損失ではなく、報酬の指標を監視してください
• トレーニングの前に報酬関数をテストしてください
• 小さく始め（num_generations=4）、徐々に規模を拡大してください
• チェックポイントは頻繁に保存してください（100 ステップごと）

このスキルはエキスパートレベルの実装を想定しています。初心者は GRPO に取り組む前に、まず教師ありファインチューニングから始めてください。