企业官网建设流程全解析

CANN 序列级均衡调度引擎：RL On-Policy 推理长尾场景的 Rollout Rebalance 实战解析

【免费下载链接】cann-recipes-train本项目针对LLM与多模态模型训练业务中的典型模型、加速算法，提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-train

导读

本文以 CANN（cann-recipes-train）仓库中 docs/features/rollout_rebalance.md 为骨架，系统讲解 RL On-Policy（RLHF/GRPO/DAPO 等）训练中 Rollout 阶段的"木桶效应"与长尾负载均衡问题。文章覆盖该特性的使能方式、全部配置项、基于 BatchSize 档位预设的最大档位最小化均衡算法、Request（SEQ）级跨 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_optimize，qwen3 场景位于 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), reverse=True): 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_init=False, use_cached_graph=True, graph_batch_sizes=graph_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_ENABLE=1使能。在 deepseek 场景中，llm_rl/deepseek/verl_patches/workers/megatron_workers.py 的actor_rollout_ref_init_model（对应 verl 原版ActorRolloutRefWorker.init_model的补丁实现）开头已内置使能逻辑：

@register(dispatch_mode=Dispatch.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.py：

class 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 # 打印间隔步长

各配置项的作用与建议：

4. 实现方案：序列级均衡调度引擎

本优化的核心目标：在 On-Policy 场景中，针对部分 Rollout 提前结束导致各 Rank 间负载不均时，对未结束的 Rollout 进行负载均衡的策略分析和重调度，从而提升计算资源的利用率和长尾状态下的推理效率。

方案包含三大关键功能：

1. Rebalance 条件检测与调度策略生成：周期性采集各 Rank 全局状态，计算最优请求迁移清单；
2. Request（SEQ）级的数据搬迁与恢复（含对应 KV Cache）：将序列连同其 KV Cache 在 DP 组内跨 Rank 迁移；
3. Rollout 后的结果还原：将迁移到其他 Rank 完成的序列输出，通过 all_gather 归还给源 Rank，并还原被临时拆散的父子请求关系。

以下代码均位于 llm_rl/deepseek/verl_patches/features/rollout_optimize/rollout_rebalance.py（qwen3 版本路径为 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( rank=self.rank, req_ids=list(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(), group=self.dp_group) return group_states

值得注意的是，RolloutRebalanceEngine.__init__中会先执行一步关键预处理：移除 request 的父子关系（req_state.parent_req = None）。这是因为 Rebalance 需要让单个 Seq 能够独立搬迁；被迁移序列的父子关系会在最后的recover阶段统一还原。

4.2 基于 BatchSize 档位预设的最大档位最小化均衡算法

calc_balancing_tasks负责计算最优的请求迁移任务清单。其最优策略定义为三层目标：

1. 主要目标：将整个 DP 组所需的最大档位（max_bs）降至最低；
2. 次要目标：在满足 1 的前提下，使迁移的请求数量（cost）最少；
3. 补充目标：在满足 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(rank=rank, surplus=delta, req_ids=rank_req_cnt_map[rank][:delta])) elif delta < 0: receivers.append(dict(rank=rank, capacity=-delta)) # 生成迁移任务清单：为了让迁移速度更快，尽可能将请求均衡分发给目标worker， # 但优先分发给最空闲的worker（即容量最大的worker） balancing_tasks = [] while True: donor_index = 0 for receiver in sorted(receivers, key=lambda r: r['capacity'], reverse=True): donor = donors[donor_index] num_to_move = min(donor['surplus'], receiver['capacity']) balancing_tasks += [dict(from_rank=donor['rank'], to_rank=receiver['rank'], req_id=req_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-robin，donor_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_cache=False），只有剩余无等待态请求可迁时，才动用正在运行（running）的请求（need_kv_cache=True，需搬迁 KV Cache）。这一分层搬迁思路在现版本中通过load_by_req_id统一采集 KV Cache 实现。

4.3 序列请求的跨 Rank 发送与接收（含 KV Cache）

得到迁移任务清单后，all_to_all_v_tasks执行实际的序列搬迁。核心流程为：

1. 采集并 abort：对属于本 Rank 的每个迁移任务，用RebalanceRequestTask.load_by_req_id采集该请求的完整状态（prompt 与已生成 token、logprobs 缓存、KV Cache 块），并调用trigger_abort中止其在本地引擎中的执行；
2. 元数据 all_to_all：将每个目标 Rank 要接收的请求元数据（get_transfer_dict，经 pickle 序列化后转为 uint8 tensor）通过dist.all_to_all_single交换——先交换各 Rank 的数据尺寸，再交换实际数据；
3. KV Cache 直传：send_kv_caches对每个需要 KV Cache 的请求，逐层逐块dist.send；接收侧load_received_tasks中按形状dist.recv接收；
4. 加载还原：load_by_transfer_info+trigger_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), dtype=torch.uint8).to(self.device) for data in objects_to_send] local_sizes = torch.tensor([len(t) for t in tensor_list], dtype=torch.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(), dtype=torch.uint8).to(self.device) dist.all_to_all_single( output_tensor, input_tensor, output_split_sizes=remote_sizes.tolist(), input_split_sizes=local_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, dst=to_rank) rank_log_info(f'[TaskSendKvCache][ToRank={to_rank}][ReqId={request_task.req_id}]')

