编辑|杜伟
AI 算力的下一个浪费重灾区,可能就在价值数百万美元的整柜系统内部。
过去,GPU 集群让人担心的是带宽不足、通信太慢。现如今,华为 Atlas 900 A3 SuperPoD、NVIDIA NVL72 等机柜级(rack-scale)系统把数十颗 GPU 放进同一个高速互连域,芯片、互连、交换网络和整柜系统被共同设计成了一个更大的计算单元。
把 GPU 高速连接起来解决了数据快速移动的问题,但算力是不是真的能跑满,还要看运行时来不来得及处理不断变化的负载。这一问题在大规模 MoE(混合专家模型)训推中尤为突出。
MoE 模型大量专家分散到不同 GPU 上,每个 token 经过 Gate 路由到少数几个专家。这种稀疏激活机制导致每张 GPU 的实际负载依赖 token 的路由结果。一旦少数专家成为热点,大量 token 集中涌向承载它们的 rank。其他 rank 即使提前完成计算,也要等待最慢的一方。
热点造成的负载倾斜同时拖慢专家计算和 token all-to-all 通信,还可能抬高过载 rank 的激活显存峰值。在真实训推中,这种不均衡可能让实际吞吐与理想状态拉开多达 2 倍的差距。更棘手的是,热点很难提前猜准。专家负载会随输入、网络层和路由 bias 更新快速变化,当前热门的专家在下一批数据中可能迅速降温。
上个月,小红书 dots infra 团队联合北京大学等提出 UltraEP,首次将 MoE 专家负载均衡变成了一项逐 microbatch、逐层实时执行的系统能力,并且在 MoE 训练的实际生产中落地。
此前 EPLB 等代表性系统方案通常根据历史路由数据估计专家负载,再周期性调整专家副本与放置。UltraEP 直接依据当前层已经产生的真实负载进行决策:Gate 完成后,每个专家究竟接收多少 token 一目了然,随即在 token dispatch 之前完成副本求解,将热点专家的权重复制到预留的冗余 slot,并把 token 分流到固有专家(main expert)和副本上。
- 技术报告地址:https://arxiv.org/pdf/2606.04101
- 项目地址:https://github.com/Dots-Infra/UltraEP
UltraEP 更值得关注的地方,在于它将实时专家均衡真正放进 MoE 的执行关键路径,相关操作的额外暴露开销控制在了约 300 微秒(µs)。
在 Qwen3-235B 上,关键路径总开销基本都落在 300 微秒以内。
UltraEP 的价值最终体现在了模型吞吐上,其中在 106B 到 671B 参数的主流 MoE 上,真实动态负载下的训推吞吐平均达到 force-balanced 理想性能的94.3%。相较于业界 SOTA 训推框架(Megatron-LM/SGLang)平均提升1.49 倍。更直观地看,原本最忙 rank 的负载最高可达到平均水平的 4.01 倍,接入 UltraEP 后,这一比例被压缩到最高1.04 倍,由少数热点 rank 造成的拥堵基本被消除。
高带宽互联的价值得到进一步释放,它允许运行时以 microbatch 和层为粒度重新调度专家计算资源。对于专家数量和专家并行规模持续扩大的 MoE 模型,这类实时调度能力可能成为大规模 AI 系统充分利用 GPU 算力的关键一环。
300 微秒内,UltraEP 如何完成一次实时均衡
UltraEP 最直接的工程挑战,是必须在极短时间内完成实时均衡。任何过高的计算或通信开销,都可能抵消实时均衡带来的性能收益。
为此,它首先将一个完整的 EP Group 部署在了同一个高带宽 Scale-up 域内,使热点专家的权重复制通过机架内高速互联完成,避免进入带宽更低的跨机 Scale-out 网络。这一硬件条件提供了通信基础。
此外 UltraEP 分别从控制面和数据面进行了针对性设计:控制面采用 GPU-native 在线求解器,根据当前层的真实负载联合求解专家复制和 token 重路由方案;数据面则通过设计专用通信 Kernel,完成专家权重和梯度在不同 rank 之间的动态传输。
具体来看,一次前向实时均衡包含三个紧密衔接的环节:前向阶段依次获取真实负载、生成均衡方案,在 token dispatch 开始前准备好所需的专家副本;反向阶段则复用前向缓存的方案,将副本产生的梯度归约回固有专家。
这些操作都受到严格的时间约束。结果显示,UltraEP 在前向阶段引入的额外开销约为 300 微秒,反向阶段暴露的额外开销接近于零。
UltraEP 与 EPLB 等预测性均衡方案的对比,包括负载精确性、决策时机、操作频率三项指标。
在 Gate 之后获取当前层真实负载
如上所述,UltraEP 的介入时机位于 Gate 结束之后、token dispatch 开始之前。 此时,系统才能看到当前 microbatch、当前层中每个专家实际接收的 token 数量。不同于基于历史统计预测下一批负载,这一信息对应的是即将执行的真实计算需求。
UltraEP 首先收集整个 EP group 的路由元数据,形成当前 microbatch 和当前层的全局负载矩阵。
各 rank 获得同一份负载数据后,它们会在 GPU 上独立执行相同的确定性求解过程,生成一致的专家副本和 token 重路由方案,无需为此额外增加一次全局同步。
前向时,专家复制求解和实际传输开销落在关键路径上。
Quota 同时决定专家副本与 token 分流
拿到真实负载后,UltraEP 需要同时回答两个问题:哪些热点专家需要增加副本,以及每个固有实例和临时副本分别应该承担多少 token。
UltraEP 引入了额度(Quota),也就是重路由后各专家实例收到的 token 负载。通过直接求解 quota,UltraEP 将专家复制和 token 重路由放进同一个优化过程。每一步探索既能实例化新的专家副本,也能更新目标 quota 分布。后续重路由则按照这些 quota,将各源 rank 的 token 需求分配到具体专家实例。
求解过程还考虑多项工程约束,包括每个 rank 的冗余 slot 数量、同一逻辑专家禁止在同一 rank 上重复部署、新副本必须获得最低有效 quota。团队将最低 quota 设为 1024,避免为了少量 token 创建收益有限的专家副本。
Quota 确定后,UltraEP 优先让 token 使用同一 rank 上的本地专家实例。本地 quota 用完后,剩余 token 按照其他实例的剩余额度发送到远程 rank,满足负载上限的同时减少跨 rank 通信。最后将聚合 quota 落实到具体 token,确定每个 token 应被发送到哪个专家实例。
与同样使用精确负载、但沿用 EPLB 式复制求解和 Round-robin 重路由的 EPLB + 相比,UltraEP 将平均负载不均衡度从 1.19 降至 1.03,求解延迟降低 27.4%,平均使用的冗余专家 slot 数量减少 57.9%。
共享冗余 slot 管理动态专家副本
求解完成后,UltraEP 已经知道当前层需要复制哪些专家、分别复制到哪些 rank,以及每个副本应该承担多少 token。下一步则是在 token dispatch 开始之前,将热点专家的权重分发到目标 rank。
为了避免在每个 microbatch、每一层中频繁申请和释放显存,UltraEP 在每个 EP Rank 上预留固定数量的冗余 slot。固有专家的位置不变,系统只需将当前层需要的专家权重写入这些预留空间,在其他 rank 上形成临时副本。
这样的布局形成了一对多的映射关系:每个逻辑专家始终拥有一个固有实例,也可以根据当前负载在其他 rank 上临时增加零个或多个副本。而冗余 slot 不维护独立的优化器状态,主要用于存放当前执行所需的专家副本权重和梯度 buffer。
UltraEP 专家内存管理示意:系统包含 8 个逻辑专家,采用 EP4 配置,每个 Rank 预留 1 个冗余专家 Slot。
UltraEP 还采用了冗余 slot 的跨层复用机制,不同 MoE 层共享同一组冗余 slot 和 buffer。执行当前层时,系统会将这一层需要的专家权重写入共享 buffer。前向执行完当前层后,这组空间可以被后续层复用;进入反向阶段时,系统再根据前向缓存的复制方案重新分发对应副本权重。这样一来,冗余副本的显存开销不随 MoE 层数线性增长。
以 Qwen3-235B-A22B 为例,如果每层分别预留冗余 slot,每个 rank 需要 9.9 GB 额外显存。采用跨层复用后,每个 rank 只需 36MB 权重空间和 72MB 梯度空间,总计108MB。
流式通信算子加速动态权重分发
Rack-Scale 系统提供了高带宽 Scale-up 互联,但硬件峰值带宽并不会自动转化为有效带宽。UltraEP 的通信模式随 microbatch 和模型层不断变化:需要复制的专家、源 rank 和目标 rank 组合也在变,形成了一张稀疏、动态且高度不对称的通信图。
UltraEP 设计了Persistent Tile Streaming,它将专家权重或梯度切分成固定大小的 tile,再将当前复制方案转换为设备端传输任务。 持久化 Kernel 持续从任务流中领取下一个 tile,无需为每个专家副本单独启动一次通信。
Tile 之间采用双缓冲流水线,当一个 tile 正在写入远程副本时,Kernel 开始获取下一个 tile 的任务索引,并执行地址解析和数据加载。任务查找、动态寻址与同步等控制开销被并入数据搬运过程,减少每次传输单独暴露的额外时延。
Persistent Tile Streaming 提高了单次动态传输的效率,但当一个热点专家需要复制到多个 rank 时,固有专家所在 rank 仍可能出现一对多发送瓶颈。UltraEP 进一步引入Chunk Streaming Relay,也就是分片流式中继。
源 rank 根据当前复制方案中的发送负载,从目标副本所在 rank 中选择压力较低的部分设备作为中继,将原本集中在源端的多播(multicast)流量分摊到多个 rank。中继以 chunk 为单位边接收边转发,无需等待完整专家权重到达,也无需在源端发送与中继转发之间设置全局 barrier,减少阶段间的空闲等待。
热点专家 multicast 的中继方案:rank 0 上的一个专家需要被复制到 rank 1–9,其中 rank 2, 5, 8 被选为中继。图中展示源和中继 rank 沿时间线的收发状态。
在 EP64、256 专家的端到端测试中,UltraEP 将不同初始负载下的不均衡度均降至 1.01。采用 Adaptive Relay 后,权重同步延迟进一步稳定在 0.19 毫秒左右,并在高负载不均衡场景下带来最高约 1.7 倍加速。
当专家权重全部写入目标冗余 slot 后,token dispatch 即可按照前一步求出的重路由方案开始。到这里,一次前向实时均衡完成。
复用前向方案,完成权重恢复与梯度归约
进入反向传播后,UltraEP 不会重新求解复制方案,而是直接复用前向阶段缓存的副本布局和 token 重路由信息。
由于冗余 slot 跨层复用,反向执行到当前 MoE 层时,UltraEP 需要根据前向方案,将对应专家的权重重新分发到冗余 slot,使副本恢复到与前向计算一致的状态。各个副本完成反向计算之后,UltraEP 再通过 Tile Streaming Kernel,将副本梯度归约回固有专家的梯度 buffer。梯度归约完成后,冗余梯度 buffer 会被清空,供后续 MoE 层继续复用。
反向传播中,专家权重的重分发和专家副本的梯度归约可以和其他反向计算 overlap。
在反向阶段,UltraEP 还控制通信 Kernel 占用的 SM(流式多处理器)数量和共享内存空间,为同时运行的反向计算 Kernel 保留足够资源。借助计算与通信重叠,反向权重恢复和梯度归约带来的暴露开销接近于零。
至此,UltraEP 完成了前向专家复制、Token 分流与反向副本梯度归约的完整流程。MoE 负载均衡不依赖过去一段时间的统计结果,也不必隔一段时间才调整一次,使专家计算任务随路由结果实时调整。
瞄准 MoE Infra 尚未成熟的实时调度层
把 UltraEP 的整套流程放到整个 MoE 基础设施栈中,它所承担的角色也随之明确。
过去几年,MoE Infra 领域逐渐形成了一批聚焦不同系统环节的基础设施组件。以 DeepSeek 提出的 DeepEP 为代表,它聚焦的是 MoE token dispatch 和 combine 之间的跨卡通信,用于提高 token 分发与计算结果回传的效率。
小红书 dots infra 团队的 UltraEP 关注的是 Gate 之后暴露出的另一类基础问题,即实时负载偏斜:真实 token 负载产生之后,系统如何立即为热点专家增加副本,并把部分 token 分配给负载较低的 rank。
两项工作中,DeepEP 优化 MoE token 的通信效率,UltraEP 实时调度专家资源。这一定位为它划出了比较清晰的能力边界。
Megatron-LM、SGLang 等训推框架组织执行流程,DeepEP 等通信库负责 token 高效传输,UltraEP 则以独立运行时的形式,专门处理 Gate 之后的实时负载均衡,并可以与现有框架、通信组件组合使用。
UltraEP 让我们看到,逐层、逐 microbatch 的实时专家调度具备了生产可行性。若后续完成更广泛的硬件适配和框架集成,它有很大的潜力成为大规模 MoE 系统中可复用的基础设施组件。
小红书 Dots Infra 系统工程能力的「牛刀初试」
从 UltraEP 可以看到,系统能否真正落地往往取决于多个基础设施环节能否协同工作。在线求解、通信算子、显存管理、框架集成和生产验证缺一不可,这对团队的系统设计与工程落地能力提出了更高要求。
因此,UltraEP 只是小红书 dots infra 整体技术实力的一个缩影,它证明了该团队具备把复杂系统问题拆解为通用基础设施模块,并进行端到端生产验证的能力。
作为支撑小红书 dots 系列模型的核心基础设施团队,dots infra 正在推进分布式训练稳定性、低精度训推、RL System、Agentic Rollout、KV Cache 管理和 Tool 调度等方向。
包括此前开源的 BigMac 和此次发布的 UltraEP 在内,dots infra 团队在多个关键基础设施环节拿出了具有竞争力的成果。这些底层能力又最终服务于更完整的小红书 dots 技术体系。
dots infra 团队所属的 Dots Studio,由小红书升级内部大模型技术团队而来,目前覆盖了预训练、后训练、推理、多模态和 Agent 等大模型完整链路。
官方地址:https://github.com/studio-dots-ai
UltraEP 作为目前 dots Infra 团队披露出来的一个代表性模块,它释放出一个清晰的信号,小红书对大模型的投入深入到了训练、推理与智能体运行所依赖的系统底座。
热门跟贴