[VeRL] AgentLoop源码走读

最近 RL sys 圈子的吴锡斌老师在 verl 上设计了将 rollout 与 tool 调用解耦的 AgentLoop,实现了自由灵活的 mutli-turn RL。在每个 AgentLoop 内部,rollout engine 只对外提供一个 token-in-token-out 的接口,而 tool 调用则通过 ToolAgentLoop 来实现。我个人比较喜欢这样解耦的设计,同时,AgentLoop 的代码结构也比较清晰。我个人学习了一次整个代码后,觉着 AgentLoop 的设计甚是不错,但是 ActorRolloutRefWorker 的历史包袱还是很重。 本文简单分析了 agent loop 的源码,并给出了一些自己的看法。 如果我们把整个 ActorRolloutRefWorker 当做一个 sgl.Engine 的话,AgentLoop 里面包装的两层 AsyncSGLangServer 和 AsyncLLMServerManager。AsyncSGLangServer 相当于在 sgl.Engine 上包装了 fastapi 成了 server,而 AsyncLLMServerManager 是在 server 上包了一层 router 做 load balance,相当于 sglang 的 router。这两层设计都是合理的,主要麻烦的是 ActorRolloutRefWorker,层层调用,最后一共经过 7 个 class 才调到 sgl.Engine,最近 verl 团队也在致力于对这块 worker class 的重构,敬请期待。最后,AgentLoopManager,AgentLoopWorker 和 AgentLoop 这三层,我觉得 AgentLoopWorker 可能未必有必要,其他两层挺合理的。 ...

August 14, 2025 · 15 min · 3051 words · Me

[VeRL] 参数速览

VeRL框架的参数众多,基于当前(2025.8.5)主线分支整理,附带了相关的理解,一些描述不一定完全正确,供学习参考。 Batch Size 参数名称 详细解释 data.train_batch_size 作用:定义了单次训练发送给 Rollout Engine 的样本数量,也即这是在每个 PPO 迭代开始时,从训练数据集中采样的提示 (Prompt)数量。 详细解释:这个值是 RL 训练中的基本样本数量。例如,设置为 1024 意味着在一次迭代中会: 1. 从数据集中随机抽取 1024 个 prompt。 2. 将这 1024 个 prompt 发送给当前的 Rollout Engine 中,从而得到 1024 组完整的 trajectories(prompt, response)。 3. 接下来,这 1024 个 trajectories 进行经验计算(make experience),后续用于 Actor 和 Critic 模型的更新。 影响与权衡:影响总共训练的样本量。 data.val_batch_size (Deprecated) 作用:在 Validation 阶段使用的批次大小。 详细解释:这与 train_batch_size 类似,但仅用于评估模型性能,不参与训练。如果设置为 null,会使用验证集的大小作为默认值。Note: 已经deprecated,推荐设置为 null。此时,整个 validation dataset 一次性发给 SGLang engines,自行进行内存管理。 actor_rollout_ref.actor.ppo_mini_batch_size critic.ppo_mini_batch_size 作用:定义了 PPO 训练更新中的 mini-batch 大小。 详细解释:data.train_batch_size 收集到的全部经验数据将被分割成多个 mini-batch,每块的大小就是 ppo_mini_batch_size。模型每处理完一个 mini-batch,才会进行一次参数更新。 例如,如果 train_batch_size = 1024,ppo_mini_batch_size = 256,那么在一个 PPO Epoch 中,模型会进行 1024 / 256 = 4 次参数更新。 影响与权衡:增大 mini-batch,单次更新的梯度更稳定,但更新频率更低,更新次数减少。 actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu critic.ppo_micro_batch_size_per_gpu 作用:定义了在单个 GPU 上进行一次 forward/backward 的数据大小。 详细解释:这是实现梯度累积的核心参数。mini-batch 会被再次切分为若干个 micro-batch。例如,在单卡上,ppo_mini_batch_size = 256,ppo_micro_batch_size_per_gpu = 32,那么梯度累积的步数就是 256 / 32 = 8。这意味着模型会运行 8 次 forward 得到 loss,然后 backward 的到 gradient。每次处理 32 个样本,直到累积完整个 mini-batch 计算出的梯度。此时,使用累积的总梯度,对模型参数进行一次更新(optimizer.step())。这个值必须根据显存大小来严格调整,是防止 OOM 的关键。 影响与权衡:增大此值,减少了梯度累积的次数,可以提高训练的吞吐量,增大显存消耗。 actor_rollout_ref.actor.ppo_micro_batch_size critic.ppo_micro_batch_size(Deprecated) 作用:已弃用,被 per_gpu 版本取代,因为它能更好地适应分布式训练环境。 Dynamic Batch Size 当样本长度差异很大时,按样本数量划分批次可能导致不同批次的计算量极不均衡,而基于 token 总数来控制 batch size 是一种平衡每个 batch 训练时间的方案。 ...

