强化学习基础设施(RL Infra)优化

一、RL Infra 核心问题与挑战

1.1 为什么RL Infra比SFT/预训练更难？

维度	预训练/SFT	RL训练
计算模式	纯forward+backward	generate → reward → update 交替
GPU利用率	高（连续计算）	低（生成阶段GPU空闲）
数据流	静态数据集	在线采样，数据不断变化
内存	模型+优化器状态	模型×多个（policy/ref/reward/value）
通信	梯度同步	模型同步+经验传输
调度	简单	actor/learner分离，复杂调度

1.2 RL Infra优化的核心目标

最大化硬件利用率 = 最大化有效token吞吐量
                 = 减少GPU空闲时间
                 + 减少内存占用
                 + 减少通信开销
                 + 提高采样效率

二、RL Infra 技术方向全景

2.1 方向总览

RL Infra 优化
├── 1. 系统架构
│   ├── Actor-Learner 分离
│   ├── 生成-训练解耦
│   ├── 异步流水线
│   └── 资源调度
├── 2. 生成优化
│   ├── 推理引擎 (vLLM/SGLang)
│   ├── 批量生成
│   ├── 采样策略
│   └── KV Cache管理
├── 3. 训练优化
│   ├── PPO/GRPO/DPO
│   ├── 梯度累积
│   ├── 混合精度
│   └── 参数高效更新
├── 4. 内存优化
│   ├── 多模型管理
│   ├── 激活检查点
│   ├── Offloading
│   └── LoRA/QLoRA
├── 5. 采样优化
│   ├── 经验回放
│   ├── 优先级采样
│   ├── 课程学习
│   └── 数据过滤
├── 6. 环境加速
│   ├── GPU加速环境
│   ├── 批量环境
│   ├── 异步环境
│   └── 模拟器优化
├── 7. 通信优化
│   ├── 模型同步
│   ├── 梯度压缩
│   ├── 参数服务器
│   └── 流水线通信
└── 8. 鲁棒性
    ├── 容错
    ├── 弹性伸缩
    ├── 训练稳定性
    └── 奖励hacking防护

三、必读论文列表

3.1 经典分布式RL系统

#	论文	年份	核心贡献	方向
1	IMPALA: Scalable Distributed Deep-RL with Importance Weighted Actor-Learner Architectures	2018	Actor-Learner分离架构，V-trace修正	系统架构
2	SEED RL: Scalable and Efficient Deep-RL with Accelerated Central Inference	2020	GPU集中推理，LSTM+ResNet高效推理	生成优化
3	Distributed Prioritized Experience Replay (Ape-X)	2018	分布式优先经验回放，Actor生成+Learner训练	采样优化
4	RL²: Fast Reinforcement Learning via Slow Reinforcement Learning	2016	元学习RL，循环策略	算法
5	R2D2: Recurrent Experience Replay in Distributed Reinforcement Learning	2021	循环策略的分布式经验回放	采样优化
6	Gorila: General Reinforcement Learning Architecture	2015	参数服务器架构，分布式DQN	系统架构

3.2 高吞吐RL系统

#	论文	年份	核心贡献	方向
7	Sample Factory: Egocentric 3D Control from Pixels at 100000 FPS with Asynchronous Reinforcement Learning	2020	极高吞吐异步RL，100K FPS	系统架构
8	TorchBeast: A PyTorch Platform for Distributed RL	2020	IMPALA的PyTorch实现	系统框架
9	Moolib: A Library for Distributed RL	2023	Facebook分布式RL库	系统框架
10	TorchRL: A modular, primitive-first, PyTorch RL library	2023	模块化RL库，支持分布式	系统框架
11	CleanRL: High-quality Single-file Implementations of RL Algorithms	2022	高质量单文件RL实现	系统框架

3.3 RLHF/LLM-RL 系统

#	论文	年份	核心贡献	方向
12	Training language models to follow instructions with human feedback (InstructGPT)	2022	RLHF完整流程：SFT → RM → PPO	RLHF系统
13	DeepSpeed-Chat: Easy, Fast and Affordable RLHF Training of ChatGPT-like Models at All Scales	2023	DeepSpeed RLHF系统，3D并行+ZeRO	RLHF系统
14	OpenRLHF: An Easy-to-use, Scalable and High-performance RLHF Framework	2024	开源RLHF框架，Ray调度	RLHF系统
15	ReaLHF: A Unified and Efficient Framework for RLHF Training	2024	动态重并行，统一RLHF训练	RLHF系统
16	RLHFlow: An Open RLHF Framework	2024	开源RLHF框架	RLHF系统
17	Alignment Handbook	2024	HuggingFace对齐全书	RLHF系统

3.4 RL算法优化（减少Infra负担）

