torchforge-rl-training

torchforge：PyTorch 原生 Agentic RL 库

torchforge 是 Meta 的 PyTorch 原生 RL 库，旨在将基础设施问题与算法问题分离。它通过自动处理分布式训练、推理和权重同步，让你专注于算法，从而加快 RL 研究。

何时使用 torchforge

以下情况适合选择 torchforge：

• 需要清晰分离 RL 算法与基础设施
• 需要 PyTorch 原生抽象（不依赖 Ray）
• 需要轻松试验算法（约 100 行代码即可实现 GRPO、DAPO、SAPO）
• 需要通过 Monarch actor 系统扩展训练
• 需要与 TorchTitan 集成以实现模型并行

以下情况可考虑其他方案：

• 需要可用于生产环境的稳定性 → 使用 miles 或 verl
• 希望使用 Megatron 原生训练 → 使用 slime
• torchforge 仍处于实验阶段，API 可能会变化

主要功能

• 算法隔离：实现 RL 算法时无需接触基础设施
• 可扩展性：通过 Monarch，从单个 GPU 扩展至数千个 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：训练器，GPU1：ref_model，GPU2：生成器）
• 来自 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)

补全结果

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 数量不足

症状：出现“GPU 资源不足”错误

解决方案：

# 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

Agent 操作流程

1. 检查环境。 检查 GPU 类型、内存和驱动/CUDA 版本（nvidia-smi）、框架版本，以及数据集的位置和大小。
2. 确认输入。 根据数据或用户确认格式、标识符和参数。任何会改变结果的值都应询问，不要猜测。
3. 先运行小规模版本。 进行冒烟测试：使用小型模型或子集，运行少量步骤，并确认损失下降且检查点已保存。
4. 按照上述说明和参考资料执行完整任务。
5. 验证结果。 在留出数据上跟踪指标，与基线进行比较，并记录随机种子、配置和硬件。
6. 汇报。 说明运行了什么（版本、命令、参数）、检查了什么，以及仍有哪些不确定之处。

如果出现这种情况	请这样做
CUDA 内存不足	减小批次大小，启用梯度累积/检查点或混合精度，或者对模型进行分片。
损失为 NaN 或发散	降低学习率，检查数据是否包含无效值，并启用梯度裁剪。
这些说明中的函数、标志或端点在已安装版本中不存在	查阅已安装版本自身的文档（help()、--help、官方文档），进行调整，并告知用户。切勿编造 API。
所需输入、标识符或参数含糊不清	询问用户，或者在运行前明确说明所做的假设。

诚信守则

• 切勿捏造结果、参数、标识符、引文或统计数据。如果无法运行或验证某项内容，请如实说明。
• 没有日志时，切勿声称取得了训练结果；启动大型任务前，请估算计算成本。
• 此处与特定版本相关的细节可能已经过时：请对照已安装版本的官方文档进行确认。
• 对于会产生费用、占用共享 GPU 或云配额、涉及个人或患者数据，或无法撤销的操作，请先征求同意。

相关技能

• distributed-llm-pretraining-torchtitan：使用 torchtitan 提供基于 PyTorch 的分布式大模型预训练，支持 4D 并行（FSDP2、TP、PP、CP）。
• pytorch-lightning-distributed：高级 PyTorch 框架，提供 Trainer 类、自动分布式训练（DDP/FSDP/DeepSpeed）、回调系统，以及精简的样板…