August 14, 2025 · 6 min · 1133 words · Me

[SGLang] 后端代码速览

本文档为开发者提供 SGLang 后端代码的代码梳理,按照一个请求从输入到最后输出的顺序进行讲解。下图简要介绍了这一流程: 具体而言,请求的处理过程如下: 用户启动 Server ,初始化 FastAPI App、TokenizerManager、DetokenizerManager 和 Scheduler,每个组件运行各自的无限事件循环(infinite event loop)。 用户向 FastAPI Server 发送 /v1/chat/completions 请求,Server 通过 v1_chat_completions endpoint 将请求转发到 TokenizerManager。 v1_chat_completions 函数将请求转换为 ChatCompletionRequest,再转换为 GenerateReqInput,并调用 TokenizerManager 的 generate_request 方法。 TokenizerManager 对请求进行 tokenization,并以 Python 对象(pyobj)形式将其转发给 Scheduler,同时调用 TokenizerManager 的 _wait_one_response 方法。 Scheduler 在事件循环 event_loop_normal 中处理请求: Scheduler 通过 recv_requests 接收请求,调用 process_input_requests 处理输入,通过 handle_generate_request 管理生成请求的逻辑,并将其加入 waiting_queue。 从 waiting_queue 中,Scheduler 使用 get_next_batch_to_run 为即将处理的请求创建 ScheduleBatch。 Scheduler 执行 run_batch 函数,将 ScheduleBatch 转换为 ModelWorkerBatch。 Scheduler 调用 TpModelWorker 的 forward_batch_generation,等待 logits_output 和 next_token_ids。 TpModelWorker 初始化 ForwardBatch,将其转发至 ModelRunner,并等待 logits_output。 ModelRunner 处理 ForwardBatch,调用 forward_extend 执行模型的前向计算(forward pass)。 模型通过 AttentionBackend 加速生成 logits,返回给 ModelRunner,进而返回给 TpModelWorker。 TpModelWorker 从 ModelRunner 接收 logits_output,调用 ModelRunner 的 sample 方法生成 next_token_ids,并将其发送回 Scheduler。 Scheduler 通过 process_batch_result 处理批次结果,使用 tree_cache.cache_finished_req(req) 缓存请求,并通过 check_finished 验证完成状态。对于未完成的请求,Scheduler 继续其事件循环,直到这个请求满足结束条件;对于已完成的请求,则转发到 Scheduler 的 stream_output。 在 stream_output 函数中,Scheduler 处理输出,将其包装成 BatchTokenIDOut,并发送给 DetokenizerManager。 DetokenizerManager 在其事件循环中接收 BatchTokenIDOut,处理后生成 BatchStrOut 并返回给 TokenizerManager。 ...

August 13, 2025 · 5 min · 885 words · Me

MoE环游记:2、深入负载均衡