#	论文	年份	核心贡献	方向
18	DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models (GRPO)	2024	Group Relative Policy Optimization，无需Critic模型	算法优化
19	ReMax: A Simple, Effective, and Efficient RL Method for Aligning LLMs	2023	REINFORCE变体，无需Value模型	算法优化
20	Back to Basics: Revisiting REINFORCE Style Optimization for Learning from Human Feedback (RLOO)	2024	REINFORCE Leave-One-Out，简单高效	算法优化
21	Direct Preference Optimization (DPO)	2023	无需RM和RL，直接偏好优化	算法优化
22	SimPO: Simple Preference Optimization with a Reference-Free Reward	2024	无参考模型的偏好优化	算法优化
23	KTO: Model Alignment as Prospect Theoretic Optimization	2024	基于前景理论的对齐，无需配对数据	算法优化
24	ORPO: Monolithic Preference Optimization without Reference Model	2024	无参考模型，单阶段优化	算法优化
25	Self-Play Fine-Tuning Converts Weak Language Models to Strong Language Models (SPIN)	2024	自博弈微调	算法优化

3.5 推理/生成优化

#	论文	年份	核心贡献	方向
26	vLLM: Efficient Memory Management for Large Language Model Serving with PagedAttention	2023	PagedAttention，连续批处理	推理引擎
27	SGLang: Efficient Execution of Structured Language Model Programs	2024	结构化生成，RadixAttention	推理引擎
28	TensorRT-LLM	2023	NVIDIA高性能LLM推理	推理引擎
29	DeepSpeed-FastGen: High-throughput Text Generation for LLMs via MII and DeepSpeed-Inference	2024	DeepSpeed推理引擎	推理引擎
30	Splitwise: Efficient Generative LLM Inference Using Phase Splitting	2023	prefill/decode分离	推理优化

3.6 环境加速

#	论文	年份	核心贡献	方向
31	Isaac Gym: High Performance GPU-Based Physics Simulation For Robot Learning	2021	GPU加速机器人仿真	环境加速
32	Brax: A Differentiable Physics Engine for Large Scale Rigid Body Simulation	2021	JAX可微分物理引擎	环境加速
33	EnvPool: A C++ based high-performance reinforcement learning environment execution engine	2022	C++批量环境引擎	环境加速
34	GPU-accelerated Atari Emulation for RL	2021	GPU加速Atari	环境加速
35	Mava: A research framework for distributed multi-agent reinforcement learning	2021	分布式多智能体RL	环境加速

3.7 内存与计算优化

#	论文	年份	核心贡献	方向
36	LoRA: Low-Rank Adaptation of Large Language Models	2021	低秩适配，减少可训练参数	参数高效
37	QLoRA: Efficient Finetuning of Quantized LLMs	2023	4-bit量化+LoRA	内存优化
38	ZeRO-Offload: Democratizing Billion-Scale Model Training	2021	CPU offload	内存优化
39	Mixed Precision Training	2018	FP16/BF16混合精度	计算优化
40	Activation Checkpointing	2016	梯度检查点	内存优化

3.8 奖励与评估

#	论文	年份	核心贡献	方向
41	Training a Helpful and Harmless Assistant with RLHF (Reward Modeling)	2022	奖励模型训练方法	奖励系统
42	Reward Model Ensembles Help Mitigate Overoptimization	2023	奖励模型集成	奖励系统
43	Scaling Laws for Reward Model Overoptimization	2022	奖励hacking的缩放定律	奖励系统
44	Rule-Based Rewards for Language Model Alignment (RBR)	2024	基于规则的奖励	奖励系统
45	Generative Reward Models	2024	生成式奖励模型	奖励系统

四、RL Infra 各方向详解

4.1 系统架构

4.1.1 Actor-Learner 分离

核心思想：将数据采样(Actor)和梯度更新(Learner)分离到不同进程/机器。

┌─────────────────────────────────────────────────────┐
│                    RL Infra 架构                       │
│                                                       │
│  ┌──────────┐    ┌──────────┐    ┌──────────┐        │
│  │ Actor 1  │    │ Actor 2  │    │ Actor N  │        │
│  │ (推理)    │    │ (推理)    │    │ (推理)    │        │
│  └────┬─────┘    └────┬─────┘    └────┬─────┘        │
│       │               │               │               │
│       └───────────────┼───────────────┘               │
│                       │ 经验缓冲                       │
│                       ▼                               │
│              ┌─────────────────┐                      │
│              │  Replay Buffer  │                      │
│              │  (经验存储)      │                      │
│              └────────┬────────┘                      │
│                       │                               │
│       ┌───────────────┼───────────────┐               │
│       ▼               ▼               ▼               │
│  ┌──────────┐    ┌──────────┐    ┌──────────┐        │
│  │Learner 1 │    │Learner 2 │    │Learner N │        │
│  │ (训练)    │    │ (训练)    │    │ (训练)    │        │
│  └──────────┘    └──────────┘    └──────────┘        │
└─────────────────────────────────────────────────────┘

