ARTICLE · INTELLIGENCE

战地情报 · 详情页

来自尧图项目组的一线实战观察与深度解析

MoE推理负载均衡实战:EasyBalance跨层调度让专家错峰出行

MoE推理负载均衡实战:EasyBalance跨层调度让专家错峰出行 MoE 推理这几年从论文里的概念一路杀进生产环境速度比很多人预想的都快。但真正把模型部署到多卡、多机上的时候你会发现一个很尴尬的现象明明总显存够、总带宽够、单卡算力也够可端到端吞吐就是上不去GPU 利用率像心电图一样忽高忽低。问题往往不在模型本身而在专家Expert被路由到的频率天然不均——有的专家被挤爆有的专家闲得发慌。EasyBalance 这个工作瞄准的就是这件事它把负载均衡从单层内部调度抬升到跨层协同的视角用一句话概括就是让 MoE 推理学会错峰出行。这篇博文我会把 Cross-Layer Load Balancing 这件事拆开讲透包括它到底解决什么、为什么单层均衡不够用、跨层调度怎么落地、实测中会踩哪些坑以及你可以直接抄的配置思路。1. 先搞清楚 MoE 推理到底卡在哪1.1 稀疏激活带来的账面很美、实际很堵MoEMixture of Experts的核心卖点是用稀疏激活换参数量一个 Token 只走 Top-K 个专家理论上计算量只和激活的专家数相关而不是和总参数量相关。听起来是用更少的算力撬动更大的模型。但这里有个被很多人忽略的前提——负载必须均匀。现实是路由网络Router/Gate是根据 Token 的语义特征来分发的而真实语料里语义分布极度不均衡。某些通用高频模式比如标点、常见虚词、高频句式会反复命中同一批热门专家而一些冷门专家可能整批请求都没被激活几次。结果就是热门专家所在的那张卡排队排到爆冷门专家所在的卡空转。你花钱买的是 N 张卡的算力实际有效利用率可能只有 40%~60%。我在实际压测里见过最夸张的情况8 卡部署一个中等规模 MoE理论吞吐能到 X实测只有 0.5X 出头用 nvtop 一看两张卡 95% 利用率另外几张长期在 20% 以下晃悠。这不是模型问题是调度问题。1.2 单层负载均衡的天花板传统做法是在每一层内部做均衡常见手段有这么几类容量因子Capacity Factor给每个专家设一个最大 Token 容量超了就丢弃或走残差。简单粗暴但会掉精度。辅助损失Auxiliary Loss训练时加一个负载均衡损失逼 Router 把 Token 分得更均匀。这是训练期手段推理期管不了已经训好的路由偏好。Token 重排 / All-to-All 优化在通信层面把发往同一专家的 Token 打包减少通信次数。这些手段在单层内确实有效但它们有个共同的盲区只看当前层不看全局。MoE 模型通常有几十层每层都有自己的 Router。第 3 层的热门专家和第 27 层的热门专家很可能落在同一张卡上因为专家到设备的映射是静态的于是这张卡在每一层都被反复压榨而另一张卡在每一层都吃不饱。单层均衡再怎么调也解决不了这种跨层叠加的设备级热点。这就是 EasyBalance 要解决的核心矛盾均衡的粒度不能停在层内专家之间必须上升到跨层设备之间。1.3 为什么叫错峰出行错峰出行这个比喻很贴切。城市早高峰堵车不是因为路不够宽而是因为所有人都在同一时间挤同几条路。解决办法不是把路修得更宽加卡而是把出行时间错开跨层调度。MoE 推理里每一层的计算和通信其实是有先后顺序的。如果第 3 层把某张卡压满而第 27 层本来也要压这张卡那我们能不能在第 27 层把一部分 Token 引导到别的卡上让两张卡的负载在时间轴上错开EasyBalance 的 Cross-Layer 思路正是基于这个观察不同层的负载高峰在时间上是可分离的只要调度器能看到全局就能把热点摊平。2. Cross-Layer 负载均衡的核心机制拆解2.1 从静态专家映射到动态跨层视图大多数 MoE 推理框架里专家到设备的映射是静态的专家 0~7 在卡 0专家 8~15 在卡 1以此类推。这个映射在编译期或加载期就定死了推理时不会变。静态映射的好处是实现简单、通信模式可预测坏处是它完全无视运行时的实际负载分布。EasyBalance 的第一步是建立一个跨层的全局负载视图。具体来说它在每一层推理时收集每个专家的实际 Token 数、每张卡的实际排队深度、以及 All-to-All 通信的实时带宽占用然后把这些信息汇总成一个跨层的负载矩阵。这个矩阵的维度大致是层数 × 设备数每个元素代表该层该设备的压力值。有了这个矩阵调度器就能回答一个关键问题当前这批 Token 在第 L 层应该优先发给哪些设备才能让整条流水线的设备负载最平注意这里的目标不是让第 L 层最平而是让所有层加起来最平。这是 Cross-Layer 和单层均衡的本质区别。2.2 调度决策的三个输入维度EasyBalance 的调度器在做决策时主要看三个维度维度含义为什么重要层内专家热度当前层各专家的 Token 命中数决定单层内的即时压力跨层设备累积压力该设备在前后若干层的总负载避免设备级热点叠加通信代价Token 发往目标设备的 All-to-All 开销均衡不能以牺牲通信为代价这三个维度是有优先级的。我的理解是通信代价是硬约束跨层累积压力是主目标层内热度是即时修正项。因为如果你为了均衡把 Token 发到一个通信代价极高的设备上省下的计算时间还不够补通信的得不偿失。实际实现里这三个维度会被加权成一个综合代价函数调度器选择代价最小的设备作为 Token 的目标。权重怎么定这取决于你的硬件拓扑。NVLink 全互联的机器通信权重可以调低跨机走网络的通信权重必须调高。这个后面在实操部分会细说。2.3 为什么跨层比层内更难也更有价值层内均衡是个局部优化问题每层独立求解实现简单但容易陷入局部最优、全局次优。跨层均衡是个全局优化问题难度高得多因为状态空间大层数 × 设备数的组合搜索空间随规模指数增长。时序耦合第 L 层的决策会影响第 L1 层的可用容量不能只看当前。实时性要求推理是延迟敏感的调度器不能花几十毫秒去算一个最优解。EasyBalance 的价值就在于它用近似算法 增量更新把这个问题压到了可接受的延迟内。它不追求每批 Token 的全局最优而是维护一个滚动更新的负载视图每批只做局部调整让系统整体趋近均衡。这种够用就好的工程取舍恰恰是它能落地的关键。3. 落地实操把 EasyBalance 思路接进你的推理栈3.1 环境准备与前置检查在动手之前先确认你的推理栈满足几个前提。这些不是 EasyBalance 独有的要求而是任何跨层调度方案都需要的专家映射可配置你的框架必须允许在运行时调整专家到设备的映射或者至少允许在 Token 分发时指定目标设备。如果专家映射是硬编码在 CUDA kernel 里的那跨层调度无从谈起。有可观测的负载指标你需要能拿到每层每设备的 Token 数、队列深度、通信耗时。大多数推理框架如 vLLM、TensorRT-LLM、DeepSpeed-Inference都提供了部分指标但可能需要自己加埋点。All-to-All 通信可控跨层调度会改变 Token 的流向通信模式必须能跟着变。如果你的 All-to-All 是固定 buffer 的需要改成动态的。提示如果你的框架暂时不支持动态专家映射可以先从跨层负载观测做起把负载矩阵打出来看看热点到底在哪。很多时候光是看清问题就能靠调整专家初始映射解决一大半。3.2 负载矩阵的采集与更新这是整个方案的地基。我建议用一个轻量的环形缓冲区来存负载矩阵每层推理完就更新对应行。伪代码大概长这样# 跨层负载矩阵shape [num_layers, num_devices] load_matrix np.zeros((num_layers, num_devices)) def update_load(layer_id, device_id, token_count, comm_cost): # 指数移动平均避免单批波动导致调度抖动 alpha 0.3 load_matrix[layer_id][device_id] ( alpha * (token_count comm_cost) (1 - alpha) * load_matrix[layer_id][device_id] ) def get_cross_layer_pressure(device_id, current_layer, window5): # 看当前层前后 window 层的累积压力 start max(0, current_layer - window) end min(num_layers, current_layer window) return load_matrix[start:end, device_id].sum()这里有几个经验点用 EMA 而不是瞬时值单批 Token 分布波动很大如果直接用瞬时值调度会导致调度器疯狂抖动反而降低性能。EMA 的 alpha 取 0.2~0.4 比较稳。窗口大小要调window 太小退化成单层均衡太大则反应迟钝。我实测 5~8 层是个不错的区间具体看你的模型层数。通信代价要归一化token_count 和 comm_cost 量纲不同必须先归一化再加权否则通信项会被淹没。3.3 调度决策的工程实现采集到负载矩阵后调度器要在每层分发 Token 时做决策。核心逻辑是对每个 Token计算它发往各候选设备的综合代价选最小的。def select_device(token, layer_id, candidate_experts): best_device None best_cost float(inf) for expert in candidate_experts: device expert_to_device[expert] # 三个代价项 intra_cost current_layer_load[layer_id][device] cross_cost get_cross_layer_pressure(device, layer_id) comm_cost estimate_comm_cost(token, device) # 加权求和权重需按硬件拓扑调 cost 0.3 * intra_cost 0.5 * cross_cost 0.2 * comm_cost if cost best_cost: best_cost cost best_device device return best_device权重0.3 / 0.5 / 0.2只是我常用的起点不是标准答案。调权重的原则是机器内 NVLink 全互联comm_cost 权重可以降到 0.1cross_cost 提到 0.6。跨机网络comm_cost 权重必须提到 0.4 以上否则均衡省下的时间全被通信吃掉。层数特别深60 层cross_cost 的窗口要相应放大否则跨层信号太弱。注意调度决策本身有开销。如果你的 batch 很小、单层计算只有几毫秒调度器算一次要 1ms那就得不偿失。这种情况下应该批量决策——一次算好一批 Token 的目标设备而不是逐 Token 算。3.4 和现有推理框架的对接方式EasyBalance 是个思路不是一个必须整体替换的框架。你可以按侵入性从低到高选择接入方式接入方式侵入性适用场景实现难度只做负载观测极低先定位热点低调整专家初始映射低热点固定且可预测低动态 Token 重路由中负载波动大中全跨层调度高大规模多机部署高我的建议是从低侵入性方案起步。先用负载观测跑一周看看热点是固定的还是漂移的。如果热点固定直接调专家映射就够了根本不需要动态调度。只有热点随时间漂移、且单层均衡压不住的时候才值得上跨层调度。4. 实测中的坑与排查链路4.1 调度抖动越均衡越慢的怪现象我第一次接跨层调度时遇到一个反直觉的问题开了调度之后吞吐不升反降。排查过程是这样的第一步看 GPU 利用率。发现利用率确实更平了但整体水位下降了。这说明均衡生效了但引入了额外开销。第二步看调度器耗时。加了个计时器发现每层调度花了 2~3ms而单层计算只有 5ms。调度开销占了 40%这就是元凶。第三步定位根因。逐 Token 调度 每 Token 都重算负载矩阵计算量爆炸。解决方案是批量决策 降低更新频率每 8 个 Token 决策一次负载矩阵每层只更新一次而不是每 Token 更新。改完之后调度开销降到 0.3ms吞吐立刻回升并超过了基线。这个坑的教训是均衡的收益必须大于调度的成本。在小 batch、浅层模型上跨层调度可能根本不划算。4.2 通信放大均衡把 Token 打散了第二个坑更隐蔽。跨层调度为了均衡会把原本发往同一设备的 Token 拆到多个设备。结果 All-to-All 的通信包变小、变多通信效率暴跌。排查方法是对比调度前后的 All-to-All 耗时。如果通信耗时涨了 30% 以上说明 Token 被打得太散。解决办法是加一个聚合约束调度时优先把 Token 发往已经积累了足够多同目标 Token 的设备避免为了微小均衡收益而拆散通信包。具体做法是在代价函数里加一项目标设备当前批次已有 Token 数的负向奖励——已有 Token 越多边际通信成本越低。4.3 冷启动阶段的负载矩阵失真系统刚启动时负载矩阵全是 0调度器没有信息可依据前几批 Token 的调度基本是随机的。如果这几批恰好赶上大请求就会造成明显的延迟毛刺。我的处理方式是预热启动后先用一批合成请求或者真实请求的前若干批跑一遍把负载矩阵填起来再开启调度。预热批次不用多通常 3~5 批就能让矩阵收敛到可用状态。4.4 专家映射变更后的缓存失效如果你用的是动态专家映射专家会在设备间迁移要注意 KV Cache 和专家权重的缓存会失效。专家迁移一次相关缓存全部作废重新加载权重的开销可能高达几十毫秒。所以专家迁移的频率必须严格控制通常只在负载严重失衡比如某设备压力是均值 2 倍以上时才触发且迁移要选在请求间隙做。5. 参数调优与效果验证5.1 关键参数速查表参数推荐起点调整方向影响EMA alpha0.3波动大调小反应慢调大调度稳定性跨层窗口 window5~8层数深调大跨层信号强度intra 权重0.3单层热点强调大即时均衡cross 权重0.5设备热点强调大全局均衡comm 权重0.2跨机部署调大通信开销决策批大小8调度开销大调大调度成本5.2 怎么判断均衡真的生效了不要只看端到端吞吐那个指标太粗。我一般看三个细粒度指标设备利用率标准差跨层调度生效后各卡利用率的方差应该明显下降。如果方差没降说明调度没起作用。All-to-All 耗时分布均衡后通信耗时的 P99 应该更接近均值长尾减少。每层热点设备的变化如果热点设备在层间是漂移的说明跨层调度在起作用如果热点始终固定在几张卡上说明调度力度不够。5.3 一个真实的调优案例我拿一个 32 层的 MoE 模型在 8 卡上做过对比。基线无跨层调度的端到端吞吐是 100%归一化设备利用率标准差 0.28。开启跨层调度后第一版逐 Token 调度吞吐 85%标准差 0.15。均衡了但变慢了。第二版批量决策 EMA吞吐 118%标准差 0.12。这才是想要的结果。第三版加通信聚合约束吞吐 127%标准差 0.11。通信优化带来额外收益。这个案例说明跨层调度的收益不是自动来的必须把调度开销和通信开销压下去均衡的收益才能释放出来。很多人第一次尝试失败就是卡在开销这一关。6. 这套思路还能往哪延伸Cross-Layer 负载均衡的本质是用全局视角做局部决策这个思想不局限于 MoE 推理。我最近在几个方向上做了延伸尝试效果还不错第一个方向是训练期的跨层均衡。训练时同样存在专家热点而且训练是长跑热点叠加的累积效应更严重。把跨层负载视图接进训练循环动态调整辅助损失的权重比固定权重的辅助损失更有效。第二个方向是多租户场景的跨层隔离。多个请求流共享一个 MoE 集群时不同租户的语义分布不同热点也不同。跨层调度可以顺便做租户隔离避免一个租户的热门专家拖垮其他租户。第三个方向是和 PagedAttention 类内存管理的结合。KV Cache 的分页管理和专家负载均衡其实有协同空间——如果某个设备的专家压力大可以把它的 KV Cache 分页迁到压力小的设备上进一步摊平内存和计算压力。最后分享一个我踩过的小坑跨层调度的负载矩阵如果只统计 Token 数会忽略不同 Token 的计算量差异。长序列的 Token 和短序列的 Token走同一个专家计算量可能差好几倍。后来我在负载矩阵里加了序列长度加权均衡效果明显更准。这个细节很多论文不会写但实际部署时很关键。如果你也在做 MoE 推理部署建议先把负载观测做起来把热点看清楚。很多时候问题没你想的那么复杂一张负载热力图就能告诉你答案。跨层调度是重武器用之前先确认轻武器是不是已经够用了。
RELATED READING

延伸阅读

更多一线实战笔记与深度复盘,助您持续精进