中午和相关方一起吃饭, 聊着聊着对方忽然问了我一个问题: 数据库和 RL, 你觉得哪个更复杂一点? 我当时几乎没犹豫: 数据库, 而且不是一个量级. 就我接触过的两个角度 —— RL 训练和推理引擎 —— 来说, 链路的深度、组件的数量、状态的变更与维护, 都远远低于一个数据库系统.
这个回答其实有点反直觉. 一个 RL step 横跨 dataset、rollout 引擎、reward、训练引擎好几个系统, 上千张卡, TP/PP/CP/DP 一堆并行策略叠在一起, 光是把 job 拉起来就够折腾了; 而数据库那边, 就算是分布式的那一套, 一台笔记本也能拉起来调. 所以这篇文章想把我的理由讲清楚: 我拿什么衡量复杂度, 以及我判断这件事的依据 —— 一个 RL step 里数据到底怎么流.
我拿什么衡量一个系统的复杂度
先摊开我的度量标准: 不是代码行数, 不是组件数量, 而是有多少状态需要在并发和故障边界上被维护正确. 我在数据库那几年, 印象最深的 bug 基本都是这个形状的:
- 两阶段事务与 WAL 回收. 一个 prepared 事务没结束, 它对应的 prepare xlog record 所在文件就不能删. GP 用一个 hash table 维护所有尚未结束的 prepared 事务的 xid 和 LSN, 取其中最老的 LSN 作为 WAL 回收的水位线. 水位线一旦算错, mirror 侧的 WAL 就会无限堆积. 见 消失的两阶段事务.
- GP 判定分布式事务先后次序的逻辑, 在 21 亿次事务之后会构造出错误的分布式快照, 于是之前的事务插入的数据对之后的事务不再可见. 见 21亿次事务之后…. 注意这个 bug 的触发条件是”跑够久”, 不是”某个请求写错了”.
- 社区修了上面那个 bug 之后, 我重新读了一遍修复代码, 发现修复姿势本身也有问题: 它改变了
CreateCheckPoint()获取两阶段事务信息的时机, 在 recovery 和 timeline 切换的边界上可能导致不一致. 见 两阶段事务一个问题.
这几个例子的共同形状很清楚: 状态跨事务、跨会话、跨进程、跨崩溃重启存活; 不变量被并发破坏一次, 代价是数据错或数据丢; 而且往往在很久以后才浮出来 —— 21 亿次事务之后.
所以我一直坚持一个判断: 把状态的语义定义清楚, 把状态维护对, 整个系统的复杂度就能降一个数量级. 我上一篇 一道有现实意义的代码笔试题 讲的其实就是这件事 —— 那道题从头到尾只有一个难点, 就是”一个时间戳表达不了当前有几个并发持有者”这一个状态语义问题; 语义一旦说清, 伪代码就是照着不变量抄.
我接手一个新系统的第一件事
既然复杂度都在状态和链路上, 那我接手一个新系统时, 第一件事就固定是同一件: 不先读论文, 不先读 config, 而是把一条数据从入口到出口每一跳的形态、owner、生命周期完整画出来. 每个字段是谁写进去的、被谁读掉、shape 是什么、什么时候消失 —— 这张图画完, 这个系统基本就归我了.
我之前跟同事开过一个玩笑: 熟悉我们 RL 框架的一定没有我了解 SGLang 深, 熟悉 SGLang 的一定没有我了解 RL 框架深, 两个都了解的一定没有我了解 Mooncake 深.
玩笑背后是个挺认真的判断: 真正的优化机会几乎都长在链路的接缝处, 而接缝恰好是没有人完整拥有的地方. 每个组件的 owner 都把自己那段做到了局部最优, 于是剩下的空间全在”上一跳交给下一跳的时候, 到底交了什么”上面. 靠着对这条链路的掌控, 我把 DSA R3 的数据流从 GB 粒度降到了 KB 级别, 让一个原本被数据通路压死的特性以近乎 zero overhead 跑了起来.
RL, 近乎无状态
回到 RL. 整体形状大概是这样, 注意它是一个环:
dataset
-> 派发: 按在途水位持续滴灌轨迹, 派出去时打上当前权重版本号
-> rollout: 一条轨迹跑完就单独回收, 不等同批
-> 就绪缓冲: 同一个 prompt 的 n 条都回来了, 这个样本才算就绪
-> 凑 batch: 从就绪样本里凑一个 global batch, 但要先算版本差
-> reward -> advantage -> (old log prob) -> training -> optimizer.step()
-> 权重同步: 参数推回推理引擎, 版本号 +1, 开一个新的缓冲
-> 回到"派发"
先分清三个身份: env group, trajectory, leaf trajectory
后面所有 shape 里的 B 都出自这里, 一半的退化分支也出自这里, 所以先把它定下来.
- dataset 里的一行是一个 prompt. 在框架视角里它是一个 env group —— 一组要在同一个环境设定下反复采样的东西.
- 一个 env group 按
n展开成 n 条 trajectory. 同组的 n 条共享一个uid, 组内再用traj_uid互相区分. - 一条 trajectory 是一串多轮交互: 每一轮一次 LLM 调用, 一个 request/response pair 是一个节点. 如果中途出现分叉, 也就是同一个前缀往下走出了多条路, 这条 trajectory 会被展平成多条 leaf trajectory; 共享前缀的那些节点在后面的 leaf 里 loss mask 会被置 0, 免得同一段 response 被训练好几次.
所以最终喂给训练的 batch 里, 一行是一条 leaf trajectory. 于是 B 既不等于 prompt 数, 也不等于 prompt 数乘 n; 同一个 uid 下到底有几行, 是运行时才知道的. 这件事不搞清, 后面每一步都会错: reward 该挂在哪一行, baseline 该拿哪些行算, mask 该置 0 还是置 1, 全都依赖”一行是什么”这个定义. 我读一个新链路时, 第一个要确认的就是它.
派发: 不是取一个 batch, 是喂一个水位
同步 RL 的心智模型是: 取一个 batch, 全部送去 rollout, 全部回来, 训练. 异步下不是这样. driver 维护的是一个在途轨迹的目标水位: 每次有轨迹跑完、或者某个推理实例空出容量, 就再派一批出去, 让在途数量往目标水位上靠. 目的很直接 —— 推理侧一刻都不要闲着. 几个容易想错的点:
- 派发的单位是 trajectory, 不是 env group. 同一个 prompt 展开的 n 条, 完全可能被派到不同的推理实例上各自独立地跑.
- 每条轨迹在注册进就绪缓冲的那一刻, 会打上当时的权重版本号. 注意是注册那一刻, 不一定等于它实际 rollout 用的版本 —— 注册之后如果又同步过一次权重, 这条轨迹后半段用的其实已经是新参数. 这个偏差后面算版本差时会被一起带进去.
- 水位是自适应的, 而且双向: 就绪样本堆了一堆却凑不出 batch, 说明派太多了, 下调; 训练侧在等 batch, 说明供给不够, 上调.
rollout: 一条轨迹怎么变成一行 token
推理侧就是普通的 HTTP 调用, 一轮一个请求. 有意思的是回来之后怎么摊平成一行. 一条三轮的轨迹长这样:
input_ids: [S, U1, A1, U2, A2, U3, A3, PAD ...]
responses: [ U1, A1, U2, A2, U3, A3, PAD ...] # 就是 input_ids[1:]
response_mask: [ 0, 1, 0, 1, 0, 1, 0 ...] # 与 responses 同形, 只在 assistant 段为 1
rollout_log_probs: [ 0, lp, 0, lp, 0, lp, 0 ...]
S 是 system prompt, Un 是第 n 轮的 user 或 tool 消息, An 是第 n 轮模型自己的输出. 只有 An 那些位置是模型生成的, 所以也只有它们参与梯度. 最后是收口: 把这条变长序列 pad 或 truncate 到定长, 顺手算出 position_ids —— 纯文本是 [B, S], 多模态走 mRoPE 是 [B, 3, S]. 多模态还有一条自己的支线 (视频抽帧、grid 尺寸、把多张图的 patch 摊平 concat 成一个大张量再靠 grid 还原分组), 但那条支线对本文的主题没有增量, 就不展开了.
凑 batch: 版本差决定谁能上车
一条轨迹跑完, 它自己变成 ready; 一个 uid 下的 n 条都 ready 了, 这个样本才算 ready. 然后就该凑 global batch 了 —— 但”就绪数量够了”并不等于”能凑”. 因为凑完这一批就要训练, 训练完权重版本就往前推一格, 那些还在 rollout 中、版本已经偏老的轨迹就可能直接超出允许的版本差, 白跑一场. 所以凑之前要先 dry-run 一遍. 这里的判断是先不急着凑, 宁可等一等, 让那批更老的在途轨迹先落地, 优先把旧样本消费掉. 样本状态的完整流转是四步:
running -> 已派发, 还在推理
ready -> 这一组 n 条都回来了
consumed -> 已经被放进某个 global batch
finished -> 那个 global batch 训完了 (被丢弃的样本也算在这里)
版本差的定义是 batch_id - version_tag, 也就是这个样本实际参与训练的那一步, 减去它注册时的版本号. 还有一个容易忽略的放大器: 训练队列里积压的 batch 越多, 一个新入队的样本从入队到真正被训练之间经历的模型更新次数就越多. 也就是说 staleness 会被队列深度放大, 所以队列深度本身就是个需要盯的指标, 而不只是个缓冲.
reward: 一个标量, 挂在最后一个 token 上
reward 这一跳返回的东西比想象中简单. 结构上是: 每条 leaf trajectory 生成一个待评的 job, 按任务类型分组批量去算, 算完按 batch 内的下标归位. 背后是规则判题、模型打分还是一整套评测服务, 对链路来说没有区别 —— 出来的就是每条一个标量 r_i，存成 token_level_rewards，只放在这条轨迹最后一个有效 token 的位置上，其余全 0。
advantage: 从 token 维回到标量, 再回到 token 维
以 GRPO 为例, 这一跳的形状变换是个来回:
s_i = Σ_t token_level_rewards[i][t] # [B, R] -> [B]
ŝ_i = (s_i - μ_group) / (σ_group + ε) # 按 uid 分组归一化
A[i][t] = ŝ_i # [B] -> [B, R], tile 回 token 维
A[i][t] = A[i][t] * response_mask[i][t] # 最后乘 mask
先压回标量, 是因为 baseline 只在轨迹级别有意义; 再 tile 回 token 维, 是因为 loss 是 token 级的. 算法那一侧的理解我写在 对 PPO-clip/penalty 一种理解, 更基础的部分在 Reinforcement-Learning 学习笔记, 这里不重复.
training: 从一个 batch 到上千个 rank
到这里字段齐了. 一个 global batch 长这样:
prompts: [B, 1] 上面那个右 pad 模式下只剩第一个 token
responses: [B, R] R = S - 1
input_ids: [B, S]
attention_mask: [B, S]
position_ids: [B, S] 多模态 mRoPE 下 [B, 3, S]
response_mask: [B, R]
rollout_log_probs: [B, R]
token_level_scores: [B, R]
token_level_rewards: [B, R]
advantages: [B, R]
returns: [B, R]
is_correct: [B]
剩下的事情全是切分和搬运, 一共七步.
-
pad 到能整除. 两个约束叠在一起: 数据要按 DP 切, 所以
B得能被dp_size整除; 每个 DP shard 内部还要切成固定份数的 mini-batch, 所以B还得能被dp_size × mini-batch 数整除. -
按 token 数重排. 按每条轨迹的有效 token 数重排整个 batch, 让后面切给各个 rank 的连续区间 token 总量尽量接近. 重排用的粒度比
dp_size更细, 这样后续按 DP 连续切时, 每个 rank 拿到的都是若干个已经均衡过的小段. 不重排的话, 长 response 一旦挤在同一个 rank 上, 其他 rank 就得干等它. -
切给 DP. 数据维只按 DP 切, 同一份 shard 复制给这个 DP 组内所有 TP/PP/CP rank —— 因为 Megatron 的 forward 需要这些 rank 全都参与同一批数据.
-
几百 MB 的张量到底怎么搬过去. 这一步现在很少有人讲, 但它是真实开销所在. 这么大的 batch 不走 RPC 序列化, 而是先写进一块共享内存池, RPC 里只带 key; worker 侧按 key 取回, 一次 RPC 结束再把这批临时对象删掉. 这套池子在 RL 下的演进我写在 RL 下 Mooncake Store 演进分析. 顺带一个真实代价: 这些引用的生命周期实际上覆盖了整次 actor update, 而不是”读完就放” —— 因为拿到的是零拷贝视图, 数据直接进了后面的训练容器.
-
切 micro-batch. mini-batch 内部再按有效 token 数和每卡 token 预算做 packing. 所以每个 micro-batch 的样本数
b_i是不相等的 —— 短序列多装几条, 长序列少装几条, 只保证 token 数接近, 且Σ b_i等于 mini-batch size. -
forward 之后怎么取值. 模型输出要按
[:, -R-1:-1]切出 response 段的 log prob 和 entropy. 这个-1偏移不是随手写的: 位置t的输出预测的是位置t+1的 token, 错一位就是在训下一个 token, 而 loss 依然收敛得像那么回事. -
loss 归一化. 第 5 步那个”样本数不等”意味着不能再按 micro-batch 的样本数归一, 否则同一批数据换个切法就得出不同量级的梯度. 处理办法是把 micro-batch 数和 mini-batch size 显式补回去, 让最终梯度等价于对整个 mini-batch 做一次 mean. 一句话: 切分方式变了, 数学结果不能变. 这一条一旦成立, 前面六步就全部退化成性能问题, 不再是正确性问题.
权重同步: 唯一一次反向流动
前面所有跳都是 dataset 往训练引擎方向流. 只有这一跳是反向的: 参数从训练引擎流回 rollout 引擎. 训练侧把参数摊平成 bucket 广播出去; rollout 侧接的动作比想象中细致:
pause_generation(retract) # 暂停调度, 并把在跑的请求撤回等待队列
# 已生成的 token 保留, 客户端视角只是卡住
release kv cache # 释放 KV cache 占的物理显存, 但保留虚拟地址
update_weights(weight_version=...)
resume kv cache # 重新占回显存
continue_generation # 撤回的旧请求和期间进来的新请求继续排队被调度
注意这些动作的性质: 全是资源腾挪加请求重排, 没有一条是在维护跨 step 的语义状态. 而且它们全都是幂等或可重来的 —— 最坏情况是这一次同步失败, 重试一次. 同步完成之后, 最新版本号加一, 开一个新的样本缓冲 —— 于是回到前面那节「派发」. 整个 RL 训练就是这个环在不停地转.
后记: RL 其实不简单
写完上面这些, 我得给上一篇那个标题打个补丁.「RL 其实很简单」说的是这条链路的骨架: 状态少, 每一跳基本是纯函数, 跨 step 活着的东西一只手数得完. 但骨架简单不代表每一跳里面简单. 恰恰相反, 上面每一节里我一句话带过的地方, 单独拎出来都够写一篇, 而且是我短期内写不完的那种:
- 轨迹树. 一条 trajectory 怎么分叉成多条 leaf, 共享前缀怎么去重, 拆成多行之后 loss 的分母该怎么算. 这里改一个系数, 梯度量级就变了.
- 派发与采样. 在途水位定多少才不浪费, 要不要 oversampling, 要不要按历史通过率动态筛 prompt; 筛掉之后组内还剩几条、baseline 还成不成立.
- rollout 引擎本身. 上面我只写了”普通的 HTTP 调用, 一轮一个请求”, 但推理引擎自己就是一整个世界: 调度、KV cache、prefix cache、投机采样、大 EP、显存管理. agentic 之后还要叠上工具调用、环境交互、超时与失败重试.
- off-policy 到什么程度还能收敛. 版本差那一节我只讲了记账怎么做, 一个字没讲那个数该设成几. 而这恰好是算法和系统的交界处: staleness 换来的吞吐到底值不值.
- reward. 我用了一整节讲它怎么”摆”, 一句没讲它怎么”算”. 判题准确率、reward hacking、长度偏置、多任务怎么配比 —— 这些直接决定训出来的东西像不像人话.
- advantage 的每一个选项. 组内均值还是 token 加权, std 到底该不该除, 要不要按长度归一化. 每一个都在改训练动力学, 而不只是改一个数. 我在 对 PPO-clip/penalty 一种理解 里也只是摸到了这个话题的门口.
- 训推一致性. 同一条序列, 两个引擎算不出同一个 log prob. 为什么、差多少、要不要修、修到什么程度, 这个能单独写好几篇.
- 切分与显存. dynamic packing 和 CP/PP 怎么互相影响, loss 归一化的等价性怎么证, optimizer offload 的时序怎么排, 显存峰值到底卡在哪一步.
- 权重同步. 参数怎么摊平成 bucket, 同步期间在飞的请求怎么处理, 新版本什么时候对推理侧才算真正可见.
所以这两个标题并不矛盾, 它们说的是两件事: 骨架简单, 是我敢一个人把整条链路端到端拿下来的前提; 每一跳都不简单, 是这件事拿下来之后才配开始的部分. 顺序不能反 —— 先挑一跳往下钻, 你连自己钻的这一跳在整条链路里承担什么都说不清, 更别说判断它值不值得钻.