在上一篇文章中,我们介绍了MoE的一个几何诠释,旨在通过Dense模型的最佳逼近出发来推导和理解MoE。同时在文末我们也说了,给出MoE的计算公式仅仅是开始,训练一个实际有效的MoE模型还有很多细节补,比如本文要讨论的负载均衡(Load Balance)问题。 负载均衡,即"不患寡而患不均",说白了就是让每个Expert都在干活,并且都在干尽可能一样多的活,避免某些Expert浪费算力。负载均衡既是充分利用训练算力的需求,也是尽可能发挥MoE大参数量潜力的需求。 问题分析 我们知道,MoE的基本形式是 $$ \boldsymbol{y} = \sum_{i\in \mathop{\text{argtop}}_k \boldsymbol{\rho}} \rho_i \boldsymbol{e}_i $$ 对于传统MoE,$\boldsymbol{\rho}$是一个概率分布(Router),$\boldsymbol{e}_i=\boldsymbol{v}_i$,$\boldsymbol{v}_i$是一个小型FFN(Expert)的输出;而对于我们上一篇推导的几何MoE,$\boldsymbol{\rho}$没有归一化的要求,它预测的是Expert的模长,而$\boldsymbol{e}_i=\boldsymbol{v}_i/\Vert\boldsymbol{v}_i\Vert$预测的是Expert的方向。 不管哪种格式的MoE,实际表现都差不多,只是理解视角的不同。但要注意,虽然MoE的公式给人的感觉是"每遇到一个Token,就去找相应的Expert来计算",但实际训练时其实是反过来的:先给每个Expert分配好相应的算力,然后将Token分配(Route)到所属的Expert中并行计算,这也就为什么负责打分的$\boldsymbol{\rho}$被称为Router。 这样一来,如果Expert的分配不均衡,就可能出现如下局面:某些Expert(Dead Expert)几乎一直闲置,浪费算力;某些Expert要处理的Token太多,根本忙不过来,只能Token Drop(即放弃处理部分Token)。从理论上来说,出现Dead Expert意味着MoE没有达到预期的参数量,即花了大参数量的显存,结果只训出来小参数量的效果。 所以,不管是从训练还是性能角度看,我们都希望保证Expert的负载均衡。 辅助损失(Auxiliary Loss) 促进负载均衡的常规思路是添加与之相关的损失函数,我们通常称之为"Aux Loss(Auxiliary Loss)",目前主流用的Aux Loss最早可以追溯到2020年的《GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding》。 介绍Aux Loss之前,我们需要先引入一些新概念。首先,我们已经提到对于一般的MoE来说,$\boldsymbol{\rho}$未必是概率分布,我们将归一化的$\boldsymbol{\rho}$记为$\boldsymbol{p}=[p_1,p_2,\cdots,p_n]$,以及它Top-$k$版为$\boldsymbol{f}=[f_1,f_2,\cdots,f_n]$,其中 $$ p_i = \frac{\rho_i}{\sum_{i=1}^n \rho_i},\qquad f_i = \begin{cases}1/k, & i\in \mathop{\text{argtop}}_k \boldsymbol{\rho} \\ 0, & i\not\in \mathop{\text{argtop}}_k \boldsymbol{\rho}\end{cases} $$ 接着我们定义$\boldsymbol{P}=\mathbb{E}[\boldsymbol{p}],\boldsymbol{F}=\mathbb{E}[\boldsymbol{f}]$,这里的$\mathbb{E}$是指对所有样本的所有Token做平均。不难看出,$\boldsymbol{F}$就是Expert当前的负载分布,而$\boldsymbol{P}$则相当于$\boldsymbol{F}$的一个光滑近似。 有了这些记号,我们就可以写出Aux Loss为: $$ \mathcal{L}_{\text{aux}} = \boldsymbol{F}\cdot \boldsymbol{P} = \sum_{i=1}^n F_i P_i \tag{1} $$ ...

August 10, 2025 · 5 min · 1055 words · Me

MoE环游记:1、从几何意义出发

