torchforge-rl-training

用途

torchforge 是 Meta 的 PyTorch 原生强化学习（RL）库，其目标是把基础设施关注点与算法关注点分离开，从而支持快速 RL 研究。使用者可以专注于算法，而分布式训练、推理和权重同步由库自动处理。

主要能力

• 算法隔离：无需改动基础设施即可实现 RL 算法。
• 可扩展性：通过 Monarch 从单 GPU 扩展到上千 GPU。
• 现代技术栈：TorchTitan（训练）、vLLM（推理）、TorchStore（同步）。
• 内置损失函数：GRPO、DAPO、CISPO、GSPO、SAPO。
• 架构分层：应用层（奖励模型、损失函数、采样）→ Forge API 层（Episode、Group dataclass、异步服务接口）→ 分布式服务层（Monarch：Trainer 基于 TorchTitan FSDP、Generator 基于 vLLM、Reference Model 冻结 KL 基线、Reward Actors 计算奖励）。

使用方式

安装：

# 创建环境
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+ 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 参考模型、GPU2 生成器）、来自 HuggingFace Hub 的模型、训练数据集（GSM8K、MATH 等）。 步骤包括：创建配置（示例含模型、数据集、训练参数、grpo 参数如 n_samples、clip_low、clip_high、beta、temperature 及 services 配置）、定义奖励函数（可从 forge.data.rewards 引入 MathReward、ThinkingReward，或自定义）、启动训练、监控进度（W&B 损失曲线、熵下降、KL 散度保持有界）。

工作流 2：自定义损失函数

创建损失类（继承 nn.Module，forward 接收 logprobs、ref_logprobs、advantages、padding_mask，实现重要性比率、裁剪策略梯度、KL 惩罚与掩码聚合），然后集成到应用中（如 apps/custom/main.py 引入并调用）。

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

配置并行度（tensor_parallel_degree、pipeline_parallel_degree、data_parallel_shard_degree）与服务进程数，可通过 SLURM 提交（sbatch --nodes=2 --gpus-per-node=8 run_grpo.sh）或本地多 GPU 启动（python -m apps.grpo.main --config config/distributed.yaml --trainer.procs 4 --generator.procs 4）。

核心 API

• 训练批次格式：字典式批次，inputs 为含 tokens 张量的字典列表，targets 含 response、ref_logprobs、advantages、padding_mask；trainer.train_step(inputs, targets) 返回 float 损失。
• Completion dataclass：text、token_ids、logprobs、metadata。
• 内置损失函数位于 forge.losses 模块，如 SimpleGRPOLoss、ReinforceLoss（ReinforceLoss 可选 clip_ratio）。

适用场景

• 需要算法与基础设施清晰分离的 RL 研究。
• 需要 PyTorch 原生抽象（无 Ray 依赖）。
• 需要快速进行算法实验（GRPO、DAPO、SAPO 等约 100 行）。
• 需要借助 Monarch actor 系统进行可扩展训练。
• 需要与 TorchTitan 集成进行模型并行。

限制与替代方案

• torchforge 是实验性的，API 可能变化。
• 需要生产级稳定性时，资料建议使用 miles 或 verl。
• 需要 Megatron 原生训练时，资料建议使用 slime。

常见问题与解决

• GPU 不足：减少服务需求或对参考模型使用 CPU。
• 生成时 OOM：减小 batch size 或序列长度。
• 权重同步慢：启用 RDMA（export TORCHSTORE_USE_RDMA=1）或降低同步频率。
• 策略坍塌：增大 KL 惩罚或加入熵奖励。

资源

文档：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