首页 / 资讯中心 / 文章详情

大模型强化学习中的推理基础设施:从采样到训练的工程实践

大模型强化学习中的推理基础设施:从采样到训练的工程实践 ★ FEATURED ARTICLE
1. 先别急着谈“训练”“推理infra”才是RL真正的主舞台不管你是做对话系统还是做模型后训练只要真正碰过一轮完整的大模型强化学习都会发现一件反直觉的事训练那部分的计算量往往只占一小半真正吃资源、吃GPU、吃排障精力的是推理infra。RL里的采样环节每轮都要让当前策略模型批量生成成千上万条输出这个生成量要跑数十甚至上百轮累计起来的推理负担比平时对外服务的API还要重得多。我最早接触这个问题的场景是某次做基于规则奖励的强化学习实验。当时我满脑子想的都是怎么调整策略更新参数、怎么写奖励函数结果真正跑起来之后卡瓶颈的却是采样速度太慢actor模型每更新一次要等rollout生成几个小时GPU利用率还上不去。后来我才意识到RL infra本质上就是把“推理服务”反着用服务场景要求低延迟、高并发、结果稳定RL场景要求高吞吐、长序列、结果多样且随时要回传训练器。两者的底层能力高度重叠但调度逻辑、批处理策略、显存管理、数据流设计全都不一样。这篇文章要说的事就是我从纯推理服务的infra一路改造成能承载RL训练的整套体系时的认知和方法。整理出来既是给自己留个存档也是想提醒那些刚入局的人想把RL跑起来、跑得稳与其天天盯着策略更新公式不如先把自己的推理infra这块地基夯实。1.1 为什么RL里的“推理”和“服务”不完全是一回事很多人一听到“用推理引擎去给RL做采样”第一反应是这不就是调个API或者部署个推理服务吗实际远没那么简单。对外提供服务的推理核心指标一般看延迟和并发用户发一个请求系统尽快吐出来一个答案用的往往是短prompt短的生成长度。RL里的rollout生成则完全反过来一次要喂进去几千条prompt每一条都要生成几百到几千token并且所有样本的生成结果都不需要立即返回给用户而是统一汇集到一个池子里等训练器消费。所以服务推理和RL里的推理看起来调用的是同一个模型但实际上对系统行为的要求差别很大这里我列一个对比维度服务推理RL采样推理核心目标低延迟、高并发高吞吐、批量生产输入长度通常较短几十字到几百token为主可能很长尤其是代码/数学场景输出长度几百token以内居多常常需要生成1k~8k甚至更长实时性要求严格p99延迟必须可控不严格可以接受几分钟一批批处理策略动态拼批为主请求随到随处理尽可能把整个rollout集合压成大批次结果去向直接返回给用户回传训练器做奖励计算和策略更新系统瓶颈首字延迟、排队、超时显存、生成吞吐、KV缓存复用率从表格能看出来服务推理更在意“体感”RL推理更在意“产量”。如果直接拿服务场景的推理引擎来跑RL采样很容易出现两个问题引擎为了追求低延迟会优先调度短请求或提前终止结果长序列的采样进度被拖后生成吞吐上不去。采样的返回路径和训练器的数据接口没打通生成结果写盘、解析、回传的链路太长形成事实上的数据阻塞。所以RL infra里说的“推理”并不是换个prompt模板重新调一遍接口而是要把推理引擎重新理解成一个数据生产器它喂给训练器的不是“一次对话的答案”而是“一个完整分布的回传”。1.2 RL infra不等于“训练框架plus”它更像一条流水线现在业内聊RL infra容易陷入两种极端。一种是光聊PPO、GRPO这些算法把infra理解成“把训练跑起来”另一种是只关注推理引擎的优化比如吞吐、显存、调度觉得做好了采样就万事大吉。我的理解是RL infra是一条完整的流水线从采样到奖励到更新再到下一轮采样形成了一个闭环。为了说明白这件事我一般把整条链路拆成四段来看采样段策略模型根据一批状态生成响应。奖励段对生成的响应计算奖励可能是规则奖励、模型奖励也可能是和环境交互的得分。训练段用强化学习算法把奖励回传成梯度更新策略。数据段采样结果、奖励值、模型指标全部回传清洗供下一轮训练使用。推理infra在其中负责的就是采样段以及和数据段交界的那一大部分。以前我做服务推理的时候只要考虑GPU上的显存效率、批大小、并发请求就够了但在RL infra里采样段产生的数据是整个流水线的粮食一旦它断了、慢了、脏了后面的奖励和训练全部要跟着停。所以真正把RL infra搞顺之后你会发现很多精力其实不是花在训练参数上调来调去而是花在解决这样的问题上怎么让采样段和训练段并行跑起来中间通过队列或者共享内存做异步衔接怎么让推理引擎不要因为某一条超长生成堵住整批任务怎么在断点续训时把已经采出来的样本保存下来不至于因为一次OOM全部重来。1.3 适合谁看这篇文章如果你是以下这类角色这篇文章对你的帮助会比较大已经在跑LLM推理服务想把能力延伸到强化学习后训练但不清楚需要改动哪些组件。正在实践RLHF/RLVR一类的训练流程发现瓶颈出现在采样和回传环节。或者刚接手某个RL训练系统被一堆“异步”“队列”“吞吐”“采样”这些词砸晕需要一套用来理解的框架。如果你对这些背景都不太熟也能看我会尽量把底层逻辑讲清楚避免满篇黑话。但前提是你得先明白一个基础概念RL更新的是策略模型的参数而采样是多次前向推理来获取行为数据。前向推理的快慢和稳定程度直接决定了RL迭代的速度上限。2. 推理infra的三大核心能力在RL里分别扮演什么角色既然推理infra是RL的基石我们就要把“基石”落实到具体的技术点上。我在做系统迁移的时候最关注的就是三件事动态批处理、KV缓存管理、请求调度策略。这三个能力分别对应RL里的采样吞吐、长序支持、任务公平/优先级调度。2.1 动态批处理让采样吞吐从“够用”变成“溢出来”动态批处理是推理引擎最核心的优化手段之一。简单说它把不同请求的prefill和decode阶段尽量重叠起来GPU始终被高占用地填满。我们平时在推理服务里用动态批处理是为了让每个请求都能更快地被响应。到了RL里同样的机制发挥的作用完全不同它是为了把整批rollout请求在最短时间内生成完。举个例子假设你做一轮强化学习需要采样1024条回复。如果推理引擎一次只能同时吃16个请求那这1024条回复要分批跑64轮如果动态批处理的调度足够激进每轮能同时容纳256个请求那么四轮就能搞定时间差距是数量级的。这里有一个常被忽视的点动态批处理的“批大小”不是越多越好。批越大显存占用越高而且如果某个请求太长会把其他请求一起拖住。RL场景里生成长度分布往往非常不均有的回答只有100字有的要写到2000字。所以好的RL采样引擎需要在“吞下更多请求”和“别让长尾请求拖垮整批”之间找一个平衡点。最常见的解法是给批处理加一个“最长生成时间预算”或者把长生成请求单独拆出来走一个慢路径。在拆的时候如果KV缓存做得好拆出来不会带来额外存储开销如果KV缓存本身是连续集中保存的那拆批就意味着要重新拷贝显存数据成本非常高。提示如果你的采样吞吐一直上不去先不要急着加GPU卡。先观察一下批处理在长请求出现时是不是把整个batch的生成步长都拉长了如果是优先引入长请求隔离策略。2.2 KV缓存从“省着用”到“用出策略”KV缓存是推理过程中缓存了历史token的Key和Value张量它的作用就是避免每生成一个token都把整个序列重算一遍。KV缓存的大小和序列长度成正比一旦请求变长显存消耗会呈线性上涨。推理服务通常会把KV缓存做成分页式的类似操作系统的分页内存这样多个请求可以共享物理显存避免碎片化。这个设计在RL里更关键因为RL采样天生就是“多输入、长输出”KV缓存的总需求会非常高。我在实际跑RL实验时碰到的典型问题是几轮采样跑完之后显存眼看着要爆但每一条单独看都还有空间。后面排查才发现是因为推理引擎在若干轮采样之间没有清理缓存而且批处理算法倾向于按请求到达顺序来排布KV页导致页表碎片化严重。解决方式一般有两层引擎层开启KV缓存的“按轮次清理”或者“按batch生命周期管理”每处理完一批rollout立刻释放页表。调度层放弃“先来先服务”改成“按预测长度聚类”把预测长度相近的请求放在同一个批次里这样KV分配更整齐显存利用率更高。尤其在做长度奖励、代码生成为主的RL任务时序列长度极不均匀。我测过一组数据同样的1024条采样任务不做长度聚类时显存溢出概率很高做了之后相同显存物理容量下吞吐提升了近一倍。2.3 调度策略RL采样有自己的优先级权重推理服务里常见的调度策略是优先处理“延迟敏感的短请求”其次再照顾长请求。但RL采样不一样。RL里有些请求就是注定要跑很久的你不能让它们一直让路。我的处理方法是给采样任务分优先级第一优先级当前策略更新之后、等待采样的“新鲜数据”。这类数据决定了训练器下一步的输入必须优先保证。第二优先级探索型数据比如让模型生成一些超长回答、低概率回答用于增加样本多样性。第三优先级历史重启或补偿型采样比如前面一批数据因为训练器卡住没消费完需要重新生成。这样做的好处是训练器永远有数据可吃不会出现“训练器干等采样结果”的情况。如果你把所有请求都一视同仁丢给推理引擎那结果往往是短的请求一直在插队长的、需要大量探索的请求被不断延后最后整个RL流程被人为拉长。3. 从推理到RL infra要做的不只是“套壳”很多团队一上来就在开源推理引擎外面包一层宣称自己有了RL infra。实际跑起来发现该卡的地方还是卡。原因很简单RL infra不是推理引擎的皮而是要把引擎和训练器焊成一个能协同的系统。简单说你要在推理引擎和训练框架之间加一套“RL编排层”让采样结果能自动喂给奖励计算和策略更新并把模型参数同步回去。3.1 核心组件rollout worker、reward pipeline、actor sync我理解的RL编排层至少包含下面几个组件rollout worker负责承接一批采样任务。它从任务队列里拉取prompt调用推理引擎做批量生成拿到输出文本之后做基本的格式清洗比如把EOS token去掉、把答案截断到合理长度然后写入统一的样本缓冲区。rollout worker可以跑在单独的GPU上也可以和推理引擎放在同一个实例上但要保证不能因为采样占用太多显存而拖慢后续的模型更新。reward pipeline采样结果不是直接进训练器的。要先算奖励奖励可以来自规则函数比如答案是否包含正确代码片段也可以来自奖励模型通常是一个小的分类/回归模型。奖励管道的设计要尽量做成异步的也就是说rollout worker产出一批结果reward pipeline立刻消费一批不让数据在中间停留。奖励计算的延迟同样会被放大成整个RL流程的瓶颈。我见过不少项目采样很快、训练很快但奖励计算是单线程逐个跑结果整条链路还是卡死。actor同步每隔若干个训练步策略模型更新一次参数。更新之后rollout worker要能加载到新模型参数继续采样。这个同步链路如果做的不好就会出现“训练更新了一版新模型但采样还在拿旧模型生成样本”的情况。正常说这不是大问题RL算法本身允许一定滞后但如果滞后太多会导致策略优势估算失真。我通常会设置一个上限最多允许采样端落后训练端若干个模型版本超过之后优先清空采样队列强制同步最新模型。3.2 数据通路设计别用文件系统当队列最简单的做法是把每条采样结果存成JSON落盘到本地训练器再读盘做训练。数据量小的时候没关系但一旦采样规模达到每轮几万条、几十万条文件系统的小文件读写会彻底拖垮系统。我建议大家思考RL数据通路时至少考虑下面三条路径中的一种共享内存队列适合单机方案几块GPU放同一台机器采样worker生成数据后直接推入内存队列训练进程消费。分布式消息队列适合多机方案采样端在A组机器训练端在B组机器用消息队列做缓冲便于动态扩缩容。列式文件批写如果非要落盘不要一条一条地写而是攒够一批按parquet或其他列式格式整批写入读取时也按批读取减少IO次数。我个人的经验是在单集群规模还不到几十台机器时共享内存队列优先最简单也最稳。到了需要水平扩展的阶段再去上消息队列。过早引入分布式消息系统等于在还没遇到流量之前先给自己加了一堆运维负担。3.3 并行与异步让采样、奖励、训练变成三条互不阻塞的流水线RL任务天然存在“数据依赖”策略更新之后又要重新采样看起来没法并行。但实际上可以拆成两个层次来做异步。第一个层次采样和训练之间的异步。策略更新到第n轮时第n-1轮采样的结果已经到了奖励管道第n-2轮的数据已经训练完毕。相当于流水线上三个环节同时在进行。第二个层次采样内部的异步。一批rollout中有的已经生成完毕有的还在跑。可以把生成完毕的结果提前送去做奖励计算不需要等整批rollout全部跑完再统一走下一个环节。这种异步流水线的设计对系统稳定性要求很高。一旦某一环节出问题比如奖励模型崩了或者某个采样worker显存溢出会导致上游数据堆积。所以在设计阶段一定要为每个环节预留“全排队”状态并加上监控告警队列堆积超过阈值立刻报警而不是等到训练器饿死才发现。4. 实操迁移三个步骤把推理服务改造成RL采样能力这部分我想直接说“怎么做”。以一个原本只提供API推理服务的团队为例假设你手上有推理引擎、有训练框架另外有一批prompt数据想做强化学习我会建议你按下面三个步骤改。4.1 第一步先把推理服务压成“批吞吐优先”在没做任何架构改动前先把推理服务切成一个专门的“批量采样模式”。这个模式要做的事有三件关闭实时请求入口只接收预先生成好的大规模prompt列表。把连续批处理的并发上限调到最高允许请求数量远超显存能真正支撑的量让引擎自己按KV缓存走调度。设置一个较大的“最大生成token长度”容忍度避免在代码或数学场景里因为默认长度限制导致输出被截断。完成这一步你用最少的代码把推理服务变成了一个普通的“批量生成器”。这时采样的吞吐大概率已经比原来快不少因为你不再被单个请求的延迟绑架了。4.2 第二步做数据桥打通采样端和训练端批量采样模式稳定之后接着要解决的就是“采样数据怎么送进训练器”。这时你需要写一个bridge模块。基本流程是推理引擎产出完整结果包括输出文本、token个数、生成时的概率日志如果需要计算KL散度的话。bridge把生成结果封装成统一格式prompt_id, response_text, response_tokens, log_prob_sum, finish_reason等字段。把封装好的记录写进共享内存队列。训练器从队列里拉数据组装成训练batch。注意这里的log_prob_sum非常关键。强化学习算法通常需要计算当前策略生成某条序列的概率和建议策略之间的KL散度。如果在推理引擎处不能导出logprob后面就得重推一遍模型来计算资源直接翻倍。我见过一些团队为了图省事采样的时候不存logprob等到训练阶段重新forward一遍最后整个训练流程的时间硬生生延长一半。这就是典型的“省一时亏一路”。4.3 第三步把训练器接到流水线上并加上断点恢复训练器接到流水线上之后第一个要测的不是效果而是稳定性。怎么测呢我一般先跑一个小的stability smoke test拿50条样本跑10轮强化学习确保每一轮训练器都能在合理时间内拿到新数据不会出现队列空转、不会出现数据重复消费。然后就是断点恢复。RL训练动不动跑几十个小时中间一旦某个GPU崩掉或者显存不够导致节点重启如果没有断点恢复之前采的几十万条样本全部作废。所以在设计阶段至少要做到rollout worker的采样状态定期快照记录已处理到哪个prompt批次。奖励管道的计算结果要和样本一起落盘这样重训时不需要重新算奖励。训练器能够从最新的模型权重和最近一次数据批次恢复而不是从头开始。我就是因为这些细节踩过很深的老坑。有一阵子做数学推理任务强化学习一个实验跑了两天中间某个采样worker OOM结果整个实验从零开始。之后把所有中间状态都加了checkpoint再后面跑同样的实验哪怕半夜出问题第二天起来恢复一下就能继续省下的时间量级完全不一样。5. 一套可参考的RL采样参数配置与监控指标体系光说过程不给配置总觉得像是纸上谈兵。下面我按自己在实际项目中验证过的方案给出一套比较稳的默认配置以及配套的监控指标。你完全可以拿来做起点再按自己的模型规模和显存容量调整。5.1 采样参数从一个保守但稳定的配置开始参数建议值说明每轮采样批次大小256~1024条视显存和训练batch决定太小则策略梯度噪声高太大则单轮时间过长单个生成最大token长度2048或4096代码/数学场景建议4096top_p0.9~1.0探索性强别设太低temperature0.7~1.0RL训练一般在探索阶段略高之后可降温采样worker数量等于推理引擎可并行实例数避免过多worker抢显存训练与采样比例1:1或1:2一个训练步对应一到两轮完整采样避免采样瓶颈累积logprob存储必定开启否则KL散度计算需二次forward这组配置不一定最优但一定比较稳妥。对于那些想做超长代码生成任务的团队我会强烈建议把单个生成最大token长度调高到8192并检查KV缓存页表是否能支持这个极值。否则生成到一半显存耗尽、请求被强制中止采样结果全变成截断垃圾训练效果会非常诡异。5.2 监控指标别只盯loss要盯“机器在干嘛”训练loss固然要看但它只能告诉你模型参数更新得顺不顺看不出infra层面的问题。我自己在RL infra上常用的监控指标分三类产出类指标每轮rollout生成的有效样本数、平均生成token数、平均生成时长、生成失败率。资源类指标推理引擎的GPU显存使用率、KV缓存命中率、动态批处理的平均批大小、队列积压数量。质量类指标奖励均值、奖励方差、策略更新后的KL散度、生成响应长度方差。这三个维度合在一起才能真正定位问题。比如你看到训练loss一路下降但最终评估效果很差那可能就是采样分布出了问题样本多样性不够、长度分布畸形、奖励信号几乎全是噪声。如果只是盯着loss你永远不会发现是rollout采样阶段就已经扼杀了模型的探索空间。我有个实用的经验每轮采样完成后把这一轮的响应长度分布、奖励分布、token利用率画成直方图对比上一轮。一旦发现长度分布突然变得极端或者奖励分布集中到某个极窄区间马上检查prompt构造和reward pipeline早发现问题比精细调模型参数重要得多。5.3 实验对比怎么跑才有效很多人做强化学习实验对比时会出现“看起来涨了但不知道涨在哪”。原因往往是实验之间混入了太多的变量没有把infra层面的差异控制住。我的习惯是先固定推理引擎和采样参数只改策略更新的算法或奖励函数。这样A/B对比时至少能确信性能变化来自算法侧而不是因为这次滚动batch大小不同导致采样质量变化。等确定算法侧稳定后再反过来调infra侧的吞吐参数追求“同效果下更快的实验迭代”。6. 实际踩过的坑关于推理到RL迁移的几个反直觉问题这部分写几个真实碰到过、网上不太容易搜到答案的问题。每个问题我都会给出排查思路和最终解法希望能帮你少走弯路。6.1 奖励函数写好了但模型就是不学东西我之前做一个代码生成强化学习任务奖励函数很明确编译通过给正分不通过给负分。跑了好几千轮模型生成结果却一直原地踏步看不出任何升分趋势。排查后发现问题不在奖励而在采样生成的长度分布太离散一些代码长度只有几十个token一些却冲到几千token。训练器做batch组装时因为长短差异太大pad token占比过高有效计算比例被稀释到极低。解决方法是引入“长度分桶采样”把代码生成长度接近的样本分到同一批。我发现按长度聚类后训练的吞吐和有效性都有非常明显的提升而且奖励曲线更平稳。这个问题的本质是推理infra带来的序列差异被传导到了训练效率上。6.2 采样日志堆积训练端总是拿不到最新数据异步流水线设计好之后出现过一个现象训练端显示队列积压为零但实际训练batch迟迟组不起来。查了很久才发现是rollout worker写队列的时候把“生成完毕”的标记放在文本数据之后训练端要等文本全部写成功才能读而文本解析的速度远慢于队列写入速度。这不是什么高深bug纯粹是数据流设计上的小失误。但教训很实任何环节之间都最好带一个独立的“活性”缓冲区并打上时间戳。一旦上游产出速度快于下游消费速度立刻就能看出来是卡在解析、网络还是训练端加载。6.3 推理引擎版本更新后采样分布悄悄变了这个坑尤其隐蔽。某次我把推理引擎升级了结论是推理速度变快了但同一组prompt、同样的采样参数生成结果的分布却不一样了。因为新版引擎对“提前终止”的判断逻辑有改动或者随机数种子、beam搜索细节变了。在RL实验期间我强烈建议把推理引擎版本、采样参数、随机种子全部锁定并且记录到实验元数据里。否则你换了个引擎版本结果模型效果变了但你会误以为是算法改进带来的。提示分布漂移类的问题可以通过对比“同一prompt集合、同一参数下的生成回答长度分布和奖励分布直方图”来快速发现。两个分布差距过大立刻检查引擎或配置变动。6.4 训着训着显存缓慢上涨最终OOM这个问题的根子通常不在模型本身而在于rollout worker缓存了太多历史样本不释放。更隐蔽的是推理引擎和训练器共用了同一块GPU显存池训练器的激活值逐步增长推理侧KV缓存却保留了历史页表没有回收两者叠加在一起把显存挤爆。排查思路很简单看显存占用曲线是不是呈阶梯状上升而不是脉冲式波动。如果是阶梯上升去查KV缓存池回收策略和历史样本缓冲区释放逻辑。最终我解决问题的方式是在每个rollout批次结束时强制清理KV缓存并且提高历史批次样本的压缩存储密度。7. 最后讲一点自己使用时的真实经验做了一段时间从推理infra到RL infra的迁移后我的体会是强化学习训练的门槛很多时候并不在算法细节而在你能不能稳定、高速地产出采样数据。算法再先进样本不稳定、质量不统一、吞吐跟不上最后都很难收敛或很难复现。我还有一个小技巧每轮实验把rollout采样的“有效token数/GPU时”作为第一优先级指标。它比训练loss更早告诉我们系统有没有滑向低效状态。只要你把这个指标盯住了推理infra上的大部分资源浪费问题都会自己浮出水面。另外真的别在RL已经跑起来之后随手升级推理引擎哪怕它号称“性能提升很多”。在实验闭环里稳定远比性能提升更重要。RL infra这一块说实话还有很多可以细化的方向包括更细粒度的显存复用、基于采样结果质量的反向调度、甚至是把奖励反馈直接接进推理引擎的prefill阶段。我这个阶段踩过的坑比较多先把路径上的坑填好后面跑起来就能更安心地往前走了。
阅读完成 · 觉得有帮助?
咨询建站