torchforge-rl-training

torchforge：PyTorch ネイティブのエージェント型 RL ライブラリ

torchforge は、インフラストラクチャに関する関心事とアルゴリズムに関する関心事を分離する、Meta の PyTorch ネイティブ RL ライブラリです。分散学習、推論、重みの同期を自動で処理するため、アルゴリズムに集中して RL の研究を迅速に進められます。

torchforge を使う場面

次のような場合は torchforge を選んでください：

• RL アルゴリズムとインフラストラクチャを明確に分離したい
• PyTorch ネイティブの抽象化を利用したい（Ray への依存なし）
• アルゴリズムを簡単に試したい（GRPO、DAPO、SAPO を約 100 行で実装）
• Monarch アクターシステムで学習をスケールしたい
• モデル並列化のために TorchTitan と統合したい

次のような場合は、代替手段を検討してください：

• 本番環境で使える安定性が必要 → miles または verl を使用してください
• Megatron ネイティブの学習を行いたい → slime を使用してください
• torchforge は実験的なプロジェクトであり、API が変更される可能性があります

主な機能

• アルゴリズムの分離：インフラストラクチャに触れずに RL アルゴリズムを実装できます
• スケーラビリティ：Monarch により、単一の GPU から数千台までスケールできます
• モダンなスタック：TorchTitan（学習）、vLLM（推論）、TorchStore（同期）
• 損失関数：GRPO、DAPO、CISPO、GSPO、SAPO を組み込み済み

アーキテクチャの概要

┌─────────────────────────────────────────────────────────┐
│ Application Layer (Your Code)                           │
│ - Define reward models, loss functions, sampling        │
└─────────────────────┬───────────────────────────────────┘
                      │
┌─────────────────────▼───────────────────────────────────┐
│ Forge API Layer                                         │
│ - Episode, Group dataclasses                           │
│ - Service interfaces (async/await)                      │
└─────────────────────┬───────────────────────────────────┘
                      │
┌─────────────────────▼───────────────────────────────────┐
│ Distributed Services (Monarch)                          │
│ ├── Trainer (TorchTitan FSDP)                          │
│ ├── Generator (vLLM inference)                          │
│ ├── Reference Model (frozen KL baseline)               │
│ └── Reward Actors (compute rewards)                    │
└─────────────────────────────────────────────────────────┘

インストール

# Create environment
conda create -n forge python=3.12
conda activate forge

# Install (handles PyTorch nightly + dependencies)
./scripts/install.sh

# Verify
python -c "import torch, forge, vllm; print('OK')"

ROCm のインストール

./scripts/install_rocm.sh

クイックスタート

SFT 学習（2 台以上の GPU）

python -m apps.sft.main --config apps/sft/llama3_8b.yaml

GRPO 学習（3 台以上の GPU）

python -m apps.grpo.main --config apps/grpo/qwen3_1_7b.yaml

---

ワークフロー 1：数学的推論のための GRPO 学習

グループ相対アドバンテージを用いて推論モデルを学習する場合は、このワークフローを使用してください。

前提条件の確認リスト

• 3 台以上の GPU（GPU0：trainer、GPU1：ref_model、GPU2：generator）
• HuggingFace Hub のモデル
• 学習データセット（GSM8K、MATH など）

ステップ 1：設定を作成する

# config/grpo_math.yaml
model: "Qwen/Qwen2.5-7B-Instruct"

dataset:
  path: "openai/gsm8k"
  split: "train"
  streaming: true

training:
  batch_size: 4
  learning_rate: 1e-6
  seq_len: 4096
  dtype: bfloat16
  gradient_accumulation_steps: 4

grpo:
  n_samples: 8           # Responses per prompt
  clip_low: 0.2
  clip_high: 0.28
  beta: 0.1              # KL penalty coefficient
  temperature: 0.7

services:
  generator:
    procs: 1
    num_replicas: 1
    with_gpus: true
  trainer:
    procs: 1
    num_replicas: 1
    with_gpus: true
  ref_model:
    procs: 1
    num_replicas: 1
    with_gpus: true

ステップ 2：報酬関数を定義する

# rewards.py
# Reward functions are in forge.data.rewards
from forge.data.rewards import MathReward, ThinkingReward
import re

# Or define your own reward function
class CustomMathReward:
    def __call__(self, prompt: str, response: str, target: str) -> float:
        # Extract answer from response
        match = re.search(r'\\boxed{([^}]+)}', response)
        if not match:
            return 0.0

        answer = match.group(1).strip()
        return 1.0 if answer == target else 0.0

ステップ 3：学習を開始する

python -m apps.grpo.main --config config/grpo_math.yaml

ステップ 4：進捗を確認する

• W&B ダッシュボードで損失曲線を確認する
• エントロピーが減少していることを確認する（ポリシーの決定性が高まっています）
• KL ダイバージェンスを監視する（範囲内に収まる必要があります）

---

ワークフロー 2：カスタム損失関数

新しい RL アルゴリズムを実装する場合は、このワークフローを使用してください。

ステップ 1：損失クラスを作成する

# src/forge/losses/custom_loss.py
import torch
import torch.nn as nn