MoE(Mixture of Experts)架构的流行自不必多说,近来火出圈的DeepSeek-V3便是MoE架构,传言GPT-5也是MoE架构,国内最近出的一些模型(Qwen3系列相关)也有不少用上了MoE。然而,虽然MoE的研究由来已久,但其应用长时间内都不愠不火,大致上是从去年初的《Mixtral of Experts》开始,MoE才逐渐吸引大家的注意力,其显著优点是参数量大,但训练和推理成本都显著低。 但同时MoE也有一些难题,如训练不稳定、负载不均衡、效果不够好等,这也是它早年没有流行起来的主要原因。不过随着这两年关注度的提升,这些问题在很大程度上已经得到解决,我们在接下来的介绍中会逐一谈到这些内容。 问题定义 我们知道,Transformer模型由Attention层和MLP层组成,MoE替换的是模型中MLP层。MLP层又分FFN(FeedForward Network)和GLU(Gated Linear Unit)两种,主流的是GLU,但简单起见我们还是以FFN为例: $$y=f(xW^{(A)})W^{(B)}$$其中$x\in\mathbb{R}^d$ 是输入向量(行向量),$W^{(A)}\in\mathbb{R}^{d\times{D}}$, $W^{(B)}\in\mathbb{R}^{D\times{d}}$ 是两个参数矩阵,$f$是Element-wise的激活函数,设$n$是一个能整除$D$的整数,那么上面的FFN可以用分块矩阵等价: $$ \begin{equation}\boldsymbol{y} = f\big(\boldsymbol{x}\begin{bmatrix}\boldsymbol{W}^{(A)}_1 & \boldsymbol{W}^{(A)}_2 & \cdots & \boldsymbol{W}^{(A)}_n\end{bmatrix}\big)\begin{bmatrix}\boldsymbol{W}^{(B)}_1 \\ \boldsymbol{W}^{(B)}_2 \\ \vdots \\ \boldsymbol{W}^{(B)}_n\end{bmatrix} = \sum_{i=1}^n \underbrace{f(\boldsymbol{x}\boldsymbol{W}^{(A)}_i)\boldsymbol{W}^{(B)}_i}_{\boldsymbol{v}_i}\end{equation} $$ 其中 $W^{(A)}_i = W^{(A)}_{[:,(i-1)c:ic]}$, $W^{(B)}_i = W^{(B)}_{[(i-1)c:ic,:]}$, $c= D/n$,这里的切片按照Python规则来。由此可见,FFN可以等价表示成n个向量 $\boldsymbol{v}_1,\boldsymbol{v}_2,\cdots,\boldsymbol{v}_n$ 之和,每个向量代表了一个小模型$f(\boldsymbol{x}\boldsymbol{W}^{(A)}_i)\boldsymbol{W}^{(B)}_i$的输出,每个小模型计算量相同,这些小模型就是MoE中的“Expert”。 MoE提出的问题是: 能否只挑k个向量的和来逼近n个向量的和呢?这样就可以将计算量降低到k/n了。 模长排序 要解决上述的问题,实质上是要解决低秩近似的问题,数学公式就是: $$\begin{equation}\mathop{\text{argmin}}_{\lambda_1,\lambda_2,\cdots,\lambda_n\in\{0,1\}}\left\Vert\sum_{i=1}^n \lambda_i \boldsymbol{v}_i - \sum_{i=1}^n\boldsymbol{v}_i\right\Vert^2\quad\text{s.t.}\quad \sum_{i=1}^n \lambda_i = k\end{equation}$$ 记$\gamma_i = 1 - \lambda_i$,那么它又可以写成: $$\begin{equation}\mathop{\text{argmin}}_{\gamma_1,\gamma_2,\cdots,\gamma_n\in\{0,1\}}\left\Vert\sum_{i=1}^n \gamma_i \boldsymbol{v}_i\right\Vert^2\quad\text{s.t.}\quad \sum_{i=1}^n \gamma_i = n - k\end{equation}$$ 这个问题的精确求解是比较困难的(NP Hard),但有一个简单的近似解:当$v_i$两两正交时,我们有 $$\begin{equation}\left\Vert\sum_{i=1}^n \gamma_i \boldsymbol{v}_i\right\Vert^2 = \sum_{i=1}^n \gamma_i^2 \Vert\boldsymbol{v}_i\Vert^2 = \sum_{i=1}^n \gamma_i \Vert\boldsymbol{v}_i\Vert^2\end{equation}$$ 上式最优解显然就是让模长$\Vert\boldsymbol{v}_i\Vert$最小的$n-k$个$\gamma_i$等于1,这又等价于说挑出模长最大的$k$个向量来逼近$n$个向量之和。当$v_i$不满足两两正交的条件时,我们依然用它来作为一个近似解。它的几何意义也很直观,模长越大的向量,在求和过程中越不容易被抵消,从而作用越突出。 ...

August 8, 2025 · 1 min · 147 words · Me