当新一代混合专家模型的总参数量正式突破万亿级门槛，专家并行已成为训练与推理中的标准配置。然而，一个容易被忽视的现实是：即便模型在预训练阶段已部署了负载均衡算法，在真实的训练与推理场景中，不同GPU之间的计算负载仍可能相差数倍。这意味着，部分GPU早早完成计算却只能空转等待，token all-to-all通信也会成为瓶颈，高负载GPU的峰值显存急剧攀升，极易触发OOM，理想吞吐与实际吞吐之间的差距甚至可能达到一倍之多。

针对这一难题，小红书与北京大学联合提出了UltraEP，首次将基于精确路由信息的实时负载均衡引入生产系统——在每个microbatch和每一层动态复制热点专家，使真实训练与推理场景也能逼近理想性能。UltraEP平均达到理想性能的94.3%，相比业界SOTA训练推理框架提升1.49倍，并已成功应用于大规模预训练的实际生产环境中。

论文地址：
https://arxiv.org/abs/2606.04101

技术报告：
https://dots-infra.github.io/UltraEP/zh/

代码仓库：
https://github.com/Dots-Infra/UltraEP

随着MoE模型参数量持续膨胀，大规模专家并行在生产中越来越普遍：专家被分散在不同设备上，token通过all-to-all通信在专家之间进行交换。

专家负载不均衡是影响实际训练和推理吞吐量的关键因素。由于路由是动态的，不同专家、不同设备接收到的token数量天然就不均衡。现有方案中，最具代表性的是DeepSeek的EPLB（2025），它根据上一个时间窗口的历史路由数据，周期性地重新摆放专家。这类预测性方法隐含一个前提——负载必须相对静态。

但现实情况恰恰相反。当下主流的细粒度MoE（几百个小专家）负载高度动态，历史信息很快过时，冷热专家的预测频频失准，均衡操作甚至可能变成负优化。

UltraEP 为什么能做到「最优」负载均衡

UltraEP选择了一条看似激进但最直接的路——基于门控后的真实负载，在每个microbatch的每一层实时进行专家重均衡。

这消除了预测偏差，但代价也很明显：预测式方法可以提前把开销掩盖或平摊掉，而UltraEP的开销全暴露在关键路径上，还是在microbatch这个最细粒度上高频发生。基于一个通信前提和一系列控制面-数据面优化，UltraEP能把关键路径开销压到300µs以内，并达到近乎最优的均衡效果。

首先，UltraEP限定热点专家复制在高带宽scale-up域内进行，避免跨机专家搬运。这是实现百微秒量级专家复制的通信前提。更关键的是，UltraEP设计了一个高效的均衡方案在线求解算法，以及一套高度优化的专家权重/梯度通信算子，将均衡操作本身的开销降到最低。

关键设计：让最优负载均衡走进生产

UltraEP的定位是生产级专家均衡库，其设计原则包括：

• 独立的Python/CUDA运行时，与DeepEP或训练推理框架解耦。
• GPU-native的计算/通信过程，和host没有数据交互。
• 保持数学等价性，以及与其他常用训练推理配置的兼容性。
• 高效的显存管理，将额外显存开销压到最低。

在这些原则的基础上，UltraEP的关键设计归纳如下：

UltraEP给每个rank预留固定slot放“冗余专家”（热点专家的副本），运行时无动态显存分配。副本不需要维护优化器状态，梯度会在反向中实时归约回原专家。

更关键的是权重/梯度buffer采用跨层复用：以Qwen3-235B为例，单个冗余slot的额外显存开销能从9.9 GB降到108 MB。

实时负载均衡面临的最大挑战，是新增的计算和通信是否会拖慢训练或推理过程。

前向传播中，UltraEP需要在门控完成后获取全局负载，才能进行复制方案求解（replication planning）和专家权重分发（weight distribution）。重路由（reroute）负责在同一个固有专家的多个副本间分流token。相比于前两个操作，重路由更轻，且基本都可以被权重分发掩盖。

在控制面的均衡方案求解中，此前方法把“专家放置”和“token重路由”看作解耦的两阶段，容易互相拖累。UltraEP直接求解每个专家实例最终会接到多少负载（quota），把两阶段耦合起来：每一步探索既能实例化新副本，也能更新负载分布。借助warp-level并行，EP64下求解时间仍保持在100µs以内。

在数据面上，专家权重分发和梯度归约都是高度动态、稀疏的通信，经典集合通信库和in-switch卸载都吃不下这种非规则模式。更棘手的是，少数热专家副本极多，其所在rank的对外多播会成为新瓶颈。为了打满scale-up物理带宽，UltraEP基于持久化算子（persistent kernel）实现，将专家权重或梯度切分成若干tile，利用内存语义和TMA进行异步的卡间数据搬运。

为了消除通信热点，UltraEP设计了分片流式中继（chunk streaming relay）的通信策略，按实时流量构建两阶段中继树，让低流量rank帮忙分摊、转发热点流量，并通过chunk级流式转发避免全局barrier和通信bubble。

反向传播时，专家重分发和梯度归约（gradient reduction）可以和其他反向计算overlap。UltraEP精细控制它们的SM占用和共享内存footprint，避免拖慢反向计算，并用保序累加保证梯度归约的确定性。

除了均衡算法，UltraEP还提供了一套可视化profiler，对均衡前后的负载情况进行层次化分析：既能一览全局分布，也能看清每个microbatch中具体的冷热rank和专家负载，从而全面评估均衡效果和剩余瓶颈。

在Qwen3-235B、GLM4.5/4.7、DeepSeek-V3等模型上，研究团队分别基于Megatron-LM（训练）和SGLang（推理prefill）进行了评估。训练统一采用EP64的专家并行，推理则根据模型专家数采用EP64或EP40。

训练实验结果，包括吞吐（TFLOPS/GPU）和总体均衡度。

推理实验结果，包括TTFT随每秒请求数（RPS）的变化，以及总体均衡度。

主要结果如下：

• 训练：平均达到理想吞吐的94.6%，相比Megatron-LM提升42%；
• 推理prefill：达到理想吞吐的90%–97%，相比SGLang提升1.56倍；
• EPLB、LPLB因历史滞后、复制预算受限，效果始终落后于UltraEP。

值得一提的是，UltraEP把rank间不均衡从1.30–4.01稳定压到1.01–1.04。它与理想上限之间剩下的差距，主要来自实际路由下各专家负载的固有非一致性和少量控制开销，而不是残余不均衡或关键路径开销。

UltraEP已经在正式生产中部署。在一个288B参数MoE模型的完整预训练中，UltraEP保持了不低于理想性能92%的水平，显著提升并稳定了长周期训练吞吐。

UltraEP的核心判断是：随着通信带宽的提升，实时、精确的系统侧负载均衡将成为专家并行的基础能力。算法侧均衡负责训练稳定性和专家特化（specialization），系统侧均衡则负责把每个microbatch中已经发生的负载偏斜重新摊平；二者目标不同，但可以自然叠加。

下一步很自然的扩展是RL场景。由于面向特定领域数据，且没有预训练中的算法侧均衡调控，专家负载往往表现出ReLibra（https://arxiv.org/abs/2605.08639）观察到的，类似推理prefill中的强动态性。因此，UltraEP有机会成为MoE RL基础设施中统一的负载调节层。