尧图网页设计 
CANN 序列级均衡调度引擎RL On-Policy 推理长尾场景的 Rollout Rebalance 实战解析【免费下载链接】cann-recipes-train本项目针对LLM与多模态模型训练业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-train导读本文以 CANNcann-recipes-train仓库中 docs/features/rollout_rebalance.md 为骨架系统讲解 RL On-PolicyRLHF/GRPO/DAPO 等训练中 Rollout 阶段的木桶效应与长尾负载均衡问题。文章覆盖该特性的使能方式、全部配置项、基于 BatchSize 档位预设的最大档位最小化均衡算法、RequestSEQ级跨 Rank 数据搬迁与 KV Cache 恢复、Rollout 结果还原的完整实现并结合 rollout_rebalance.py 源码给出底层调用链分析最终附上 Atlas A3 集群上的实测性能数据。读者读完后可直接在基于 verl vllm_ascend 的 NPU 训练环境中开启该特性并理解其收益来源与适用前提。1. 背景Rollout 阶段的长尾分布与木桶效应RL 训练中的 Rollout推理采样阶段存在典型的木桶效应由于输入 Prompt 所生成的响应Response长度呈长尾分布少数极长的生成任务会拖慢整个 DP数据并行组的进展。处理短序列的节点在完成计算后不得不进入长时间的闲置等待造成算力浪费。长尾问题优化的本质是 RL 训练系统的负载均衡。对于单轮推理的同步场景本特性将优化目标定位为提升进入长尾状态后的推理效率——即当部分 Rollout 提前结束、各 Rank 间出现负载不均时对仍未结束的 Rollout 进行负载均衡的策略分析和重调度从而提升计算资源利用率和长尾状态下的整体吞吐。从仓库结构上看该特性属于 verl 侧的独立补丁目录deepseek 场景位于 llm_rl/deepseek/verl_patches/features/rollout_optimizeqwen3 场景位于 llm_rl/qwen3/verl-mindspeed/patches/verl/features/rollout_optimize两套实现共享同一套设计思想。2. 前置依赖多档位预编图Multi-Graph本方案依赖 vllm_ascend 提供的能力在torchair_graph_config中配置use_cached_graph与graph_batch_sizes支持提前预编译多档位 BatchSize 的推理图。推理过程中随着剩余序列数减少引擎会自动匹配最小 BatchSize 的图进行推理从而避免用大 BatchSize 图处理少量序列造成的算力浪费。Rebalance 与多档位编图的配合逻辑是Rebalance 负责主动调度让各 Rank 的剩余序列数尽量收敛到某个低档位对应的 BatchSize多档位编图负责执行侧匹配让实际推理自动落到与剩余序列数匹配的最小档位图。这也是配置项multi_graph被描述为如果不开启rebalance 依然会按预编图的档位做均衡调度但不会形成明显的性能收益的原因——没有多档位图均衡调度只是让序列分布均匀却无法转化为执行侧的档位降低收益。从源码看rollout_rebalance.py 中的enable_rollout_rebalance通过hook.before(LLM, __init__)拦截 LLM 初始化additional_config kwargs.get(additional_config) max_batch_size additional_config[torchair_graph_config][graph_batch_sizes][0] # 滤除定义的graph_batch_sizes中超出max_batch_size的部分 graph_batch_sizes [max_batch_size] for bs in sorted(set(config.graph_batch_sizes), reverseTrue): if bs max_batch_size: graph_batch_sizes.append(bs) States.graph_batch_sizes graph_batch_sizes if config.multi_graph: additional_config[torchair_graph_config].update( graph_batch_sizes_initFalse, use_cached_graphTrue, graph_batch_sizesgraph_batch_sizes, )可以看到配置中声明的档位如[64, 32, 16, 8, 4]会先以torchair_graph_config.graph_batch_sizes的第一个值最大 BatchSize为上限过滤再降序排列最终写入torchair_graph_config并打开use_cached_graph。这保证了预编译的档位图不会超出硬件/引擎允许的最大 BatchSize。3. 使用说明3.1 初始化配置特性通过环境变量ROLLOUT_REBALANCE_ENABLE1使能。在 deepseek 场景中llm_rl/deepseek/verl_patches/workers/megatron_workers.py 的actor_rollout_ref_init_model对应 verl 原版ActorRolloutRefWorker.init_model的补丁实现开头已内置使能逻辑register(dispatch_modeDispatch.ONE_TO_ALL) def actor_rollout_ref_init_model(self): if os.getenv(ROLLOUT_REBALANCE_ENABLE, 0) ! 0: from features.rollout_optimize.rollout_rebalance import enable_rollout_rebalance enable_rollout_rebalance() ...在 qwen3/verl-mindspeed 场景中对应的补丁为 0011-verl-feature-enable_rollout_rebalance.patch其入口函数为init_rollout_rebalance由 features/rollout_optimize/init.py 导出。官方文档给出的通用接线方式为在 verl 的verl/workers/megatron_workers.py文件开头追加from patches.verl.features.rollout_optimize import init_rollout_rebalance init_rollout_rebalance()即只需设置环境变量并在 worker 初始化入口调用一次初始化函数后续所有逻辑均通过 hook 自动注入无需改动 verl 其余代码。3.2 配置项介绍可以直接在config.py中修改配置也可以将以下配置写入 verl 启动的 yaml 中在初始化位置进行配置提取并传入init_rollout_rebalance/enable_rollout_rebalance方法。配置类定义见 llm_rl/deepseek/verl_patches/features/rollout_optimize/config.pyclass RolloutRebalanceConfig: enable True # RolloutRebalance特性总开关 check_interval 1000 # 间隔多少个step进行一次rebalance检查 multi_graph True # 是否开启多档位编图如果关闭rebalance依然会按预编图的档位做均衡调度但是不会形成明显的性能收益 graph_batch_sizes [64, 32, 16, 8, 4] # 预编图的档位设置 profile True # 是否打印过程中的性能数据 profile_interval 100 # 打印间隔步长各配置项的作用与建议配置项默认值含义调优建议enableTrue文档版默认读取环境变量ROLLOUT_REBALANCE_ENABLE默认0特性总开关不使能时enable_rollout_rebalance直接返回不注入任何 hookcheck_interval1000间隔多少个 step 进行一次 Rebalance 检查检查越频繁调度越及时但 all_gather / all_to_all 开销越大长尾场景下建议按 step 耗时调整multi_graphTrue是否开启多档位编图关闭时 Rebalance 仍会按预编图档位均衡调度但无法形成明显的性能收益graph_batch_sizes[64, 32, 16, 8, 4]预编图的档位设置档位越低长尾阶段匹配到的执行图越小、单步耗时越低档位列表会以引擎最大 BatchSize 为上限自动过滤profileTrue是否打印过程中的性能数据用于观察各 Rank 剩余序列数SeqCntMap与当前最大档位MaxBS的变化profile_interval100打印间隔步长配合profile使用4. 实现方案序列级均衡调度引擎本优化的核心目标在 On-Policy 场景中针对部分 Rollout 提前结束导致各 Rank 间负载不均时对未结束的 Rollout 进行负载均衡的策略分析和重调度从而提升计算资源的利用率和长尾状态下的推理效率。方案包含三大关键功能Rebalance 条件检测与调度策略生成周期性采集各 Rank 全局状态计算最优请求迁移清单RequestSEQ级的数据搬迁与恢复含对应 KV Cache将序列连同其 KV Cache 在 DP 组内跨 Rank 迁移Rollout 后的结果还原将迁移到其他 Rank 完成的序列输出通过 all_gather 归还给源 Rank并还原被临时拆散的父子请求关系。以下代码均位于 llm_rl/deepseek/verl_patches/features/rollout_optimize/rollout_rebalance.pyqwen3 版本路径为 llm_rl/qwen3/verl-mindspeed/patches/verl/features/rollout_optimize/rollout_rebalance.py。4.1 全局状态感知DP 组内的请求分布采集Rebalance 决策的前提是知道每个 Rank 上还剩多少序列。get_current_state从llm_engine.output_processor.request_states中提取本 Rank 当前的请求 ID 列表sync_group_states通过dist.all_gather_object在 DP 组内同步得到全局状态快照def get_current_state(self): return dict( rankself.rank, req_idslist(self.llm_engine.output_processor.request_states.keys()), ) def sync_group_states(self): if self.world_size 1: return [self.get_current_state()] group_states [None for _ in range(self.world_size)] dist.all_gather_object(group_states, self.get_current_state(), groupself.dp_group) return group_states值得注意的是RolloutRebalanceEngine.__init__中会先执行一步关键预处理移除 request 的父子关系req_state.parent_req None。这是因为 Rebalance 需要让单个 Seq 能够独立搬迁被迁移序列的父子关系会在最后的recover阶段统一还原。4.2 基于 BatchSize 档位预设的最大档位最小化均衡算法calc_balancing_tasks负责计算最优的请求迁移任务清单。其最优策略定义为三层目标主要目标将整个 DP 组所需的最大档位max_bs降至最低次要目标在满足 1 的前提下使迁移的请求数量cost最少补充目标在满足 1 和 2 的前提下使得各 Rank 间的数据搬迁流向尽可能均匀避免多对一阻塞。算法流程以仓库当前 deepseek 版本源码为准其输入为remaining_reqs_by_rank即 4.1 节采集到的全局状态def calc_balancing_tasks(self, remaining_reqs_by_rank): if not remaining_reqs_by_rank or len(remaining_reqs_by_rank) 1: return [] rank_req_ids_map {item[rank]: list(item[req_ids]) for item in remaining_reqs_by_rank} rank_req_cnt_map {rank: len(ids) for rank, ids in rank_req_ids_map.items()} ranks sorted(rank_req_cnt_map.keys()) # 寻找所有可行的优化方案 max_bs_before max(self._get_bs(bs) for bs in rank_req_cnt_map.values()) avg_bs sum(rank_req_cnt_map.values()) / len(rank_req_cnt_map.values()) max_bs_next None for target_bs in States.graph_batch_sizes[::-1]: # 优先匹配出最小的BS if avg_bs target_bs max_bs_before: max_bs_next target_bs break if not max_bs_next: return [] # 分析各个worker的可调度量 donors [] receivers [] for rank in ranks: req_cnt rank_req_cnt_map[rank] delta req_cnt - max_bs_next if delta 0: donors.append(dict(rankrank, surplusdelta, req_idsrank_req_cnt_map[rank][:delta])) elif delta 0: receivers.append(dict(rankrank, capacity-delta)) # 生成迁移任务清单为了让迁移速度更快尽可能将请求均衡分发给目标worker # 但优先分发给最空闲的worker即容量最大的worker balancing_tasks [] while True: donor_index 0 for receiver in sorted(receivers, keylambda r: r[capacity], reverseTrue): donor donors[donor_index] num_to_move min(donor[surplus], receiver[capacity]) balancing_tasks [dict(from_rankdonor[rank], to_rankreceiver[rank], req_idreq_id) for req_id in donor[req_ids][:num_to_move]] donor[req_ids] donor[req_ids][num_to_move:] donor[surplus] - num_to_move receiver[capacity] - num_to_move donor_index (donor_index 1) % len(donors) donors [x for x in donors if x[surplus]] if not donors: break receivers [x for x in receivers if x[capacity]] return balancing_tasks算法关键点解读档位降级判定max_bs_before是当前 DP 组所需的最大档位avg_bs是平均剩余序列数。算法从最小档位向上遍历States.graph_batch_sizes找到能容纳平均剩余数且低于当前最大档位的最小目标档位max_bs_next即均衡后的最大档位。若找不到如平均值本身就高于当前档位说明不值得迁移直接返回空任务。_get_bs(size)映射将序列数映射为不小于它的最小档位倒序遍历档位列表返回第一个size bs的档位若超出最大档位则返回原始 size。这决定了 TPOT 档位评估的依据。迁移方向判定req_cnt max_bs_next的 Rank 为捐赠方donor有多余序列req_cnt max_bs_next的 Rank 为接收方receiver有剩余容量每个接收方分配到的序列数不超过其容量。均匀分发策略接收方按容量降序最空闲优先排序捐赠方轮流round-robindonor_index (donor_index 1) % len(donors)供给从而让数据搬迁流向尽可能均匀避免多对一阻塞。注关联文档中给出的calc_balancing_tasks为早期实现版本含running/waiting状态、block_usage容量估算与need_kv_cache标记仓库当前源码已演进为按req_ids计数的精简版本二者策略目标最大档位最小化 → 迁移成本最小 → 流向均匀完全一致。早期版本中block_usage 0.8的 Rank 不作为接收方的约束在现版本中体现为仅按剩余序列数与档位容量做纯调度决策。在早期文档版本中还会针对等待态请求waiting做特殊处理由于等待态请求尚未进入计算、无需搬迁 KV Cache因此优先将其迁出need_kv_cacheFalse只有剩余无等待态请求可迁时才动用正在运行running的请求need_kv_cacheTrue需搬迁 KV Cache。这一分层搬迁思路在现版本中通过load_by_req_id统一采集 KV Cache 实现。4.3 序列请求的跨 Rank 发送与接收含 KV Cache得到迁移任务清单后all_to_all_v_tasks执行实际的序列搬迁。核心流程为采集并 abort对属于本 Rank 的每个迁移任务用RebalanceRequestTask.load_by_req_id采集该请求的完整状态prompt 与已生成 token、logprobs 缓存、KV Cache 块并调用trigger_abort中止其在本地引擎中的执行元数据 all_to_all将每个目标 Rank 要接收的请求元数据get_transfer_dict经 pickle 序列化后转为 uint8 tensor通过dist.all_to_all_single交换——先交换各 Rank 的数据尺寸再交换实际数据KV Cache 直传send_kv_caches对每个需要 KV Cache 的请求逐层逐块dist.send接收侧load_received_tasks中按形状dist.recv接收加载还原load_by_transfer_infotrigger_load将请求重新装载到目标 Rank 的 LLMEngine、Scheduler 与 ModelRunner 中。def all_to_all_v_tasks(self, schedule_tasks): objects_to_send [[] for _ in range(dist.get_world_size())] send_tasks [] for schedule_task in schedule_tasks: if self.rank schedule_task[from_rank]: request_task RebalanceRequestTask(self.llm_engine).load_by_req_id(schedule_task[req_id]) send_tasks.append((request_task, schedule_task[to_rank])) objects_to_send[schedule_task[to_rank]].append(request_task.get_transfer_dict()) request_task.trigger_abort() tensor_list [ torch.frombuffer(pickle.dumps(data), dtypetorch.uint8).to(self.device) for data in objects_to_send] local_sizes torch.tensor([len(t) for t in tensor_list], dtypetorch.long).to(self.device) remote_sizes torch.empty_like(local_sizes).to(self.device) dist.all_to_all_single(remote_sizes, local_sizes) input_tensor torch.cat(tensor_list).to(self.device) output_tensor torch.empty(remote_sizes.sum().item(), dtypetorch.uint8).to(self.device) dist.all_to_all_single( output_tensor, input_tensor, output_split_sizesremote_sizes.tolist(), input_split_sizeslocal_sizes.tolist(), ) received_tensor torch.split(output_tensor, remote_sizes.tolist()) received_tasks [] for rank_data in received_tensor: received_tasks pickle.loads(rank_data.to(cpu).numpy().tobytes()) self.send_kv_caches(send_tasks) self.load_received_tasks(received_tasks)KV Cache 的发送侧deepseek 版本源码按层堆叠后逐块发送def send_kv_caches(self, send_tasks): if States.reprefill_mode: return for request_task, to_rank in send_tasks: for kv_cache_block in request_task.layers_kv_cache_blocks: dist.send(kv_cache_block, dstto_rank) rank_log_info(f[TaskSendKvCache][ToRank{to_rank}][ReqId{request_task.req_id}])接收侧在load_received_tasks中按layers_kv_cache_shapes预分配空 tensor 后dist.recv并通过日志forceTrue强制打印记录ReceivedTask、ReceivedKvCache、ReceivedTaskLoaded三个阶段的耗时毫秒便于定位搬迁开销。注意源码中接收侧使用.cuda()申请设备内存States.reprefill_mode为 False 时States.reprefill_mode为 True 时跳过 KV Cache 搬迁走重新 prefillreprefill路径——对应recover_scheduler_request中将num_computed_tokens置 0、状态置为PREEMPTED的分支。4.4 KV Cache 搬迁与恢复RebalanceRequestTask.load_by_req_id中完成 request 级 KV Cache 采集通过model_runner.input_batch的 block_table 找到该请求占用的 KV Cache 块索引再对每层每块做索引采样并堆叠得到该请求专属的 KV Cache 张量列表# request级的kvCache采集 self.layers_kv_cache_blocks [] for cache_block_index in range(len(self.global_kv_caches[0])): self.layers_kv_cache_blocks.append( torch.stack([layer[cache_block_index][request_block_table] for layer in self.global_kv_caches]) )接收侧在recover_model_runner中还原先为请求分配新的 KV Cache 槽位allocate_slots重建CachedRequestState并加入input_batch然后把搬迁来的 KV Cache 块写入新槽位# kvCache还原 for layer_index, layer_caches in enumerate(self.model_runner.kv_caches): reload_indexes list(range(len(new_block_ids))) for i, cache_block in enumerate(self.layers_kv_cache_blocks): layer_caches[i][new_block_ids] cache_block[layer_index][reload_indexes]配合recover_request_state还原logprobs_processor的cumulative_logprob与 logprobs 列表、detokenizer的 token 缓存、is_prefillingFalse与recover_scheduler_request还原max_tokens、已生成 token、num_computed_tokens并重新挂入scheduler.running目标 Rank 即可无缝续推该序列仿佛它一直运行在本机。4.5 Rollout 后的结果还原迁移到其他 Rank 的请求其request_id会被改写为src_rank_{源Rank}_{原req_id}前缀形式见trigger_load中new_req_id的构造。recover方法负责将输出归还_split_rebalance_outputs从LLMEngine.step输出中分离出带src_rank前缀的 rebalance 输出解析源 Rank 与原始 req_id抽取 token_ids、logprobs 等与本地正常输出get_rebalance_outputs通过dist.all_gather_object在 DP 组内同步取回属于本 Rank 的 rebalance 输出_build_rebalance_request_output将其还原为标准的RequestOutputCompletionOutputLogprob父子关系还原current_outputs_map以去除序号前缀后的request_id为 key 归并因为同一请求可能被拆成多个src_rank_*子请求分别完成将子输出挂到父RequestOutput.outputs下最终返回还原后的输出集合。def recover(self, outputs): # 将进行了rebalance迁移的outputs通过all_gather还原到源rank rebalance_outputs, current_outputs self._split_rebalance_outputs(outputs) for rebalance_output in self.get_rebalance_outputs(rebalance_outputs): rank_log_info(f[RecvReqId{rebalance_output[req_id]}], forceTrue) current_outputs.append(self._build_rebalance_request_output(rebalance_output)) current_outputs_map {} for request_output in current_outputs: request_id request_output.request_id if _ in request_id: request_id request_id.split(_, 1)[-1] parent_request current_outputs_map.get(request_id) if not parent_request: parent_request RequestOutput( request_idrequest_id, promptNone, prompt_token_idsrequest_output.prompt_token_ids, prompt_logprobsNone, outputs[], finishedTrue) current_outputs_map[request_id] parent_request parent_request.outputs request_output.outputs return current_outputs_map.values()此外enable_rollout_rebalance中的hook.after(LLMEngine, step)会把 step 输出中带src_rank前缀的RequestOutput暂存到States.outputs_cache待全部迁移请求完成后统一参与结果还原避免中间输出被打散。而LLM._run_engine前 hook 中还会缓存首个请求的sampling_paramsStates.sampling_params供迁移加载时copy复用采样策略。4.6 检测节流与性能剖析check()由LLMEngine.step的 after hook 每次触发但内部通过CheckCounterutils.py控制频率check_interval步做一次 Rebalance 检测、profile_interval步做一次性能采样且二者按阈值各自计数、互不阻塞从而控制 all_gather / all_to_all 的通信开销。profile()打印BSMap各 Rank 映射后的档位、SeqCntMap各 Rank 剩余序列数、CurrentMaxBS、ProfileStepCost、TPOT由step 耗时 * 1000 / threshold估算并在最大档位发生变化时输出MaxBSChanged提示。日志默认只在 rank 0 打印rank_log_info的forceFalse分支排障时可用forceTrue观察每个 Rank 的收发细节。5. 使能效果长尾场景下的实测数据关联文档在Atlas A3 集群 64 卡环境上进行了验证实验开启本方案后单轮推理耗时从约 10200s 优化到约 6100s性能收益达 60% 左右。实验配置如下配置项值模型Qwen3 235B数据集deepscalerdata.train_batch_size512data.max_response_length32768actor_rollout_ref.rollout.n16TP / DPTP4DP32性能收益主要来源于单个 step 的 TPOTTime Per Output Token差距默认场景下长尾阶段的 TPOT 会从 125ms 上升到 200ms大 BatchSize 档位图处理少量剩余序列算力利用率低使能 Rebalance 并配合多档位编图后能在 1~2K 推理长度时就快速将推理档位降低让单个 step 的 TPOT 降低到60ms 的量级长尾阶段持续的时间越长如max_response_length32768的极端长尾任务性能差距被持续放大总收益越显著。仓库中 llm_rl/deepseek/verl_patches/features/rollout_optimize/README.md 还记录了 DeepSeek 场景的另一组对照实验Atlas A3 128 卡、DeepSeekV3、open-r1/OpenR1-Math-220K、TP2/DP128单轮推理耗时从约 6200s 优化到约 2300s性能收益为 57%~62%与 Qwen3 场景结论一致。6. 适用前提与使用建议适用场景On-PolicyRLHF/GRPO 等训练中单轮推理的同步 Rollout 场景尤其是max_response_length大、响应长度长尾明显的任务推理、代码、数学等对于响应长度接近均匀分布的场景Rebalance 收益有限。前置依赖vllm_ascend 需支持torchair_graph_config中的use_cached_graph与graph_batch_sizes即多档位预编图能力multi_graph应保持开启以获得最大收益。档位设计graph_batch_sizes的档位应根据 DP 组内每 Rank 的平均序列数设计最大档位会被引擎max_batch_size自动钳制档位越低长尾阶段的单步耗时越低但预编译图的数量和显存占用也会增加。检查频率check_interval需在调度及时性与 all_gather/all_to_all 通信开销之间权衡建议结合profile输出的[RebalanceSchedule][Costxx ms]评估。迁移日志[RebalanceScheduleTasks]、[Rebalance][ReqId][Src][Dst]、[ReceivedTask/KvCache]、[MaxBSChanged]等日志可用于确认调度是否按预期发生以及各 Rank 剩余序列数[Rankx, Remainy]是否收敛到目标档位。7. 相关资源特性文档docs/features/rollout_rebalance.md核心实现deepseekllm_rl/deepseek/verl_patches/features/rollout_optimize/rollout_rebalance.py、config.py、utils.py特性说明deepseekllm_rl/deepseek/verl_patches/features/rollout_optimize/README.md使能入口deepseekllm_rl/deepseek/verl_patches/workers/megatron_workers.py使能补丁与实现qwen30011-verl-feature-enable_rollout_rebalance.patch 及 llm_rl/qwen3/verl-mindspeed/patches/verl/features/rollout_optimize训练目录总览deepseekllm_rl/deepseek/README.mdverl_patches/features/rollout_optimize即本特性的部署目录同族 RL 特性文档docs/features/history_tree_dynamic_rl.md、docs/features/length_aware_resampler.md【免费下载链接】cann-recipes-train本项目针对LLM与多模态模型训练业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-train创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 CANN 序列级均衡调度引擎：RL On-Policy 推理长尾场景的 Rollout Rebalance 实战解析
CANN 序列级均衡调度引擎RL On-Policy 推理长尾场景的 Rollout Rebalance 实战解析【免费下载链接】cann-recipes-train本项目针对LLM与多模态模型训练业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-train导读本文以 CANNcann-recipes-train仓库中 docs/features/rollout_rebalance.md 为骨架系统讲解 RL On-PolicyRLHF/GRPO/DAPO 等训练中 Rollout 阶段的木桶效应与长尾负载均衡问题。文章覆盖该特性的使能方式、全部配置项、基于 BatchSize 档位预设的最大档位最小化均衡算法、RequestSEQ级跨 Rank 数据搬迁与 KV Cache 恢复、Rollout 结果还原的完整实现并结合 rollout_rebalance.py 源码给出底层调用链分析最终附上 Atlas A3 集群上的实测性能数据。读者读完后可直接在基于 verl vllm_ascend 的 NPU 训练环境中开启该特性并理解其收益来源与适用前提。1. 背景Rollout 阶段的长尾分布与木桶效应RL 训练中的 Rollout推理采样阶段存在典型的木桶效应由于输入 Prompt 所生成的响应Response长度呈长尾分布少数极长的生成任务会拖慢整个 DP数据并行组的进展。处理短序列的节点在完成计算后不得不进入长时间的闲置等待造成算力浪费。长尾问题优化的本质是 RL 训练系统的负载均衡。对于单轮推理的同步场景本特性将优化目标定位为提升进入长尾状态后的推理效率——即当部分 Rollout 提前结束、各 Rank 间出现负载不均时对仍未结束的 Rollout 进行负载均衡的策略分析和重调度从而提升计算资源利用率和长尾状态下的整体吞吐。从仓库结构上看该特性属于 verl 侧的独立补丁目录deepseek 场景位于 llm_rl/deepseek/verl_patches/features/rollout_optimizeqwen3 场景位于 llm_rl/qwen3/verl-mindspeed/patches/verl/features/rollout_optimize两套实现共享同一套设计思想。2. 前置依赖多档位预编图Multi-Graph本方案依赖 vllm_ascend 提供的能力在torchair_graph_config中配置use_cached_graph与graph_batch_sizes支持提前预编译多档位 BatchSize 的推理图。推理过程中随着剩余序列数减少引擎会自动匹配最小 BatchSize 的图进行推理从而避免用大 BatchSize 图处理少量序列造成的算力浪费。Rebalance 与多档位编图的配合逻辑是Rebalance 负责主动调度让各 Rank 的剩余序列数尽量收敛到某个低档位对应的 BatchSize多档位编图负责执行侧匹配让实际推理自动落到与剩余序列数匹配的最小档位图。这也是配置项multi_graph被描述为如果不开启rebalance 依然会按预编图的档位做均衡调度但不会形成明显的性能收益的原因——没有多档位图均衡调度只是让序列分布均匀却无法转化为执行侧的档位降低收益。从源码看rollout_rebalance.py 中的enable_rollout_rebalance通过hook.before(LLM, __init__)拦截 LLM 初始化additional_config kwargs.get(additional_config) max_batch_size additional_config[torchair_graph_config][graph_batch_sizes][0] # 滤除定义的graph_batch_sizes中超出max_batch_size的部分 graph_batch_sizes [max_batch_size] for bs in sorted(set(config.graph_batch_sizes), reverseTrue): if bs max_batch_size: graph_batch_sizes.append(bs) States.graph_batch_sizes graph_batch_sizes if config.multi_graph: additional_config[torchair_graph_config].update( graph_batch_sizes_initFalse, use_cached_graphTrue, graph_batch_sizesgraph_batch_sizes, )可以看到配置中声明的档位如[64, 32, 16, 8, 4]会先以torchair_graph_config.graph_batch_sizes的第一个值最大 BatchSize为上限过滤再降序排列最终写入torchair_graph_config并打开use_cached_graph。这保证了预编译的档位图不会超出硬件/引擎允许的最大 BatchSize。3. 使用说明3.1 初始化配置特性通过环境变量ROLLOUT_REBALANCE_ENABLE1使能。在 deepseek 场景中llm_rl/deepseek/verl_patches/workers/megatron_workers.py 的actor_rollout_ref_init_model对应 verl 原版ActorRolloutRefWorker.init_model的补丁实现开头已内置使能逻辑register(dispatch_modeDispatch.ONE_TO_ALL) def actor_rollout_ref_init_model(self): if os.getenv(ROLLOUT_REBALANCE_ENABLE, 0) ! 0: from features.rollout_optimize.rollout_rebalance import enable_rollout_rebalance enable_rollout_rebalance() ...在 qwen3/verl-mindspeed 场景中对应的补丁为 0011-verl-feature-enable_rollout_rebalance.patch其入口函数为init_rollout_rebalance由 features/rollout_optimize/init.py 导出。官方文档给出的通用接线方式为在 verl 的verl/workers/megatron_workers.py文件开头追加from patches.verl.features.rollout_optimize import init_rollout_rebalance init_rollout_rebalance()即只需设置环境变量并在 worker 初始化入口调用一次初始化函数后续所有逻辑均通过 hook 自动注入无需改动 verl 其余代码。3.2 配置项介绍可以直接在config.py中修改配置也可以将以下配置写入 verl 启动的 yaml 中在初始化位置进行配置提取并传入init_rollout_rebalance/enable_rollout_rebalance方法。配置类定义见 llm_rl/deepseek/verl_patches/features/rollout_optimize/config.pyclass RolloutRebalanceConfig: enable True # RolloutRebalance特性总开关 check_interval 1000 # 间隔多少个step进行一次rebalance检查 multi_graph True # 是否开启多档位编图如果关闭rebalance依然会按预编图的档位做均衡调度但是不会形成明显的性能收益 graph_batch_sizes [64, 32, 16, 8, 4] # 预编图的档位设置 profile True # 是否打印过程中的性能数据 profile_interval 100 # 打印间隔步长各配置项的作用与建议配置项默认值含义调优建议enableTrue文档版默认读取环境变量ROLLOUT_REBALANCE_ENABLE默认0特性总开关不使能时enable_rollout_rebalance直接返回不注入任何 hookcheck_interval1000间隔多少个 step 进行一次 Rebalance 检查检查越频繁调度越及时但 all_gather / all_to_all 开销越大长尾场景下建议按 step 耗时调整multi_graphTrue是否开启多档位编图关闭时 Rebalance 仍会按预编图档位均衡调度但无法形成明显的性能收益graph_batch_sizes[64, 32, 16, 8, 4]预编图的档位设置档位越低长尾阶段匹配到的执行图越小、单步耗时越低档位列表会以引擎最大 BatchSize 为上限自动过滤profileTrue是否打印过程中的性能数据用于观察各 Rank 剩余序列数SeqCntMap与当前最大档位MaxBS的变化profile_interval100打印间隔步长配合profile使用4. 实现方案序列级均衡调度引擎本优化的核心目标在 On-Policy 场景中针对部分 Rollout 提前结束导致各 Rank 间负载不均时对未结束的 Rollout 进行负载均衡的策略分析和重调度从而提升计算资源的利用率和长尾状态下的推理效率。方案包含三大关键功能Rebalance 条件检测与调度策略生成周期性采集各 Rank 全局状态计算最优请求迁移清单RequestSEQ级的数据搬迁与恢复含对应 KV Cache将序列连同其 KV Cache 在 DP 组内跨 Rank 迁移Rollout 后的结果还原将迁移到其他 Rank 完成的序列输出通过 all_gather 归还给源 Rank并还原被临时拆散的父子请求关系。以下代码均位于 llm_rl/deepseek/verl_patches/features/rollout_optimize/rollout_rebalance.pyqwen3 版本路径为 llm_rl/qwen3/verl-mindspeed/patches/verl/features/rollout_optimize/rollout_rebalance.py。4.1 全局状态感知DP 组内的请求分布采集Rebalance 决策的前提是知道每个 Rank 上还剩多少序列。get_current_state从llm_engine.output_processor.request_states中提取本 Rank 当前的请求 ID 列表sync_group_states通过dist.all_gather_object在 DP 组内同步得到全局状态快照def get_current_state(self): return dict( rankself.rank, req_idslist(self.llm_engine.output_processor.request_states.keys()), ) def sync_group_states(self): if self.world_size 1: return [self.get_current_state()] group_states [None for _ in range(self.world_size)] dist.all_gather_object(group_states, self.get_current_state(), groupself.dp_group) return group_states值得注意的是RolloutRebalanceEngine.__init__中会先执行一步关键预处理移除 request 的父子关系req_state.parent_req None。这是因为 Rebalance 需要让单个 Seq 能够独立搬迁被迁移序列的父子关系会在最后的recover阶段统一还原。4.2 基于 BatchSize 档位预设的最大档位最小化均衡算法calc_balancing_tasks负责计算最优的请求迁移任务清单。其最优策略定义为三层目标主要目标将整个 DP 组所需的最大档位max_bs降至最低次要目标在满足 1 的前提下使迁移的请求数量cost最少补充目标在满足 1 和 2 的前提下使得各 Rank 间的数据搬迁流向尽可能均匀避免多对一阻塞。算法流程以仓库当前 deepseek 版本源码为准其输入为remaining_reqs_by_rank即 4.1 节采集到的全局状态def calc_balancing_tasks(self, remaining_reqs_by_rank): if not remaining_reqs_by_rank or len(remaining_reqs_by_rank) 1: return [] rank_req_ids_map {item[rank]: list(item[req_ids]) for item in remaining_reqs_by_rank} rank_req_cnt_map {rank: len(ids) for rank, ids in rank_req_ids_map.items()} ranks sorted(rank_req_cnt_map.keys()) # 寻找所有可行的优化方案 max_bs_before max(self._get_bs(bs) for bs in rank_req_cnt_map.values()) avg_bs sum(rank_req_cnt_map.values()) / len(rank_req_cnt_map.values()) max_bs_next None for target_bs in States.graph_batch_sizes[::-1]: # 优先匹配出最小的BS if avg_bs target_bs max_bs_before: max_bs_next target_bs break if not max_bs_next: return [] # 分析各个worker的可调度量 donors [] receivers [] for rank in ranks: req_cnt rank_req_cnt_map[rank] delta req_cnt - max_bs_next if delta 0: donors.append(dict(rankrank, surplusdelta, req_idsrank_req_cnt_map[rank][:delta])) elif delta 0: receivers.append(dict(rankrank, capacity-delta)) # 生成迁移任务清单为了让迁移速度更快尽可能将请求均衡分发给目标worker # 但优先分发给最空闲的worker即容量最大的worker balancing_tasks [] while True: donor_index 0 for receiver in sorted(receivers, keylambda r: r[capacity], reverseTrue): donor donors[donor_index] num_to_move min(donor[surplus], receiver[capacity]) balancing_tasks [dict(from_rankdonor[rank], to_rankreceiver[rank], req_idreq_id) for req_id in donor[req_ids][:num_to_move]] donor[req_ids] donor[req_ids][num_to_move:] donor[surplus] - num_to_move receiver[capacity] - num_to_move donor_index (donor_index 1) % len(donors) donors [x for x in donors if x[surplus]] if not donors: break receivers [x for x in receivers if x[capacity]] return balancing_tasks算法关键点解读档位降级判定max_bs_before是当前 DP 组所需的最大档位avg_bs是平均剩余序列数。算法从最小档位向上遍历States.graph_batch_sizes找到能容纳平均剩余数且低于当前最大档位的最小目标档位max_bs_next即均衡后的最大档位。若找不到如平均值本身就高于当前档位说明不值得迁移直接返回空任务。_get_bs(size)映射将序列数映射为不小于它的最小档位倒序遍历档位列表返回第一个size bs的档位若超出最大档位则返回原始 size。这决定了 TPOT 档位评估的依据。迁移方向判定req_cnt max_bs_next的 Rank 为捐赠方donor有多余序列req_cnt max_bs_next的 Rank 为接收方receiver有剩余容量每个接收方分配到的序列数不超过其容量。均匀分发策略接收方按容量降序最空闲优先排序捐赠方轮流round-robindonor_index (donor_index 1) % len(donors)供给从而让数据搬迁流向尽可能均匀避免多对一阻塞。注关联文档中给出的calc_balancing_tasks为早期实现版本含running/waiting状态、block_usage容量估算与need_kv_cache标记仓库当前源码已演进为按req_ids计数的精简版本二者策略目标最大档位最小化 → 迁移成本最小 → 流向均匀完全一致。早期版本中block_usage 0.8的 Rank 不作为接收方的约束在现版本中体现为仅按剩余序列数与档位容量做纯调度决策。在早期文档版本中还会针对等待态请求waiting做特殊处理由于等待态请求尚未进入计算、无需搬迁 KV Cache因此优先将其迁出need_kv_cacheFalse只有剩余无等待态请求可迁时才动用正在运行running的请求need_kv_cacheTrue需搬迁 KV Cache。这一分层搬迁思路在现版本中通过load_by_req_id统一采集 KV Cache 实现。4.3 序列请求的跨 Rank 发送与接收含 KV Cache得到迁移任务清单后all_to_all_v_tasks执行实际的序列搬迁。核心流程为采集并 abort对属于本 Rank 的每个迁移任务用RebalanceRequestTask.load_by_req_id采集该请求的完整状态prompt 与已生成 token、logprobs 缓存、KV Cache 块并调用trigger_abort中止其在本地引擎中的执行元数据 all_to_all将每个目标 Rank 要接收的请求元数据get_transfer_dict经 pickle 序列化后转为 uint8 tensor通过dist.all_to_all_single交换——先交换各 Rank 的数据尺寸再交换实际数据KV Cache 直传send_kv_caches对每个需要 KV Cache 的请求逐层逐块dist.send接收侧load_received_tasks中按形状dist.recv接收加载还原load_by_transfer_infotrigger_load将请求重新装载到目标 Rank 的 LLMEngine、Scheduler 与 ModelRunner 中。def all_to_all_v_tasks(self, schedule_tasks): objects_to_send [[] for _ in range(dist.get_world_size())] send_tasks [] for schedule_task in schedule_tasks: if self.rank schedule_task[from_rank]: request_task RebalanceRequestTask(self.llm_engine).load_by_req_id(schedule_task[req_id]) send_tasks.append((request_task, schedule_task[to_rank])) objects_to_send[schedule_task[to_rank]].append(request_task.get_transfer_dict()) request_task.trigger_abort() tensor_list [ torch.frombuffer(pickle.dumps(data), dtypetorch.uint8).to(self.device) for data in objects_to_send] local_sizes torch.tensor([len(t) for t in tensor_list], dtypetorch.long).to(self.device) remote_sizes torch.empty_like(local_sizes).to(self.device) dist.all_to_all_single(remote_sizes, local_sizes) input_tensor torch.cat(tensor_list).to(self.device) output_tensor torch.empty(remote_sizes.sum().item(), dtypetorch.uint8).to(self.device) dist.all_to_all_single( output_tensor, input_tensor, output_split_sizesremote_sizes.tolist(), input_split_sizeslocal_sizes.tolist(), ) received_tensor torch.split(output_tensor, remote_sizes.tolist()) received_tasks [] for rank_data in received_tensor: received_tasks pickle.loads(rank_data.to(cpu).numpy().tobytes()) self.send_kv_caches(send_tasks) self.load_received_tasks(received_tasks)KV Cache 的发送侧deepseek 版本源码按层堆叠后逐块发送def send_kv_caches(self, send_tasks): if States.reprefill_mode: return for request_task, to_rank in send_tasks: for kv_cache_block in request_task.layers_kv_cache_blocks: dist.send(kv_cache_block, dstto_rank) rank_log_info(f[TaskSendKvCache][ToRank{to_rank}][ReqId{request_task.req_id}])接收侧在load_received_tasks中按layers_kv_cache_shapes预分配空 tensor 后dist.recv并通过日志forceTrue强制打印记录ReceivedTask、ReceivedKvCache、ReceivedTaskLoaded三个阶段的耗时毫秒便于定位搬迁开销。注意源码中接收侧使用.cuda()申请设备内存States.reprefill_mode为 False 时States.reprefill_mode为 True 时跳过 KV Cache 搬迁走重新 prefillreprefill路径——对应recover_scheduler_request中将num_computed_tokens置 0、状态置为PREEMPTED的分支。4.4 KV Cache 搬迁与恢复RebalanceRequestTask.load_by_req_id中完成 request 级 KV Cache 采集通过model_runner.input_batch的 block_table 找到该请求占用的 KV Cache 块索引再对每层每块做索引采样并堆叠得到该请求专属的 KV Cache 张量列表# request级的kvCache采集 self.layers_kv_cache_blocks [] for cache_block_index in range(len(self.global_kv_caches[0])): self.layers_kv_cache_blocks.append( torch.stack([layer[cache_block_index][request_block_table] for layer in self.global_kv_caches]) )接收侧在recover_model_runner中还原先为请求分配新的 KV Cache 槽位allocate_slots重建CachedRequestState并加入input_batch然后把搬迁来的 KV Cache 块写入新槽位# kvCache还原 for layer_index, layer_caches in enumerate(self.model_runner.kv_caches): reload_indexes list(range(len(new_block_ids))) for i, cache_block in enumerate(self.layers_kv_cache_blocks): layer_caches[i][new_block_ids] cache_block[layer_index][reload_indexes]配合recover_request_state还原logprobs_processor的cumulative_logprob与 logprobs 列表、detokenizer的 token 缓存、is_prefillingFalse与recover_scheduler_request还原max_tokens、已生成 token、num_computed_tokens并重新挂入scheduler.running目标 Rank 即可无缝续推该序列仿佛它一直运行在本机。4.5 Rollout 后的结果还原迁移到其他 Rank 的请求其request_id会被改写为src_rank_{源Rank}_{原req_id}前缀形式见trigger_load中new_req_id的构造。recover方法负责将输出归还_split_rebalance_outputs从LLMEngine.step输出中分离出带src_rank前缀的 rebalance 输出解析源 Rank 与原始 req_id抽取 token_ids、logprobs 等与本地正常输出get_rebalance_outputs通过dist.all_gather_object在 DP 组内同步取回属于本 Rank 的 rebalance 输出_build_rebalance_request_output将其还原为标准的RequestOutputCompletionOutputLogprob父子关系还原current_outputs_map以去除序号前缀后的request_id为 key 归并因为同一请求可能被拆成多个src_rank_*子请求分别完成将子输出挂到父RequestOutput.outputs下最终返回还原后的输出集合。def recover(self, outputs): # 将进行了rebalance迁移的outputs通过all_gather还原到源rank rebalance_outputs, current_outputs self._split_rebalance_outputs(outputs) for rebalance_output in self.get_rebalance_outputs(rebalance_outputs): rank_log_info(f[RecvReqId{rebalance_output[req_id]}], forceTrue) current_outputs.append(self._build_rebalance_request_output(rebalance_output)) current_outputs_map {} for request_output in current_outputs: request_id request_output.request_id if _ in request_id: request_id request_id.split(_, 1)[-1] parent_request current_outputs_map.get(request_id) if not parent_request: parent_request RequestOutput( request_idrequest_id, promptNone, prompt_token_idsrequest_output.prompt_token_ids, prompt_logprobsNone, outputs[], finishedTrue) current_outputs_map[request_id] parent_request parent_request.outputs request_output.outputs return current_outputs_map.values()此外enable_rollout_rebalance中的hook.after(LLMEngine, step)会把 step 输出中带src_rank前缀的RequestOutput暂存到States.outputs_cache待全部迁移请求完成后统一参与结果还原避免中间输出被打散。而LLM._run_engine前 hook 中还会缓存首个请求的sampling_paramsStates.sampling_params供迁移加载时copy复用采样策略。4.6 检测节流与性能剖析check()由LLMEngine.step的 after hook 每次触发但内部通过CheckCounterutils.py控制频率check_interval步做一次 Rebalance 检测、profile_interval步做一次性能采样且二者按阈值各自计数、互不阻塞从而控制 all_gather / all_to_all 的通信开销。profile()打印BSMap各 Rank 映射后的档位、SeqCntMap各 Rank 剩余序列数、CurrentMaxBS、ProfileStepCost、TPOT由step 耗时 * 1000 / threshold估算并在最大档位发生变化时输出MaxBSChanged提示。日志默认只在 rank 0 打印rank_log_info的forceFalse分支排障时可用forceTrue观察每个 Rank 的收发细节。5. 使能效果长尾场景下的实测数据关联文档在Atlas A3 集群 64 卡环境上进行了验证实验开启本方案后单轮推理耗时从约 10200s 优化到约 6100s性能收益达 60% 左右。实验配置如下配置项值模型Qwen3 235B数据集deepscalerdata.train_batch_size512data.max_response_length32768actor_rollout_ref.rollout.n16TP / DPTP4DP32性能收益主要来源于单个 step 的 TPOTTime Per Output Token差距默认场景下长尾阶段的 TPOT 会从 125ms 上升到 200ms大 BatchSize 档位图处理少量剩余序列算力利用率低使能 Rebalance 并配合多档位编图后能在 1~2K 推理长度时就快速将推理档位降低让单个 step 的 TPOT 降低到60ms 的量级长尾阶段持续的时间越长如max_response_length32768的极端长尾任务性能差距被持续放大总收益越显著。仓库中 llm_rl/deepseek/verl_patches/features/rollout_optimize/README.md 还记录了 DeepSeek 场景的另一组对照实验Atlas A3 128 卡、DeepSeekV3、open-r1/OpenR1-Math-220K、TP2/DP128单轮推理耗时从约 6200s 优化到约 2300s性能收益为 57%~62%与 Qwen3 场景结论一致。6. 适用前提与使用建议适用场景On-PolicyRLHF/GRPO 等训练中单轮推理的同步 Rollout 场景尤其是max_response_length大、响应长度长尾明显的任务推理、代码、数学等对于响应长度接近均匀分布的场景Rebalance 收益有限。前置依赖vllm_ascend 需支持torchair_graph_config中的use_cached_graph与graph_batch_sizes即多档位预编图能力multi_graph应保持开启以获得最大收益。档位设计graph_batch_sizes的档位应根据 DP 组内每 Rank 的平均序列数设计最大档位会被引擎max_batch_size自动钳制档位越低长尾阶段的单步耗时越低但预编译图的数量和显存占用也会增加。检查频率check_interval需在调度及时性与 all_gather/all_to_all 通信开销之间权衡建议结合profile输出的[RebalanceSchedule][Costxx ms]评估。迁移日志[RebalanceScheduleTasks]、[Rebalance][ReqId][Src][Dst]、[ReceivedTask/KvCache]、[MaxBSChanged]等日志可用于确认调度是否按预期发生以及各 Rank 剩余序列数[Rankx, Remainy]是否收敛到目标档位。7. 相关资源特性文档docs/features/rollout_rebalance.md核心实现deepseekllm_rl/deepseek/verl_patches/features/rollout_optimize/rollout_rebalance.py、config.py、utils.py特性说明deepseekllm_rl/deepseek/verl_patches/features/rollout_optimize/README.md使能入口deepseekllm_rl/deepseek/verl_patches/workers/megatron_workers.py使能补丁与实现qwen30011-verl-feature-enable_rollout_rebalance.patch 及 llm_rl/qwen3/verl-mindspeed/patches/verl/features/rollout_optimize训练目录总览deepseekllm_rl/deepseek/README.mdverl_patches/features/rollout_optimize即本特性的部署目录同族 RL 特性文档docs/features/history_tree_dynamic_rl.md、docs/features/length_aware_resampler.md【免费下载链接】cann-recipes-train本项目针对LLM与多模态模型训练业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-train创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考