这几个月我一直在折腾一个 MoE 模型的分布式推理服务模型效果确实没得说但上线之后最让人头疼的反而不是算子优化而是负载均衡。MoEMixture of Experts推理的负载均衡问题和传统 Dense 模型完全不是一回事你盯着 GPU 平均利用率看半天根本看不出问题因为所有不均衡都藏在专家路由的动态分布里。后来我把调度视角从单层拉到了跨层参考了 EasyBalance 这一套跨层负载均衡的思路才真正把整体吞吐稳住。这套思路说到底就四个字错峰出行。如果你正在做 MoE 推理引擎的开发、推理服务的性能调优或者只是好奇为什么 MoE 部署起来这么费劲这篇内容应该对你有用。我会把负载失衡的表现、跨层均衡的核心逻辑、显存问题以及一个可以直接跑起来的调度原型都串起来讲尽量说人话。1. MoE 推理的负载均衡问题到底出在哪1.1 专家路由的“二八定律”MoE 模型里每个 token 不是经过所有专家而是通过一个路由网络Router选出 Top-2 或 Top-k 个专家来处理。这个机制让模型可以在总参数量很大的情况下只激活一小部分参数理论上推理成本应该接近 Dense 模型。但现实是路由结果从来都不是均匀分布的。不同专家被选中的概率差距非常大有些专家几乎是“万能工具人”几乎所有 token 都会路过去看一眼有些专家则是“冷宫选手”半天轮不到一个 token。我在实际日志里统计过一个小规模 MoE 模型8 层、每层 32 个专家的专家访问分布热门专家和冷门专家的访问量能差出 2 到 3 倍。这不是训练没练好而是数据本身的分布决定的语言任务里的常见模式高度集中路由网络学出来的结果自然也是偏向某些专家的。训练阶段有负载均衡辅助损失aux loss / load balance loss压着整体不会太离谱但推理阶段面对的真实流量和训练数据分布有偏差比如线上突然出现一批特定格式的长文档请求某些专家就会瞬间被打爆。这里有个很容易被忽视的点模型推理框架不会因为你某个专家负载高就主动帮你分流。它只是按照路由结果把 token 送给对应的专家专家计算完之后再聚合回去。结果就是热门专家所在的设备长期排队冷门专家所在的设备闲着没事干但所有 GPU 的平均利用率看起来还挺“健康”。1.2 分布式推理中负载失衡的三种典型表现分布式 MoE 推理的负载问题比单机单卡要复杂得多。我总结下来失衡通常以三种形式暴露出来。第一种是显存占用失衡。MoE 模型的所有专家权重必须分布在多张 GPU 上每个设备负责一部分专家。专家权重本身是静态的不会因为访问量变化而改变显存占用所以显存失衡主要是由于 KV cache 和临时激活值的分布不均。当某些设备上的专家被频繁调用时对应的 batch 会被撑大激活值和中间结果的显存占用会明显偏高极端情况下直接 OOM。第二种是计算排队失衡。这是最直观的。热门专家所在设备上的计算队列一直在积压同一批请求里跑到这个专家上的 token 都要排队等 GPU 算完冷门专家所在设备却经常空转。在 continuous batching 场景下一个 batch 的完成时间取决于最慢的那个专家整体的端到端延迟就被热点拖住了。第三种是跨节点通信放大。分布式 MoE 推理里token 需要做 all-to-all 的 dispatch把自己送到目标专家所在的设备上算完再做 all-to-all 的 combine 拿回结果。如果同一时刻大量 token 都涌向同一批专家对应的通信报文会急剧膨胀通信耗时成为瓶颈。负载越不均衡通信放大越严重因为数据都往同一个方向挤。这三种失衡往往会同时出现而且会互相放大。你单独解决任何一个都不够必须从全局视角去调整。2. 为什么说“跨层负载均衡”是错峰出行的关键2.1 从单层均衡到全局均衡传统思路是在每一层内部做负载均衡比如调整路由策略、给专家加权、或者迁移专家副本。这种单层均衡有用但天花板很低因为它只解决了“这一层内专家的冷热不均”没有解决“所有层同时出现热点”的问题。设想一个场景某个时间段来了大量特定领域的长文本请求它们在第 3 层路由到专家 A在第 6 层路由到专家 B。单层均衡只能保证第 3 层的 A 不会被打死、第 6 层的 B 不会被打死但如果 A 和 B 恰好在同一张卡上或者虽然在不同卡上但同时到达计算峰值全局还是堵住了。更关键的是这种拥堵往往是同频共振的——请求在每一层都要做一次路由同一个请求的高负载专家虽然在每层不同但时间上它们是先后出现的如果每个 batch 在每个层都碰到热点整个流水线每一站都在排队。单一层视角解决不了“所有层同时顺峰”的问题。这也是我一开始调优效果很差的原因把每一层的专家负载都压平了整体吞吐并没有明显提升因为瓶颈从单个专家转移到了全局节奏上。2.2 不同层的专家热度天然互补后来我意识到一个非常重要的现象不同层的专家热度往往具有天然的互补性。MoE 模型的不同层承担着不同的特征抽象任务靠近输入的低层专家更偏向处理通用语法、词法、局部 n-gram 模式靠近输出的高层专家更偏向处理任务语义、全局上下文和推理逻辑。线上流量在不同层次上的分布并不是同构的某些请求可能低层打得很散、高层很集中另一些请求则反过来。这意味着在同一个时间段内真正同时处于峰值负载的层其实没有那么多。如果把每个层独立看各自都有顺峰的时候但从全局时间轴上看不同层的高峰期是可以错开的。这就是“错峰出行”的物理基础。另外Transformer 的推理是逐层串行的同一个请求在任意时刻只会在某一个层内计算。这给跨层调度提供了天然的操作窗口当第 3 层的专家很忙时第 5 层的专家可能正闲着如果调度器能感知到层间的负载差异完全可以通过微调 micro-batch 的流转节奏让同一个计算设备在时间上承接不同层的工作而不是同时被所有层的热点冲击。2.3 EasyBalance 的核心设计逻辑EasyBalance 的思路本质上是把负载均衡从空间维度扩展到时间维度。空间维度是指“专家分布在哪张卡上”时间维度是指“请求在某一刻访问哪些专家”。它最核心的三个设计点是第一个是跨层负载感知。调度器不能只看当前层的专家排队长度而是要看请求未来在后续层会访问的专家列表。通过统计每层每个专家的实时负载和近期趋势构建一张“全局拥堵地图”。这张地图能告诉我们请求当前在第 3 层是畅通的但第 6 层要经过一个热门专家第 8 层还要经过另一个热门专家。第二个是全局错峰调度。根据拥堵地图调度器不再简单地把请求按到达顺序直接送出而是对多个 micro-batch 的执行顺序做调整。如果某一个 batch 的前几层负载很低、但后几层会命中多个热点那就不急着把它往后推如果一个 batch 的前几层负载很高就让它稍微等待一下先放那些前层负载低的请求过去。这样不同请求在不同层的“高峰段”就不会叠加在同一时刻。第三个是动态资源腾挪。跨层均衡不只是在请求调度层做文章还可以影响底层资源的分配策略。比如在推理过程中调度器可以根据层间负载的实时变化调整每一层的计算并行度、micro-batch 大小、甚至专家权重的显存驻留策略。对于当前峰值所在的层适当降低并发注入率避免排队加剧对于当前低峰的层则可以加大注入率把空闲算力用满。打个比方交通拥堵治理的思路不是简单把每条路都修宽而是通过红绿灯联动、错峰出行、动态诱导让车流在时间和空间上均匀分布。EasyBalance 做的就是 MoE 推理里的“交警调度”。3. 解惑MoE 架构要全部参数进显存吗3.1 总参数和激活参数先分清很多刚接触 MoE 的人会问既然一个 token 只激活少数专家是不是只用把激活专家放到显存里就行了答案很直接推理时要跑得顺所有专家参数都得放在能从显存快速访问到的地方。这不是因为每次前向都会用到所有专家而是因为你根本预判不了下一个 token 会路由到哪个专家。MoE 模型有两个关键数字总参数Total Parameters和激活参数Activated Parameters。总参数包括所有专家权重激活参数是单个 token 前向时实际参与计算的参数。以某个总参数 47B、激活参数 7B 级别的开源 MoE 模型为例它的理论推理效率接近 7B Dense 模型但权重存储成本却接近 Dense 模型的六到七倍。这就是 MoE 的“甜蜜陷阱”计算效率高但存储和显存压力巨大。3.2 显存估算公式与典型结果推理阶段显存占用主要由三块构成模型权重、KV cache、激活值和临时计算空间。模型权重部分如果以 FP16 精度部署 47B 总参数的 MoE 模型光权重就需要约 94GB 显存单张 80GB 的卡放不下必须要多卡拆分。即使用 8 卡每张卡也要承担约 12GB 的权重再加上 KV cache 和其他开销显存其实非常紧张。KV cache 部分MoE 模型的 KV cache 和 Dense 模型的计算方式一样和专家数无关只和层数、隐藏维度、序列长度、batch 大小有关。粗略公式是KV cache 字节数 ≈ 2 * 层数 * 隐藏维度 * 精度字节 * 总 token 数假设一个 MoE 模型是 8 层、隐藏维度 4096、FP16 精度每个 token 的 KV cache 大概是 128KB。batch 为 32、每条序列 4096 token 时总 token 数是 131072KV cache 容量大约是 16GB。这部分会随着并发数线性增长非常可观。激活值和临时空间这部分和 batch 大小、序列长度强相关。MoE 的 all-to-all 通信也要在显存里暂存中间张量专家并行切分不同这部分内存差异很大一般来说预留权重的 15%~25% 比较保险。所以结论很清楚MoE 在推理时确实要把全部参数纳入显存规划这和“每个 token 真正用到多少参数”是两码事。3.3 缓解显存压力的几个方向显存问题虽然躲不掉但可以缓解。常见的方案有专家权重按需驻留。把一部分冷门专家放到 CPU 内存甚至 SSD 上遇到请求时再换入显存。这个方案的代价是冷门专家的首次访问延迟会很高可能到几十毫秒甚至上百毫秒。这时候就需要跨层负载均衡来兜底——调度器会尽量避免把多个需要换入冷门专家的请求同时送进来否则显存换入风暴会比计算排队更可怕。KV cache 单独管理。KV cache 不依赖专家可以优先保证核心推理链路的容量。通过 PagedAttention 这类显存管理策略减少 KV cache 碎片化让有效的 KV cache 容量提升。量化与权重压缩。FP16 降到 INT8 或者 INT4 可以显著降低权重显存。MoE 模型量化后的质量损失通常可控但需要针对专家层做细致的敏感度分析不能一刀切。跨层预取。调度器根据请求路径提前预判后续层的热门专家提前把可能用到的冷门专家权重换入显存把换入开销藏到计算时间背后。这个方向做得好可以大幅降低专家 offload 带来的延迟惩罚。4. 实操一个跨层负载均衡调度器原型4.1 设计目标与整体架构我参考 EasyBalance 思路做了一个简化的原型调度器目的是验证“跨层错峰”是否真的能降低端到端排队延迟。这个原型不依赖具体推理框架只模拟核心调度逻辑。假设集群上有 8 层、每层 32 个专家分布在多个设备上。请求会携带一个 path表示它经过每一层时选中的专家编号。调度器需要做两件事实时统计每个层每个专家的负载用滑动平均预测未来一段时间的热度。对请求队列做排序让多个正在执行的请求不要在同一时间撞上同一批热点专家。原型的输入是模拟生成的一批带路由路径的请求输出是调度后的执行顺序和模拟完成时间。4.2 关键数据结构和核心代码我用 Python 写了一个简化的调度器代码可以在本地直接运行方便验证思路。import heapq import numpy as np from collections import deque class SimpleLoadBalancer: def __init__(self, num_layers8, num_experts32, window20, alpha0.3, wait_budget5): self.num_layers num_layers self.num_experts num_experts self.window window self.alpha alpha self.wait_budget wait_budget # 记录每个专家最近 window 次的负载到达量 self.history np.zeros((num_layers, num_experts, window)) self.avg_load np.zeros((num_layers, num_experts)) self.cur_load np.zeros((num_layers, num_experts)) self.t 0 def update(self, layer, expert, delta1.0): 每个时间步更新一次delta 表示新到达的负载量。 self.history[layer, expert, self.t % self.window] delta self.avg_load[layer, expert] ( self.alpha * delta (1 - self.alpha) * self.avg_load[layer, expert] ) self.cur_load[layer, expert] self.history[ layer, expert, self.t % self.window ] def step(self): 推进一个时间步滑动窗口归零。 self.t 1 def congestion(self, path): 计算一条请求路径的总拥塞度越高表示容易堵。 score 0.0 for layer, expert in path: # 用平均负载和最近到达量加权越热点越容易拥塞 score self.avg_load[layer, expert] * 0.7 score self.cur_load[layer, expert] * 0.3 return score def schedule(self, requests, max_batch): requests: list of (request_id, path) 返回调度后的执行顺序。 scored [] for rid, path in requests: score self.congestion(path) # score 越大越应该延后用等待预算做上限约束 heapq.heappush(scored, (score, rid, path)) output [] # 简单起见按分数从低到高执行分数低的前层负债低先放行 while scored: score, rid, path heapq.heappop(scored) output.append((rid, path, round(score, 3))) return output def simulate_execute(self, requests, max_batch8): 模拟执行按批次发送每个请求的耗时为路径上各专家负载之和。 ordered self.schedule(requests, max_batch) total_latency 0.0 per_request [] for rid, path, score in ordered: latency 0.0 for layer, expert in path: # 模拟专家计算时间与负载直接相关 latency 1.0 self.avg_load[layer, expert] * 0.5 total_latency latency per_request.append((rid, latency, score)) return per_request, total_latency / max(1, len(ordered))这个原型的核心逻辑在于congestion函数和schedule函数。congestion把一条请求会经过的所有专家负载累加成一个总分得分越高说明这条请求越可能撞上热点。调度器优先放行低分请求让高分请求在队列里多等一会儿。wait_budget参数用来控制“最多等多久”避免为了均衡而无限延迟。4.3 负载均衡的效果评估为了验证效果我造了两组模拟数据一组是请求路径随机、专家访问均匀的“理想数据”另一组是模拟热门专家集中访问的“顺峰数据”。分别用这个调度器调度后对比一下平均排队时间。if __name__ __main__: balancer SimpleLoadBalancer(num_layers8, num_experts32) # 模拟一批请求路径中包含热门专家 (2, 5) 和 (6, 20) hot_layers [(2, 5), (6, 20)] requests [] for i in range(200): path [] for layer in range(8): expert np.random.randint(0, 32) if layer in [2, 6] and np.random.rand() 0.35: # 部分请求会路由到热门专家 hit hot_layers[0] if layer 2 else hot_layers[1] expert hit[1] path.append((layer, expert)) requests.append((i, path)) # 模拟前序负载积累 for layer, expert in path: balancer.update(layer, expert, delta0.1) per_request, avg_latency balancer.simulate_execute(requests, max_batch8) print(平均完成延迟:, round(avg_latency, 3)) # 输出拥塞度最高的前 5 个请求确认它们被延后执行 sorted_req sorted(per_request, keylambda x: x[1], reverseTrue) for rid, lat, score in sorted_req[:5]: print(frid{rid}, latency{lat:.2f}, congestion{score:.2f})我跑了几轮对比不调度和经过错峰调度的情况平均端到端延迟大约能降低 15% 到 25%。这里的具体数值受模拟参数影响很大更重要的结论是把请求按跨层拥塞度重新排列后热点专家处的排队长度明显下降整体的尾部延迟改善比平均延迟更显著。4.4 参数调优要点与工程集成这个原型里几个参数值得细调alpha是滑动平均的平滑系数控制负载预测对最近变化的敏感度。线上流量突变比较剧烈时alpha可以调到 0.5 以上如果流量相对平稳0.2~0.3 更稳。window是统计窗口影响对周期性负载模式的捕捉能力。wait_budget是调度等待预算这个参数直接关系到服务质量设得太大部分请求可能被压很久设得太小错峰效果就没了。工程上接入真实推理框架时可以把congestion函数返回的分数作为 micro-batch 调度器的优先级权重而不是像原型里这样只做简单的二分。框架本身的 continuous batching 已经保证了高吞吐跨层均衡要做的是在“下一个该执行哪个 batch”的决策里加入“未来路径负载”这一维度。实际操作中我建议优先在 prefill 阶段做这个调度因为 prefill 阶段的计算密度高、排队效应明显decode 阶段反而以显存访问为主错峰收益没有那么大。5. 常见问题排查实录5.1 路由分配不均衡怎么快速发现线上出现吞吐下降时别急着调模型或加卡先看路由日志。绝大多数推理框架都能输出 token 到专家的路由统计直接画一张“专家访问热度柱状图”就够了。如果访问量排名前 10% 的专家承担了超过 30% 的流量基本可以认定负载失衡。更精准的诊断是看专家级排队耗时。单个专家的平均排队时间如果明显高于其他专家说明这台卡上的计算资源被某几个专家拖住。再进一步是看跨设备 all-to-all 通信时间如果通信耗时随并发升高非线性增长大概率是通信热点导致的。5.2 跨层均衡会不会导致延迟变高这是做调度优化时最容易被质疑的问题。任何调度都会引入额外的等待时间跨层均衡也不例外。关键是看它在“引入小等待”和“消除大排队”之间是否划算。我实测下来只要wait_budget设置合理牺牲的少量尾延迟完全能被热点排队缩短带来的收益覆盖。但有一个坑如果把请求压得太狠反而可能造成新的排队。因为调度器让某个请求等待是为了躲开当前热点但如果接下来那个热点一直持续所有被延后的请求最终还是会撞在一起形成“滚雪球”。所以跨层均衡必须配合动态反馈一旦观察到排队长度不降反升就立即降低等待预算、加快放行节奏。5.3 负载均衡与连续批处理怎么配合连续批处理continuous batching本身也是一种调度形式它按 token 级别动态插入新请求。跨层均衡和 continuous batching 天然互补continuous batching 决定“哪些请求可以进入执行引擎”跨层均衡决定“进入引擎后的多个 micro-batch 之间应该用什么样的节奏去发”。实际操作里我习惯把它拆成两层。第一层用 continuous batching 维护一个“就绪队列”第二层用跨层负载感知的调度器决定“从就绪队列里按什么顺序取 batch”。两层都做效果才稳定只做一层要么吞吐上不去要么尾延迟压不下来。5.4 常见问题速查表现象可能原因排查手段GPU 平均利用率高但吞吐低所有请求都在排队等待同几个热点专家查专家热度分布和专家级排队时间某张卡频繁 OOM该卡负责的专家被大量请求命中batch 膨胀查看每卡 batch 大小和中间激活显存all-to-all 通信耗时异常高多个热门专家集中在跨卡路径上查看通信拓扑和路由转发的目的端分布接入跨层均衡后尾延迟反而上升等待预算过大或热点持续存在调小 wait_budget增加负载预测更新频率冷门专家 offload 后首次访问特别慢换入换出过于频繁加上跨层预取策略避免频繁换入爆发这里也提醒一下开源框架自带的负载均衡 loss 只是训练期约束线上推理阶段没有路由损失可以做一切均衡只能靠调度器实时观察和调整。最后再分享一点个人经验MoE 推理的负载均衡没有银弹。我一开始也迷信“把每个专家都放一个副本”或者“把所有层都单独均衡一遍”结果不是显存爆炸就是通信变慢。后来想明白一个道理模型是分层的负载是流动的调度器应该把整个推理过程当成一条流水线来看而不是一堆独立专家的集合。EasyBalance 这个跨层错峰的思路最精彩的地方就是把“层”这个本来就存在的时间维度用起来了——让不同请求在不同层的高峰期互相避开而不是所有人都堵在同一层。沿着这个方向后面还可以把 KV cache 状态、prefill/decode 分离、prefix caching 都纳入调度目标错峰的空间会更大。如果你也在做 MoE 推理建议先不要急着加卡把路由日志拿出来仔细看看也许会有和我一样的发现。