作者 | 谭三爷@知乎转自 |
原文链接：https://zhuanlan.zhihu.com/p/1959094134032830632

点击下方卡片，关注“大模型之心Tech”公众号

本文只做学术分享，如有侵权，联系删文，欢迎添加小助理微信AIDriver004做进一步咨询

在过去的几个月里Simon^[1]和我一起开发了一套LLM后训练框架^[2]。这对infra同学大概很trivial，但作为一名算法，每一个模块的开发都让我领悟到强化学习中的细节，令我受益颇丰。在此向你分享一些我的乐趣。

强化学习框架需要做什么？

在LLM强化学习框架成熟之前，DPO曾盛极一时。当时非常火热的topic是DPO和PPO孰优孰劣？我的答案是：DPO是牺牲算法换取infra简易性的产物。

• 算法上，DPO假设奖励服从Bradley-Terry模型，并且只适用于单轮，这使DPO难以适用于binary奖励的推理和多轮的agentic问题。
• infra上，DPO分离了推理和训练，而这两个系统各自分别成熟，这使得DPO的实践成本很低。

所以，在我看来，强化学习框架实际上是训练和推理引擎的粘合剂。策略梯度的公式是

这其实可以看作在每一个token上施加权重的SFT

强化学习中的训练和推理引擎各自并没有本质上的新颖之处。相较于成熟的SFT和DPO，强化学习框架所做的，是针对强化学习的on-policy特性

• 管理训练和推理引擎的资源，如分卡和offloading
• 将训练引擎的参数同步到推理引擎

其他的部分，如训练和推理引擎内部的分布式和checkpointing，只需沿用先前的成熟方案。

分卡策略

强化学习的每次迭代分为三个阶段：

• 推理引擎的rollout
• 训练引擎的推理，包括Actor和RefActor计算logps、Critic计算values和RM计算rewards
• 训练引擎的训练，包括Actor和Critic的更新

在标准的sync强化学习中，这三个阶段相互依赖，需要顺序执行；每个阶段内的步骤相互独立，可以并行执行。如果阶段内的步骤顺序执行，则对应着将所有worker colocate到所有GPU上；如果阶段内的步骤并行执行，则对应着将不同worker分配到不同GPU上。不过，算法上REINFORCE、基于规则的奖励和去掉KL正则的兴起，消除了并行执行的需求。

在async强化学习中，推理引擎和训练引擎分离，各自持续工作。训练引擎每一次迭代后，从推理引擎拉取数据并向推理引擎推送参数。在我看来，这是打破强化学习on-policy特性，牺牲算法稳定性换取infra效率的产物。

负载均衡

训练引擎的负载均衡

传统的做法是为每个DP分配相同数量, i.e., batch_size_per_gpu,的sequence，但每个sequence的长度不一，分配较短sequence的DP需要等待较长的DP。在这里，我复用了veRL的get_seqlen_balanced_partitions函数，其作用是根据各sequence的长度将sequences分配到k_partitions个partitions里，使各partition的长度均衡。问题便转化为了求解合适的k_partitions。这里的约束有二

• k_partitions的数量需要是dp_size的倍数，因为partitions要均分到各个DP上
• 各partition的长度不能超过max_length_per_gpu

我的做法是，从不小于math.ceil(sum(seqlen_list) / max_length_per_gpu)的第一个dp_size倍数开始。因为如k_partitions小于该值，则各partition的长度均值都大于max_length_per_gpu，他们的长度不可能都不大于max_length_per_gpu。

从该值开始，调用get_seqlen_balanced_partitions，判断是否存在partition的长度超过max_length_per_gpu。如存在，则将k_partitions增加dp_size, i.e., 满足约束一的下一个值；如不超过，则获得满足要求的分割。

训练引擎负载均衡带来的一个有趣问题是损失放缩。传统的损失是在各个sequence之间求平均。当各个DP sequence的数量相同时，每个DP只需要计算并反传本DP上各sequence损失的均值，再在optimizer step时对各个DP的梯度求平均即可。在均衡各个DP的负载时，每个DP上sequence的数量不同。在这里，我在每个DP上将各个sequence的损失求和，并除以各个DP的sequence数量之和。

如此，每个DP计算并反传的是平均损失的分量，各DP上的损失和梯度应该被求和。由于FSDP将在optimizer step时将对各个DP的梯度求平均，我在反传前再将各个DP上的损失乘以dp_size。如此，将和FSDP对各个DP梯度求平均的操作融合为对各个DP梯度求和。

推理引擎的负载均衡

传统的做法是为每个DP分配相同数量的prompt，但每个prompt生成response的长度不一，生成较短response的DP需要等待较长的DP，并且response的长度难以事先预测。

在这里，我复用了slime的做法，不预先将prompt分配到各个DP，而是在rank 0上起router将请求转发到各个DP。如此，router将根据各个server的空闲情况转发请求，降低response长度不均导致的GPU idle。这种做法还带来了三点附加的好处

• 所有的请求都从rank 0发出，因此rank 0上的进度条可以准确反应rollout的进度
• 只有rank 0和环境进行交互，不需要将环境分配给各个DP，便于实现环境的负载均衡
• 可以使用router的高级特性，如PD分离，提高推理效率

环境的负载均衡

在代码和搜索等场景中，受限于硬件和API的并发，只有有限数量, let say 个, 环境可供交互，以生成出, let say 个, 轨迹。当时，轨迹需要排队和环境进行交互。一种直接的做法是将第个轨迹分配到第个环境上，但这样的问题是，每个轨迹的生成时间不一，当一些轨迹排队的时候另一些环境可能已经空闲。

在这里，我不预先将轨迹分配到环境上，而是使用asyncio.Semaphore控制至多个轨迹进行生成，每个轨迹在开始生成时被分配一个空闲环境，结束时释放占用环境。

Tokenization

在我开发时观察到了一个有趣的现象：

tokenize(prompt) + tokenize(response)并不一定等于tokenize(prompt + response)。

在推理时，是将tokenize(prompt)喂入推理引擎，再decode(response_tokens)得到response。此时若拼接prompt和response文本，即tokenize(prompt + response)喂入训练引擎，输入训练和推理引擎的token可能不同，破坏强化学习on-policy的特性。

在多轮的agentic场景中，亦不能在下一轮将tokenize(prompt + response + tool_result)喂入推理引擎。

正确的做法是总是拼接token而不是文本：对prompt和tool_result单独tokenize，将拼接好的tokens喂入推理引擎，并直接从推理引擎获取生成的response_tokens。

在多轮的agentic场景中tool_result常常拼接在prompt和response之后。这一方面是可以在推理时复用KV cache，一方面是可以在训练时将多轮放在一个sequence中，但强化学习并不要求状态是的prefix。

一般地，给定，环境转移到是的前缀只是一种可以将两轮放在一个sequence中的特殊情况。

因此，我的做法是，观察到后，判断是否是其前缀。

• 如果是，则继续在后拼接；
• 如果不是，则。

非concat的状态转移将导致一条轨迹可能对应多条sequence，相比每条轨迹对应单条sequence，logps和values的计算和Actor和Critic的更新将不受影响，唯一影响的是advantage的计算。注意advantage的计算对象是动作，即每一个被生成的token。因此，应将多条sequence的response_tokens抽出，依顺序拼接，再依GAE的公式和计算advantages。