⚡

torchforge强化学习训练Skill torchforge-rl-training

torchforge是Meta的PyTorch原生强化学习库，专为算法与基础设施分离设计，支持GRPO、SFT等训练方法，适用于数学推理、自定义损失函数和多GPU分布式训练。关键词：torchforge, 强化学习, PyTorch, GRPO, 分布式训练, AI应用, 算法实验

AI应用 0 次安装 0 次浏览 更新于 3/21/2026

name: torchforge-rl-training description: 提供使用torchforge进行PyTorch原生智能强化学习的指导，torchforge是Meta的库，将基础设施与算法分离。当您需要清晰的RL抽象、易于算法实验或使用Monarch和TorchTitan进行可扩展训练时使用。 version: 1.0.0 author: Orchestra Research license: MIT tags: [强化学习, PyTorch, GRPO, SFT, Monarch, TorchTitan, Meta] dependencies: [torch>=2.9.0, torchtitan>=0.2.0, vllm, monarch]

torchforge: PyTorch原生智能强化学习库

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

何时使用torchforge

在以下情况选择torchforge：

  • 清晰的RL算法与基础设施分离
  • PyTorch原生抽象（无Ray依赖）
  • 易于算法实验（GRPO、DAPO、SAPO约100行代码）
  • 使用Monarch actor系统进行可扩展训练
  • 与TorchTitan集成进行模型并行

考虑替代方案当：

  • 您需要生产就绪的稳定性 → 使用 miles 或 verl
  • 您想要Megatron原生训练 → 使用 slime
  • torchforge是实验性的，API可能更改

关键特性

  • 算法隔离：实现RL算法而无需接触基础设施
  • 可扩展性：从单GPU到数千GPU通过Monarch
  • 现代堆栈：TorchTitan（训练）、vLLM（推理）、TorchStore（同步）
  • 损失函数：内置GRPO、DAPO、CISPO、GSPO、SAPO

架构概述

┌─────────────────────────────────────────────────────────┐
│ 应用层（您的代码）                                       │
│ - 定义奖励模型、损失函数、采样                           │
└─────────────────────┬───────────────────────────────────┘
                      │
┌─────────────────────▼───────────────────────────────────┐
│ Forge API层                                            │
│ - Episode、Group数据类                                 │
│ - 服务接口（异步/等待）                                  │
└─────────────────────┬───────────────────────────────────┘
                      │
┌─────────────────────▼───────────────────────────────────┐
│ 分布式服务（Monarch）                                   │
│ ├── 训练器（TorchTitan FSDP）                           │
│ ├── 生成器（vLLM推理）                                   │
│ ├── 参考模型（冻结KL基线）                               │
│ └── 奖励演员（计算奖励）                                 │
└─────────────────────────────────────────────────────────┘

安装

# 创建环境
conda create -n forge python=3.12
conda activate forge

# 安装（处理PyTorch nightly + 依赖）
./scripts/install.sh

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

ROCm安装

./scripts/install_rocm.sh

快速开始

SFT训练（2+ GPUs）

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

GRPO训练（3+ GPUs）

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

工作流程1：用于数学推理的GRPO训练

使用此工作流程训练具有组相对优势的推理模型。

先决条件检查表

  • [ ] 3+ GPUs (GPU0: 训练器, GPU1: 参考模型, 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           # 每个提示的响应数
  clip_low: 0.2
  clip_high: 0.28
  beta: 0.1              # KL惩罚系数
  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
# 奖励函数在forge.data.rewards中
from forge.data.rewards import MathReward, ThinkingReward
import re

# 或定义自己的奖励函数
class CustomMathReward:
    def __call__(self, prompt: str, response: str, target: str) -> float:
        # 从响应中提取答案
        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:
        # 计算重要性比率
        ratio = torch.exp(logprobs - ref_logprobs)

        # 剪裁策略梯度
        clipped_ratio = torch.clamp(
            ratio,
            1 - self.clip_range,
            1 + self.clip_range
        )
        pg_loss = -torch.min(ratio * advantages, clipped_ratio * advantages)

        # KL惩罚
        kl = ref_logprobs - logprobs

        # 应用掩码和聚合
        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)

# 在训练循环中
loss = loss_fn(
    logprobs=logprobs,
    ref_logprobs=ref_logprobs,
    advantages=advantages,
    padding_mask=padding_mask,
)

工作流程3：多GPU分布式训练

使用此工作流程扩展到多个GPUs或节点。

分布式配置

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

parallelism:
  tensor_parallel_degree: 2    # 跨GPU分割模型
  pipeline_parallel_degree: 1
  data_parallel_shard_degree: 2

services:
  generator:
    procs: 2                   # TP=2的2个进程
    num_replicas: 1
    with_gpus: true
  trainer:
    procs: 2
    num_replicas: 1
    with_gpus: true

使用SLURM启动

# 提交作业
sbatch --nodes=2 --gpus-per-node=8 run_grpo.sh

本地启动（多GPU）

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

核心API参考

训练批次格式

torchforge使用基于字典的批次进行训练：

# 输入：包含torch.Tensor值的字典列表
inputs = [{"tokens": torch.Tensor}]

# 目标：包含训练信号的字典列表
targets = [{
    "response": torch.Tensor,
    "ref_logprobs": torch.Tensor,
    "advantages": torch.Tensor,
    "padding_mask": torch.Tensor
}]

# train_step返回损失作为浮点数
loss = trainer.train_step(inputs, targets)

完成

从vLLM生成的输出：

@dataclass
class Completion:
    text: str              # 生成的文本
    token_ids: list[int]   # 令牌ID
    logprobs: list[float]  # 对数概率
    metadata: dict         # 自定义元数据

内置损失函数

损失函数

损失函数在forge.losses模块中：

from forge.losses import SimpleGRPOLoss, ReinforceLoss

# 用于GRPO训练的SimpleGRPOLoss
loss_fn = SimpleGRPOLoss(beta=0.1)

# 前向传播
loss = loss_fn(
    logprobs=logprobs,
    ref_logprobs=ref_logprobs,
    advantages=advantages,
    padding_mask=padding_mask
)

ReinforceLoss

from forge.losses.reinforce_loss import ReinforceLoss

# 带有可选重要性比率剪裁
loss_fn = ReinforceLoss(clip_ratio=0.2)

常见问题及解决方案

问题：GPU不足

症状："GPU资源不足"错误

解决方案：

# 减少服务要求
services:
  generator:
    procs: 1
    with_gpus: true
  trainer:
    procs: 1
    with_gpus: true
  # 移除参考模型（使用生成器权重）

或使用CPU作为参考模型：

ref_model:
  with_gpus: false

问题：生成期间OOM

症状：vLLM中的CUDA OOM

解决方案：

# 减少批次大小
grpo:
  n_samples: 4  # 从8减少

# 或减少序列长度
training:
  seq_len: 2048

问题：权重同步慢

症状：训练和生成之间的长暂停

解决方案：

# 启用RDMA（如果可用）
export TORCHSTORE_USE_RDMA=1

# 或减少同步频率
training:
  sync_interval: 10  # 每10步同步一次

问题：策略崩溃

症状：熵降至零，奖励停止改进

解决方案：

# 增加KL惩罚
grpo:
  beta: 0.2  # 从0.1增加

# 或添加熵奖励
training:
  entropy_coef: 0.01

资源