关键挑战：

Actor和Learner之间的模型同步延迟
经验数据传输带宽
Actor推理效率

代表论文：

IMPALA (2018): V-trace修正off-policy偏差
SEED RL (2020): GPU集中推理
Sample Factory (2020): 异步高吞吐

4.1.2 生成-训练解耦（LLM-RL特有）

核心思想：LLM-RL中，生成(rollout)和训练(update)使用不同的并行策略。

生成阶段（推理优化）           训练阶段（训练优化）
┌─────────────────┐         ┌─────────────────┐
│  vLLM/SGLang    │         │  DeepSpeed      │
│  - TP并行        │   ──→   │  - 3D并行        │
│  - PagedAttention│         │  - ZeRO-3       │
│  - 连续批处理    │         │  - 梯度累积      │
└─────────────────┘         └─────────────────┘

关键挑战：

模型在推理引擎和训练框架之间转换
推理和训练的GPU利用率不均衡
长序列生成的内存管理

代表论文：

DeepSpeed-Chat (2023): 推理-训练统一框架
OpenRLHF (2024): Ray调度推理-训练分离
ReaLHF (2024): 动态重并行

4.2 生成优化

4.2.1 推理引擎选择

引擎	特点	适用场景
vLLM	PagedAttention，连续批处理	通用LLM推理
SGLang	RadixAttention，结构化生成	复杂prompt/多轮对话
TensorRT-LLM	NVIDIA优化，kernel融合	最高吞吐
DeepSpeed-FastGen	Splitwise架构	DeepSpeed生态

4.2.2 批量生成优化

问题：RL训练中，需要生成大量response，每个prompt生成多个response。

优化点：

连续批处理：动态添加/移除请求，提高GPU利用率
前缀缓存：相同system prompt的请求共享KV Cache
异步生成：生成和训练并行，不阻塞训练
批量采样：同一prompt生成多个response，复用prefill

关键数据（典型RL训练）：

每个prompt生成K=64个response
生成阶段占总时间50-80%
推理引擎吞吐直接影响训练速度

4.3 训练优化

4.3.1 RL算法选择

算法	需要Value模型	需要RM	复杂度	内存	代表
PPO	✓	✓	高	高	InstructGPT
GRPO	✗	✓	中	中	DeepSeek-R1
ReMax	✗	✓	低	低	-
RLOO	✗	✓	低	低	-
DPO	✗	✗	低	低	-
SimPO	✗	✗	低	低	-
KTO	✗	✗	低	低	-
SPIN	✗	✗	低	低	自博弈

PPO vs GRPO 内存对比（70B模型，BF16）：

PPO: Policy + Reference + Reward + Value = 4个模型
     70B × 4 × 2 bytes = 560 GB (仅参数)
     + 优化器状态 × 2 = 560 × 12 = 6.7 TB

GRPO: Policy + Reference + Reward = 3个模型（无Value）
      70B × 3 × 2 bytes = 420 GB (仅参数)
      节省 25% 内存

4.3.2 梯度优化

梯度累积：RL训练中batch size通常较大，需要梯度累积。

实际batch_size = micro_batch × gradient_accumulation × num_gpus
例如: 1 × 8 × 64 = 512

混合精度：

BF16训练（推荐，比FP16更稳定）
FP32梯度累积（防止精度损失）
FP32 master weights

4.4 内存优化

4.4.1 多模型内存管理

RL训练需要同时加载多个模型：

模型	用途	是否可卸载	是否可量化
Policy	当前训练模型	否	否
Reference	KL约束	可offload	可量化
Reward	奖励计算	可offload	可量化
Value	基线估计(PPO)	可offload	可量化

优化策略：

Reference模型offload：训练时offload到CPU，生成时加载回GPU
Reward模型量化：INT8/INT4量化减少内存
LoRA训练：仅训练adapter，base模型共享
梯度检查点：减少激活内存

4.4.2 内存估算公式

模型参数内存 = num_params × bytes_per_param
优化器状态 = num_params × 12 (Adam: 2 states + momentum)
梯度 = num_params × bytes_per_param
激活 = batch_size × seq_len × hidden_dim × num_layers × bytes

总内存 ≈ (num_params × 14 + 激活) × num_models

4.5 采样优化

4.5.1 经验回放

问题：在线RL需要不断生成新数据，生成成本高。

优化点：

优先经验回放：高价值样本优先重用（Ape-X, R2D2）
离线数据混合：在线数据+离线数据混合训练
数据过滤：过滤低质量/重复样本
课程学习：从简单到复杂逐步增加难度

4.5.2 数据效率