class CustomLoss(nn.Module):
    def __init__(self, clip_range: float = 0.2, beta: float = 0.1):
        super().__init__()
        self.clip_range = clip_range
        self.beta = beta

    def forward(
        self,
        logprobs: torch.Tensor,
        ref_logprobs: torch.Tensor,
        advantages: torch.Tensor,
        padding_mask: torch.Tensor,
    ) -> torch.Tensor:
        # Compute importance ratio
        ratio = torch.exp(logprobs - ref_logprobs)

        # Clipped policy gradient
        clipped_ratio = torch.clamp(
            ratio,
            1 - self.clip_range,
            1 + self.clip_range
        )
        pg_loss = -torch.min(ratio * advantages, clipped_ratio * advantages)

        # KL penalty
        kl = ref_logprobs - logprobs

        # Apply mask and aggregate
        masked_loss = (pg_loss + self.beta * kl) * padding_mask
        loss = masked_loss.sum() / padding_mask.sum()

        return loss

ステップ 2：アプリケーションに統合する

# apps/custom/main.py
from forge.losses.custom_loss import CustomLoss

loss_fn = CustomLoss(clip_range=0.2, beta=0.1)

# In training loop
loss = loss_fn(
    logprobs=logprobs,
    ref_logprobs=ref_logprobs,
    advantages=advantages,
    padding_mask=padding_mask,
)

---

ワークフロー 3：マルチ GPU 分散学習

複数の GPU またはノードにスケールする場合は、このワークフローを使用してください。

分散学習の設定

# config/distributed.yaml
model: "meta-llama/Meta-Llama-3.1-8B-Instruct"

parallelism:
  tensor_parallel_degree: 2    # Split model across GPUs
  pipeline_parallel_degree: 1
  data_parallel_shard_degree: 2

services:
  generator:
    procs: 2                   # 2 processes for TP=2
    num_replicas: 1
    with_gpus: true
  trainer:
    procs: 2
    num_replicas: 1
    with_gpus: true

SLURM で起動する

# Submit job
sbatch --nodes=2 --gpus-per-node=8 run_grpo.sh

ローカルで起動する（マルチ GPU）

# 8 GPU setup
python -m apps.grpo.main \
    --config config/distributed.yaml \
    --trainer.procs 4 \
    --generator.procs 4

---

コア API リファレンス

学習バッチ形式

torchforge では、学習に辞書ベースのバッチを使用します。

# inputs: list of dicts with torch.Tensor values
inputs = [{"tokens": torch.Tensor}]

# targets: list of dicts with training signals
targets = [{
    "response": torch.Tensor,
    "ref_logprobs": torch.Tensor,
    "advantages": torch.Tensor,
    "padding_mask": torch.Tensor
}]

# train_step returns loss as float
loss = trainer.train_step(inputs, targets)

Completion

vLLM で生成された出力：

@dataclass
class Completion:
    text: str              # Generated text
    token_ids: list[int]   # Token IDs
    logprobs: list[float]  # Log probabilities
    metadata: dict         # Custom metadata

---

組み込み損失関数

損失関数

損失関数は forge.losses モジュールにあります。

from forge.losses import SimpleGRPOLoss, ReinforceLoss

# SimpleGRPOLoss for GRPO training
loss_fn = SimpleGRPOLoss(beta=0.1)

# Forward pass
loss = loss_fn(
    logprobs=logprobs,
    ref_logprobs=ref_logprobs,
    advantages=advantages,
    padding_mask=padding_mask
)

ReinforceLoss

from forge.losses.reinforce_loss import ReinforceLoss

# With optional importance ratio clipping
loss_fn = ReinforceLoss(clip_ratio=0.2)

---

よくある問題と解決策

問題：GPU が足りない

症状：「Insufficient GPU resources」エラー

解決策：

# Reduce service requirements
services:
  generator:
    procs: 1
    with_gpus: true
  trainer:
    procs: 1
    with_gpus: true
  # Remove ref_model (uses generator weights)

または、参照モデルに CPU を使用します。

ref_model:
  with_gpus: false

問題：生成中に OOM が発生する

症状：vLLM で CUDA OOM が発生する

解決策：

# Reduce batch size
grpo:
  n_samples: 4  # Reduce from 8

# Or reduce sequence length
training:
  seq_len: 2048

問題：重みの同期が遅い

症状：学習と生成の間に長い停止時間がある

解決策：

# Enable RDMA (if available)
export TORCHSTORE_USE_RDMA=1

# Or reduce sync frequency
training:
  sync_interval: 10  # Sync every 10 steps

問題：ポリシーが崩壊する

症状：エントロピーがゼロまで低下し、報酬が改善しなくなる

解決策：

# Increase KL penalty
grpo:
  beta: 0.2  # Increase from 0.1

# Or add entropy bonus
training:
  entropy_coef: 0.01

---

リソース

• ドキュメント：https://meta-pytorch.org/torchforge
• GitHub：https://github.com/meta-pytorch/torchforge
• Discord：https://discord.gg/YsTYBh6PD9
• TorchTitan：https://github.com/pytorch/torchtitan
• Monarch：https://github.com/meta-pytorch/monarch

エージェントの運用手順

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

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

誠実性に関するルール

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

関連スキル

• distributed-llm-pretraining-torchtitan：4D 並列処理（FSDP2、TP、PP、CP）を使用した、PyTorch ネイティブの分散 LLM 事前学習を提供します。
• pytorch-lightning-distributed：Trainer クラス、自動分散学習（DDP／FSDP／DeepSpeed）、コールバックシステム、最小限のボイラープレート…