接收侧在load_received_tasks中按layers_kv_cache_shapes预分配空 tensor 后dist.recv，并通过日志（force=True强制打印）记录ReceivedTask、ReceivedKvCache、ReceivedTaskLoaded三个阶段的耗时（毫秒），便于定位搬迁开销。注意源码中接收侧使用.cuda()申请设备内存（States.reprefill_mode为 False 时）；States.reprefill_mode为 True 时跳过 KV Cache 搬迁，走重新 prefill（reprefill）路径——对应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_prefilling=False）与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方法负责将输出归还：

1. _split_rebalance_outputs：从LLMEngine.step输出中分离出带src_rank前缀的 rebalance 输出（解析源 Rank 与原始 req_id，抽取 token_ids、logprobs 等）与本地正常输出；
2. get_rebalance_outputs：通过dist.all_gather_object在 DP 组内同步，取回属于本 Rank 的 rebalance 输出；
3. _build_rebalance_request_output：将其还原为标准的RequestOutput（CompletionOutput+Logprob）；
4. 父子关系还原：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"]}]', force=True) 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_id=request_id, prompt=None, prompt_token_ids=request_output.prompt_token_ids, prompt_logprobs=None, outputs=[], finished=True) 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_params（States.sampling_params），供迁移加载时copy复用采样策略。

4.6 检测节流与性能剖析

• check()由LLMEngine.step的 after hook 每次触发，但内部通过CheckCounter（utils.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的force=False分支），排障时可用force=True观察每个 Rank 的收发细节。

5. 使能效果：长尾场景下的实测数据

关联文档在Atlas A3 集群 64 卡环境上进行了验证实验，开启本方案后单轮推理耗时从约 10200s 优化到约 6100s，性能收益达 60% 左右。

实验配置如下：

性能收益主要来源于单个 step 的 TPOT（Time Per Output Token）差距：

• 默认场景下，长尾阶段的 TPOT 会从 125ms 上升到 200ms（大 BatchSize 档位图处理少量剩余序列，算力利用率低）；
• 使能 Rebalance 并配合多档位编图后，能在 1~2K 推理长度时就快速将推理档位降低，让单个 step 的 TPOT 降低到60ms 的量级；
• 长尾阶段持续的时间越长（如max_response_length=32768的极端长尾任务），性能差距被持续放大，总收益越显著。

仓库中 llm_rl/deepseek/verl_patches/features/rollout_optimize/README.md 还记录了 DeepSeek 场景的另一组对照实验（Atlas A3 128 卡、DeepSeekV3、open-r1/OpenR1-Math-220K、TP=2/DP=128），单轮推理耗时从约 6200s 优化到约 2300s，性能收益为 57%~62%，与 Qwen3 场景结论一致。

6. 适用前提与使用建议

• 适用场景：On-Policy（RLHF/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][Cost=xx ms]评估。
• 迁移日志：[RebalanceScheduleTasks]、[Rebalance][ReqId][Src][Dst]、[ReceivedTask/KvCache]、[MaxBSChanged]等日志可用于确认调度是否按预期发生，以及各 Rank 剩余序列数（[Rank=x, Remain=y]）是否收敛到目标档位。

7. 相关资源

• 特性文档：docs/features/rollout_rebalance.md
• 核心实现（deepseek）：llm_rl/deepseek/verl_patches/features/rollout_optimize/rollout_rebalance.py、config.py、utils.py
• 特性说明（deepseek）：llm_rl/deepseek/verl_patches/features/rollout_optimize/README.md
• 使能入口（deepseek）：llm_rl/deepseek/verl_patches/workers/megatron_workers.py
• 使能补丁与实现（qwen3）：0011-verl-feature-enable_rollout_rebalance.patch 及 llm_rl/qwen3/verl-mindspeed/patches/verl/features/rollout_optimize
• 训练目录总览（deepseek）：llm_rl/deepseek/README.md（verl_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