关键指标：

每个token的有效训练信号
生成token与训练token的比例
数据重复利用率

优化方法：

多次复用：同一batch多次梯度更新
重要性采样：off-policy修正
优势估计优化：GAE, V-trace

4.6 环境加速（游戏/机器人RL）

4.6.1 GPU加速环境

传统方式：CPU运行环境，GPU运行模型，大量数据传输开销。

GPU加速：环境和模型都在GPU上运行。

环境	加速方式	吞吐提升
Isaac Gym	GPU物理仿真	100-1000x
Brax	JAX可微分	100-1000x
EnvPool	C++批量	10-50x
Atari (GPU)	GPU模拟器	50-100x

4.6.2 批量环境

核心思想：同时运行多个环境实例，批量处理。

传统：env1 → obs1 → model → act1 → env1 → ...
批量：[env1,...,envN] → [obs1,...,obsN] → model → [act1,...,actN] → ...

关键数据：

典型批量大小：64-2048个环境
GPU利用率提升：10-100x

4.7 通信优化

4.7.1 模型同步

问题：多Actor需要同步最新模型参数。

方案：

同步更新：所有Actor等待Learner更新后同步（简单但慢）
异步更新：Actor独立更新，使用stale gradients（快但不稳定）
半同步更新：NVIDIA A-SGD等，容忍一定延迟

4.7.2 经验传输

问题：Actor生成的大量经验需要传输到Learner。

优化：

压缩：量化observations
异步传输：边生成边传输
共享内存：同一节点共享内存，避免序列化

4.8 鲁棒性

4.8.1 训练稳定性

RL训练特有不稳定因素：

奖励hacking：模型找到获取高奖励的捷径
KL散度爆炸：策略偏离参考模型太远
优势估计偏差：Value模型不准
梯度爆炸：长序列+高方差梯度

防护措施：

KL惩罚：限制策略偏离
梯度裁剪：防止梯度爆炸
奖励归一化：稳定训练
早停：基于KL散度或奖励饱和

4.8.2 容错

大规模RL训练的容错挑战：

Actor数量多（1000+），故障率高
经验缓冲需要持久化
训练状态需要checkpoint

方案：

弹性Actor：Actor故障不影响Learner
分布式经验缓冲：持久化经验数据
异步checkpoint：不阻塞训练

五、RL Infra 关键技术点总结

5.1 生成阶段优化（最关键）

优化点                          效果
─────────────────────────────────────────
推理引擎(vLLM/SGLang)          2-5x 吞吐提升
连续批处理                      2-3x GPU利用率
前缀缓存                       30-50% 推理加速
异步生成（不阻塞训练）           50% 总时间节省
批量采样（同一prompt多response）  经验效率提升

5.2 内存优化

优化点                          效果
─────────────────────────────────────────
GRPO替代PPO（去掉Value模型）    节省25%内存
Reference模型offload到CPU       节省1个模型内存
Reward模型INT8量化              节省50%Reward内存
LoRA训练                       节省90%+优化器状态
梯度检查点                      节省60-80%激活内存

5.3 算法优化（减少Infra负担）

优化点                          效果
─────────────────────────────────────────
GRPO（无Critic）               去掉1个模型
DPO/SimPO（无RM无RL）          去掉2个模型
RLOO（REINFORCE）              简化优势估计
多次复用数据                    减少生成次数
离线数据混合                    减少在线生成需求

5.4 通信优化

优化点                          效果
─────────────────────────────────────────
模型同步策略（异步/半同步）      减少等待时间
经验压缩（量化observation）      减少传输带宽
共享内存（同节点）              避免序列化开销
流水线通信                      隐藏通信延迟

六、推荐阅读路径

入门（理解RL Infra核心问题）

InstructGPT - RLHF完整流程
IMPALA - Actor-Learner架构
DeepSpeed-Chat - 工业级RLHF系统
GRPO (DeepSeekMath) - 简化RL算法

进阶（系统优化）

OpenRLHF - 开源RLHF框架
ReaLHF - 动态重并行
SEED RL - GPU集中推理
Sample Factory - 高吞吐异步RL
vLLM - 高效推理引擎
SGLang - 结构化生成

前沿（最新进展）

DeepSeek-R1 - GRPO大规模应用
ReMax - 简化RL训练
RLOO - REINFORCE变体
SimPO - 无参考偏好优化
Rule-Based Rewards - 规则奖励

按技术栈

DeepSpeed系：DeepSpeed-Chat → ZeRO → DeepSpeed-FastGen
vLLM系：vLLM → SGLang → TensorRT-LLM
开源框架：OpenRLHF → TRL → alignment-handbook
简化算法：PPO → GRPO → ReMax → RLOO → DPO → SimPO
环境加速：Isaac Gym → Brax → EnvPool