发表时间: 2026-05 · arXiv:2605.08639 (Peking University)
原文: https://arxiv.org/abs/2605.08639
Chao Jin∗ Xinming Wei∗ Yinmin Zhong∗ Chengxu Yang† Bingyang Wu∗ Ruidong Zhu∗ Zili Zhang∗ Yuliang Liu† Xin Jin∗
∗School of Computer Science, Peking University †Independent Researcher
一句话结论 本文提出了针对大语言模型强化学习训练中混合专家模型(MoE)负载不均衡问题的系统 ReLibra,利用采样与训练阶段处理相同序列的“路由重放”先验,将负载均衡解耦为跨节点的批次间专家重排与节点内的微批次专家动态复制,在多样化强化学习任务上将训练吞吐量相比 Megatron-LM 提升最高达 1.6 倍,并达到理想均衡基线 90%–94% 的性能。
要解决什么问题 在大语言模型的强化学习后训练中,专家并行(EP,将不同专家分配至不同 GPU)面临严重的动态负载不均。一方面,强化学习任务多聚焦于数学、代码等垂直领域,各微批次激活的专家集差异巨大,且强化学习阶段移除了预训练时期的辅助均衡损失与路由偏置,使热门专家在相邻微批次间的重合率通常低于 0.5;另一方面,现有 MoE 负载均衡系统(如 EPLB、PopFetcher)面向预训练设计,依赖历史统计预测未来负载,在微批次剧烈波动下预测完全失效。这种微批次失衡导致承载热门专家的 GPU 沦为长尾,其他 GPU 则在跨节点全对全(All-to-All)通信中严重阻塞空转(空等时间可占 MoE 前向耗时的 31%,拖慢算子耗时达 1.4 倍);该长尾效应还会向数据并行与流水线并行传导,放大流水线气泡与跨节点延迟。
怎么做的 核心思路是利用强化学习的工作流特性规避预测盲区,并通过时间尺度解耦匹配集群拓扑带宽。在强化学习中,生成响应的采样(rollout)与参数更新的训练阶段处理完全相同的词元序列且模型参数一致,因此样本路由决策在采样时即可确定并在训练中精确“路由重放”;同时后处理的奖励计算与 KL 散度评估留出了充裕时间在关键路径外求解方案。为避免微批次全局调度的 NP 难复杂度,ReLibra 将其解耦为两个部件:跨节点低带宽 RDMA 匹配粗粒度的批次间专家重排,节点内高带宽 NVLink 匹配细粒度的微批次专家复制。在批次间尺度,系统以最小化 MoE 执行时间为目标:
其中 $\mathcal{G}$ 为 GPU 集合,$T_g^{\mathrm{comp}}$ 和 $T_g^{\mathrm{comm}}$ 分别为 GPU $g$ 的计算与通信时间。系统以贪心最长处理时间算法(LPT)为初解,使用基于交换的模拟退火算法进行全局搜索,并借助 Log-Sum-Exp 算子平滑最大值函数的梯度平原;随后通过第二轮退火重排样本在注意力副本上的放置以优化数据局部性。在微批次尺度,系统采用层共享副本缓冲区,单 GPU 仅分配 2 个副本槽位并在各层轮流复用,将副本内存开销降低一到两个数量级;通过异步张量内存加速器(TMA)内核将副本传输与注意力计算完全重叠掩盖。针对副本放置变量 $z_{e,g} \in \{0,1\}$ 与词元分配变量 $y_{j,e,g} \in [0,1]$,系统构建混合整数线性规划(MILP):
效果如何 实验在包含 20 个节点、160 块 NVIDIA Hopper GPU(单节点 8 卡,卡间 900 GB/s NVLink,节点间通过 8 根 400 Gbps 导轨优化 InfiniBand RDMA 互联)的集群上展开,并在独立的 128 核 CPU 节点上异步运行求解器。测试涵盖 Qwen3-30B-A3B(EP 为 16、DP 为 4)、GLM4.5-106B-A12B(EP 为 32、PP 为 2)以及 Qwen3-235B-A22B(EP 为 32、PP 为 5)三款模型(每层 128 专家,选取 top-8 路由),涵盖推理、代码、指令遵循及混合四类任务。对比基线包括:主流框架 Megatron-LM、基于历史激活模式预取的 PopFetcher、采用批次静态贪心重排并提供真实负载先验的 $\text{EPLB}^+$、引入微批次词元动态切分并提供真实负载的 $\text{LPLB}^+$,以及采用随机路由破坏模型算法的系统极限基线 Balanced。结果显示,ReLibra 较 Megatron-LM 取得 1.21–1.58 倍(最高 1.6 倍)的训练吞吐量提升,较 PopFetcher 提升 1.17–1.27 倍,较具备真实负载的 $\text{EPLB}^+$ 与 $\text{LPLB}^+$ 仍提升达 1.06–1.21 倍,达到 Balanced 基线 90%–94% 的吞吐量;跨 GPU 负载偏斜比从极度失衡下降并稳定在 1.00–1.08。系统开销方面,层共享缓冲区仅占单 GPU 参数内存的 0.13%–1.62%,批次间重排仅占批训练耗时的 1.4%–2.1%,副本同步和求解时间分别被计算和后处理完全掩盖。该方法的局限在于依赖采样与训练路由一致的重放先验;且专家复制受限于单节点内,跨节点的突发失衡仍需依赖批次间重排调节。
混合专家模型(Mixture-of-Experts, MoE)[1–3] 已成为扩展大语言模型(LLMs)参数容量的主流架构 [4–9]。在专家并行(Expert Parallelism, EP)[6, 11, 19] 下,MoE 通过将 token 稀疏路由到少量专家,实现了在不显著增加计算量的前提下提升模型容量。然而,MoE 训练普遍存在负载不均衡问题 [14–18],当某些专家被频繁激活成为“热专家”(hot experts)时,承载这些专家的 GPU 会严重过载成为落后者(straggler),拖慢整个系统的训练效率。
与预训练阶段相对稳定的负载不均衡模式不同,MoE 大模型的强化学习(RL)后训练(post-training)[12, 13] 呈现出极其剧烈且高频波动的微批次(micro-batch)级不均衡特征:
现有的 MoE 训练系统(如 EPLB [25]、PopFetcher [17])均基于“以历史预测未来”的前提设计,通过统计近期训练批次的负载来推测未来的专家需求。但在 RL 训练中,微批次维度的剧烈波动导致这一前提彻底失效。
为此,本文提出了 ReLibra,一个面向 MoE RL 训练的高效系统。该系统的核心洞察是:利用强化学习独特的 “生成采样-训练反向”(Rollout-Training)工作流 中的 路由重放(Routing Replay) 机遇。在 RL 迭代中,Rollout 阶段生成的 token 序列在进入 Training 阶段时,其使用的 MoE 参数与 token 序列完全相同,因此两者的 token-to-expert 路由决策完全一致。这意味着在训练开始前,系统就已经百分之百获知了当前批次内所有微批次的确切专家负载需求。
图 1:MoE 强化学习训练工作流与路由重放(Routing Replay)的优化机遇。
图 2:MoE 训练中的两类负载均衡机制:专家重排(Expert Reordering)与专家副本(Expert Replication)。
本文的主要创新点与系统贡献如下:
设计了高效求解算法与无额外开销的轻量系统架构:
全面的端到端评估:在涵盖 30B 到 235B 参数规模的 MoE 模型(Qwen3、GLM-4.5 等)及多领域 RL 任务上,ReLibra 相比 Megatron-LM 实现了高达 $1.6\times$ 的训练吞吐提升,相比即使注入先验上帝负载(Oracle Loads)的先进基线 EPLB 仍提升达 $1.2\times$,且吞吐达到理想绝对均衡基线(Balanced)的 $90\% - 94\%$。
MoE 训练与专家并行。MoE 架构扩展了标准的 Transformer [31],其核心是将稠密的前馈网络(FFN)替换为包含多个并行 FFN 专家及一个可训练路由模块(Router)的稀疏 MoE 层。对于每个输入的 token,Router 会根据评分预测挑选出前 $k$ 个专家(Top-$k$),这些被选中的专家独立计算后,其输出被加权聚合得到最终结果。这种稀疏激活使得大模型能够以恒定的单 token 计算预算容纳数倍到数十倍的模型容量 [5–11]。
图 3:MoE 层前向计算与专家并行(EP)的 Dispatch 与 Combine 通信阶段。
由于 MoE 模型具有极大的参数量,单一显卡无法容纳,因而通常引入专家并行(EP)[6, 11, 19]。如图 3 所示,在 EP 架构下,同一 MoE 层内部的不同专家被切分并部署在不同的 GPU 上;在输入 token 完成前置 Attention 计算并经由 Router 决定目标专家后,系统通过全局 All-to-All 通信(即 Dispatch 阶段)将 tokens 分发给对应的目标专家所在 GPU;专家完成前向计算后,再通过反向的 All-to-All 通信(即 Combine 阶段)将计算结果传送回原始 GPU 并聚合。当前前沿的大模型(如 Qwen3-235B-A22B [7] 拥有 128 个专家,DeepSeek-V3-671B [6] 拥有 256 个专家)使得跨节点的 EP(Cross-Node EP)成为工业界常态。在工业集群中,EP 还会与数据并行(DP [32])、张量并行(TP [33])以及流水线并行(PP [34, 35])协同混合编排。
大语言模型的强化学习训练工作流。RL 是增强大模型特定领域推理能力 [13] 及人类偏好对齐 [12] 的关键阶段。每一次典型的 RL 迭代都由一个两阶段循环构成(如图 1 所示):
论文选取了包括推理、代码、指令遵循以及混合领域在内的标准基准(详见下表 1),并在 32 规模的跨节点 EP 环境下(Global Batch 1024,切为 32 个微批次)深入剖析了 Qwen3-235B-A22B [7] 的不均衡动态。
表 1:实验所使用的 RL 任务领域与对应数据集。
| 领域(Domain) | 数据集(Dataset) |
|---|---|
| 推理(Reasoning) | DAPO-Math-17k [21], GPQA [36] |
| 指令遵循(Instruction Following) | IFBench [23] |
| 代码生成(Coding) | CodeForces [22] |
| 混合领域(Mixed) | 以上三者的混合,遵循 DeepSeek-V3.2 [20] 的配比 |
图 4:单个 MoE 层在不同领域连续微批次间的专家负载动态分布(相同颜色代表同一专家)及负载偏斜度变化。
观察一:专家负载在微批次维度剧烈波动且偏斜度极高。如图 4 所示,作者展示了某一 MoE 层中各专家在连续微批次间的负载占比(Load Share)与负载偏斜度(Skewness,定义为微批次内单个专家的最大 token 数除以所有专家的平均 token 数)。分析表明:
图 5:相邻微批次间负载最高的前 $k$ 个专家的平均重合比例(Intersection Ratio)。数值越低代表热专家切换越剧烈。
观察二:相邻微批次间热专家交集比例普遍低于 50%。为了严格量化时间维度上的不均衡跳变,作者定义了第 $t$ 个与第 $t+1$ 个微批次间 Top-$k$ 最热专家的重合比例(Intersection Ratio):
其中 $H_t^k$ 代表第 $t$ 个微批次中负载最高的 $k$ 个专家。如图 5 所示,针对 30B、106B 到 235B 的不同模型,在混合数据集上对所有 MoE 层求平均,Top-4 与 Top-8 专家在相邻微批次间的重合比例绝大多数时间都低于 0.5。这意味着在任意连续的两个微批次之间,超过一半的头部热专家发生了替换。
图 6:不同微批次下单个 MoE 层前向计算耗时分解对比(Balanced 均衡状态、Straggler 落后者状态、Blocked 阻塞等待状态)。
观察三:微批次不均衡导致硬件出现严重落后与空载阻塞。图 6 展示了同一块 GPU 在不同微批次下的前向传播时间分解。当遇到分配给本 GPU 的热专家负载过载时(“Straggler” 案例),不仅计算时间拉长,且由于需要接收海量远程 token,通信耗时显著激增,整体前向耗时飙升至均衡状态(Balanced)的 $1.4\times$。更严重的是,当该 GPU 在下一个微批次中仅处理冷门专家时(“Blocked” 案例),虽然其自身的本地专家计算量极小,但其前向时间的 $31\%$ 都耗费在等待远端落后 GPU 完成的同步屏障(Barrier Stall)上。在多维混合并行(DP+PP+EP)集群中,这种局部的微批次级不均衡会沿着流水线级联放大,引发大范围的 Pipeline Bubbles 与跨机架滞后。
核心机遇:路由重放(Routing Replay)与充裕的决策窗口。
核心挑战:微批次级组合优化的 NP-Hard 属性。
尽管路由信息已知,但为数十个层、数十个微批次实时规划最优的负载均衡方案在计算上极其棘手。系统可采用的手段包括:
在显存容量上限、通信带宽拓扑约束下,微批次级的联合规划构成了一个超大规模的组合优化问题。文献已证实,即使完全不考虑副本、仅考虑无显存约束的同构机器完工时间最小化(Makespan Minimization on Identical Machines, $P||C_{\max}$),该问题也是 NP-hard 的 [30]。若加上副本选择与多节点容量约束,问题等价于设施选址问题(Facility Location Problem)[44] 的变种,在微秒/毫秒级时限内求取全局精确最优解是不可能的。
核心设计原则:时间尺度解耦(Timescale Decomposition)。
针对跨节点分布式集群的不对称拓扑(节点内 NVLink 具备数百 GB/s 带宽,而跨节点 RDMA 链路带宽要低得多),ReLibra 提出了时间尺度的层次化解耦:
图 7:ReLibra 系统的总体分层架构,包含控制平面(求解器)与数据平面(路由捕获器与执行器)。
ReLibra 的整体架构划分为控制平面(Control Plane)与数据平面(Data Plane),如图 7 所示:
控制平面统筹全局 RL 工作流状态:
解耦负载均衡求解器(Load-Balancing Solver):将复杂的端到端优化拆解为两阶段求解流水线:
非关键路径重叠:控制平面的计算任务被完全调度在 Post-Rollout 处理阶段(如奖励计算、KL 散度评估期间)异步并行执行,其端到端求解耗时完全脱离训练的关键执行路径。
数据平面部署于分布式 GPU 训练集群中,由两个核心模块协同驱动:
批次间专家重排致力于依据整个训练批次聚合的路由数据,最小化整个 EP 组的全局 MoE 模块执行时间。所涉及的核心数学符号如表 2 所示。
表 2:批次间专家重排核心符号定义说明。
| 符号 | 描述与物理含义 |
|---|---|
| $\mathcal{G}$ | 专家并行(EP)组中的所有 GPU 集合 |
| $E, E_g$ | 模型某一层中的全量专家集合,以及部署在 GPU $g$ 上的专家集合 |
| $x_{j,e}$ | 在当前全局批次中,源自 GPU $j$ 并被路由至专家 $e$ 的 token 总数 |
| $\tau(j,e)$ | token 从源 GPU $j$ 到目标专家 $e$ 的通信链路拓扑类型 |
| $y_{j,e,g}$ | 来自 GPU $j$ 的属于专家 $e$ 的 token 分配给 GPU $g$ 处理的比例 |
| $z_{e,g}$ | 二值变量,指示专家 $e$ 的副本是否允许放置在 GPU $g$ 上 |
| $L_g^{\mathrm{comp}}$ | GPU $g$ 上承担的总计算负载(以处理 token 数计) |
| $C_g^{\mathrm{nvlink}}$ | GPU $g$ 在 NVLink 上注入的发送/接收通信负载 |
| $C_g^{\mathrm{rdma}}$ | GPU $g$ 在 RDMA 网络上注入的发送/接收通信负载 |
| $T_g^{\mathrm{comp}}, T_g^{\mathrm{comm}}$ | GPU $g$ 上预估的 MoE 计算时间与 All-to-All 通信时间 |
| $T_{\mathrm{MoE}}$ | EP 组内整个 MoE 模块的最终端到端执行耗时(由落后者决定) |
从路由决策到计算与通信负载映射。
将负载以处理 token 数度量。设 $x_{j,e}$ 为源 GPU $j$(负责该 token 的 Attention 计算)路由到专家 $e$ 的 token 总量。假定每个专家完整驻留在单一 GPU 上不进行张量切分。则驻留在 GPU $g$ 上的专家集合记为 $E_g$。GPU $g$ 的总计算负载为:
其中 $l_e = \sum_{j \in \mathcal{G}} x_{j,e}$ 为专家 $e$ 承载的全部 token 数。
图 8:MoE 训练中基于 Rail-Optimized(轨道优化)的多机多卡网络拓扑及三类通信流量路径。
轨道优化(Rail-Optimized)网络拓扑下的通信负载分类。
在大规模集群中普遍采用轨道优化拓扑(如图 8 所示)。同一节点内具有相同本地索引的 GPU 与网卡(NIC)连接在同一 PCIe 交换机上形成一个 Rail;不同节点间相同 Rail 的 NIC 汇聚连接到相同的 Leaf 交换机。在这种拓扑下,Token 的 All-to-All 路由会产生四种不同的源-目的物理相对路径 $\tau(j,e) \in \{\mathrm{loc}, \mathrm{nv}, \mathrm{sr}, \mathrm{cr}\}$:
据此,在 Dispatch 阶段,GPU $g$ 注入到 NVLink 上的发送(tx)与接收(rx)通信负载精确形式化为:
MoE 模块端到端执行时间估算建模。
设 Transformer 隐藏层维度为 $h$,前馈网络中间层维度为 $h'$。每个 token 经过单一专家的 FFN 计算需要执行 3 次 GEMM 操作,对应浮点运算量约为 $6hh'$ FLOPs(忽略极小开销的逐元素激活与门控权重计算)。若单卡有效计算算力为 $F$ FLOPS,则 GPU $g$ 上的专家计算时间为:
若 NVLink 与 RDMA 的实测有效双工通信带宽分别为 $B_{\mathrm{nvlink}}$ 与 $B_{\mathrm{rdma}}$,由于通信链路全双工特性(tx 与 rx 并发)以及先进 All-to-All 算子(如 DeepEP [50])对 NVLink 与 RDMA 传输流水线的重叠掩盖设计,GPU $g$ 的通信时间建模为:
由于 MoE 模型通常在单层拥有超过百个专家(如 128 或 256 个),全解搜索空间高达 $O(|E|!)$。简单的贪心长作业优先(Longest-Processing-Time, LPT)[55] 算法仅按专家累计计算负载 $l_e$ 降序排序,依次放置在当前负载最小的 GPU 上。然而,单纯按计算排序忽略了通信拓扑耦合性:移动一个专家的宿主 GPU,不仅改变其计算分布,更会彻底改变该专家与产生对应 token 的各源 GPU 之间的相对拓扑距离(例如从 NVLink 退化为 Cross-Rail RDMA),反而引发灾难性的跨节点通信拥塞。
为此,ReLibra 提出了一种联合建模计算与通信代价的基于交换的模拟退火算法。算法首先以 LPT 的解作为初始搜索起点 $\Pi_0$;随后在多核 CPU 上并行启动多个不同随机种子的退火进程,迭代执行双专家随机交换,根据平滑后的目标函数评估是否接受转移。算法伪代码规范整理如下:
# 算法 1:批次间专家重排与模拟退火搜索算法 (Algorithm 1)
# 输入:专家集合 E, GPU 集合 G, 路由统计矩阵 {x_{j,e}}, 随机种子集合 S, 降温速率 gamma, 终止阈值 epsilon
# 输出:全局最优专家重排布局方案 Pi_star
def greedy_lpt(routing_matrix, gpus):
# 根据各专家总计算量降序排序,贪心放置于当前总负载最小且未超容量的 GPU 上
pass
def incremental_moe_time(pi_prime, ea, eb):
# 仅针对被交换的两个专家 ea, eb 执行 O(|G|) 增量时间评估
pass
def anneal_search(seed, pi_0, gamma, epsilon):
pi = pi_0.copy()
pi_star_s = pi_0.copy()
t = smoothed_moe_objective(pi)
t_star_s = t
theta = t # 初始温度设为初始状态评价值
while theta > epsilon:
ea, eb = random_sample(E, 2)
pi_prime = swap_experts(pi, ea, eb) # 状态增量变迁
t_prime = smoothed_moe_objective(pi_prime) # 基于平滑目标函数评估
delta = t - t_prime
# Metropolis 准则:优化则直接接受,恶化则依概率接受以跳出局部极小
if t_prime < t or random_uniform(0, 1) < math.exp(delta / theta):
pi = pi_prime
t = t_prime
if t < t_star_s:
pi_star_s = pi.copy()
t_star_s = t
theta = gamma * theta # 温度衰减
return pi_star_s, t_star_s
# 并行调度与最优选择
def optimize_expert_reordering():
pi_0 = greedy_lpt(x, G)
results = parallel_for(s in S, anneal_search, s, pi_0, gamma, epsilon)
best_pi = min(results, key=lambda res: compute_exact_t_moe(res[0]))
return best_pi
增量状态评估加速(Incremental State Update)。
若每次交换两个专家都对全集群所有 GPU 的计算和通信负载重新计算,单步评估复杂度将达到 $O(|E| \cdot |\mathcal{G}|)$。ReLibra 指出:当交换专家 $e_a$ 与 $e_b$ 时,全集群除这两个专家所涉及的路由流量之外,其余所有专家的计算负载和通信拓扑完全不受影响。因此,算法仅重新扫描 $\{x_{j, e_a}\}$ 与 $\{x_{j, e_b}\}$ 在源 GPU 集合 $\mathcal{G}$ 上的投影并进行差量抵消。评估单次交换的时间复杂度直接压缩到 $O(|\mathcal{G}|)$,使得单核 CPU 每秒可迭代搜索数十万次。
目标函数平滑化(Objective Smoothing)。
精确目标函数 $T_{\mathrm{MoE}}$ 包含外层与内层的大量 $\max(\cdot)$ 算子。在离散搜索空间中,$\max$ 算子会导致搜索曲面出现极大的水平停滞高原(Plateaus):某次专家交换可能显著优化了次瓶颈 GPU 的负载,但只要全局第一瓶颈 GPU 耗时微未改变,目标函数值就毫无变化,模拟退火算法因此丢失了引导梯度的方向信号。为此,ReLibra 引入 Log-Sum-Exp(LSE)软近似算子替换所有的硬 $\max$ 算子:
当缩放因子 $\beta \to \infty$ 时,$\mathrm{LSE}_\beta \to \max$。平滑后的连续优化目标函数定义为:
数据局部性二次优化(Optimizing Data Locality)。
现代大规模训练通常采用跨节点 EP 搭配“数据并行 Attention”(Data Parallel Attention [6])架构。每个 GPU 上都常驻一份独立的 Attention 模块权重副本,这意味着全局批次中的任何一个样本(Sample)原则上被调度给 EP 组内的哪一块 GPU 执行 Attention 均可。这为降低通信总量创造了二次优化空间:若将样本放置在与其所激活的高频专家物理距离更近(同卡或同节点)的 GPU 上,将从根本上削减跨节点 All-to-All 流量。
因此,在获得专家最优排布 $\Pi^\star$ 后,ReLibra 保持专家位置冻结,启动第二轮基于样本交换的模拟退火:在 GPU 之间随机交换分配的样本,依据相同的平滑通信公式驱动优化。最终输出具备拓扑亲和感知(Locality-Aware)的样本-GPU 绑定映射表。
在微批次尺度上动态创建热专家副本面临着显存容量与同步延迟的双重严苛约束。
图 9:用于微批次细粒度副本调度的层共享副本缓冲区(Layer-Shared Replica Buffer)架构与前向跨层复用机制。
显存开销压缩机制。
传统思路如果为每一个 MoE 层都静态预分配专属副本显存槽位,在模型深度较大时显存开销极为沉重。设模型共有 $L$ 个 MoE 层,每张卡在 EP 下承载 $M$ 个常驻专家;若允许每层放置 $r$ 个临时副本,朴素方案需要常驻 $L \cdot r$ 个专家槽位。
ReLibra 的核心洞察是:在深度前向与反向流水线中,计算是逐层推进的(Layer-by-Layer),在任意给定的时刻,硬件只需要访问当前活动 MoE 层的专家参数。如图 9 所示,ReLibra 在每张 GPU 上仅静态开辟容纳 $r$ 个专家参数的统一物理缓冲区,并跨全模型 $L$ 个层进行复用。在前向传播计算第 1 层前,将第 1 层所需的动态副本填入缓冲区;第 1 层计算完成后,该缓冲区立即被用于接收第 2 层所需的动态副本,依此类推。这使得额外显存开销从 $O(L \cdot r)$ 骤降到 $O(r)$,在拥有数十层的大模型中显存占用直接降低一个到两个数量级。
同步通信完全隐藏机制。
ReLibra 将副本的生命周期严格限定在节点内部(同机 8 卡之间),从而完全避开慢速跨节点网络,仅利用数百 GB/s 的 NVLink 进行极速数据摆渡。更进一步,系统采用基于单程序多数据(SPMD)模式的 Push 通信策略:
混合整数线性规划(MILP)严格数学建模。
针对某一层特定微批次,ReLibra 将专家副本选择、放置节点与 Token 切分比例联合形式化为一个标准的混合整数线性规划问题。
对于给定的微批次,设输入矩阵为 $\{x_{j,e}\}$。设 $\mathcal{G}(e) \subseteq \mathcal{G}$ 表示与专家 $e$ 的宿主卡同处一个物理节点内的所有候选 GPU 集合。
引入决策变量:
综上,该微批次的 MoE 执行完工时间优化问题定义为:
受限于以下物理约束条件:
约束逻辑推导解析:
由于目标函数中的最大化算子可通过引入辅助变量进行标准线性化,整个模型在数学上是一个紧凑的 MILP。
增量贪心启发式算法(Incremental Greedy Heuristic)。
工业通用求解器(如 SCIP [56])在求解上述含大量二值变量的 MILP 时存在高昂的分支定界(Branch-and-Bound)时间开销,无法满足 RL 训练中极短的微批次求解时限要求。
因此,ReLibra 提出了一种启发式贪心与线性规划交替迭代的求解方案。该算法从“零副本”状态起步,在每次迭代中贪心寻找能带来最大吞吐改善的单一放置决策,随后将二值变量固定,调用极速 LP 算子求解连续变量 $y$。规范整理后的算法逻辑如下:
# 算法 2:微批次级专家副本与分流启发式规划算法 (Algorithm 2)
# 输入:专家集 E, 专家重排方案 Pi, GPU 集 G, 路由矩阵 {x_{j,e}}, 候选节点拓扑集合 {G(e)}, 缓冲区槽位限额 r
# 输出:最优副本布局矩阵 z_star, 最优 token 分流比例张量 y_star
def solve_token_split_lp(routing_matrix, z_placement):
# 在固定二值放置 z 的条件下,通过单纯形法或轻量 LP 求解最优连续变量 y
pass
def optimize_microbatch_replication(E, Pi, G, routing_matrix, r):
# 1. 初始化放置矩阵:本地原生专家置 1,其余全为 0
z = {}
for g in G:
for e in E:
z[e, g] = 1 if e in E_g else 0
# 2. 在无副本基准下初始化解
y = solve_token_split_lp(routing_matrix, z)
# 3. 贪心迭代扩容副本
while True:
# 检查是否还有任何 GPU 拥有空闲的副本 slot
available_gpus = [g for g in G if has_available_slot(g, z, r)]
if not available_gpus:
break
# 定位当前计算与通信综合负载最高的瓶颈落后者 GPU
g_b = argmax(G, key=lambda g: compute_gpu_moe_time(g, Pi, y, z))
# 在该落后 GPU 上,找出贡献最大绝对负载的专家
e_star = argmax(E_{g_b}, key=lambda e: compute_expert_load(e, y, z))
# 在该专家所属物理节点的同机 GPU 集合 G(e_star) 中,寻找尚有槽位且当前总负载最小的 GPU
feasible_targets = [g for g in G(e_star) if z[e_star, g] == 0 and has_available_slot(g, z, r)]
if not feasible_targets:
break
g_t = argmin(feasible_targets, key=lambda g: compute_gpu_moe_time(g, Pi, y, z))
# 在该目标 GPU 上开辟副本,并更新放置决策
z[e_star, g_t] = 1
# 重新调用 LP 算子,在新的副本拓扑下求取绝对最优的细粒度分流比 y
y = solve_token_split_lp(routing_matrix, z)
return z, y
ReLibra 基于开源强化学习框架 Slime [57] 与分布式训练框架 Megatron-LM [58] 构建,累计编写了 14,000 行 Python、C++ 和 CUDA 核心代码,Rollout 生成引擎接入 SGLang [59]。
低开销路由捕获器实现。
通过深度定制 SGLang 的自回归采样内核,在每个 Token 生成完成的瞬间挂钩 Router 模块,提取其 Top-$k$ 专家张量。为了防止向 CPU 频繁做小内存同步拷贝破坏 GPU 流水线,捕获器将连续步的路由张量追加到 GPU 上一个固定几十兆字节(Tens of MB)的循环 Ring Buffer 中,随后通过独立 CUDA Stream 异步 D2H 搬运至 Host 端常驻内存,使 Rollout 吞吐损失完全不可见。
多核并行负载均衡求解器实现。
控制平面的求解器全部使用高性能现代 C++ 实现,并使用 OpenMP [60] 进行多核多线程并行化编排:
基于持久化异步 TMA 内核的显存缓冲区通信实现。
针对 NVIDIA Hopper 架构的硬件特质,ReLibra 使用张量内存加速器(Tensor Memory Accelerator, TMA)[61] 编写了持久化异步数据传输 CUDA 内核。为了防止参数搬运抢占 Attention 计算的算力核心,该持久化内核被严格硬编码限制在极少量的流式多处理器(Streaming Multiprocessors, SMs)上执行。绝大多数 SMs 得以全速投入到计算密集的 FlashAttention 内核中,真正实现了零竞争、零抖动的通信-计算异步重叠。
互联网络拓扑:
独立求解节点:控制平面求解器部署于一台独立的 CPU-only 算力节点上,配备 128 核 Intel Xeon 高性能 CPU,由 Ray [64] 统一管理任务分发。
实验选用了来自不同机构、横跨三个数量级规模的先进生产级 MoE 模型,各模型的切分拓扑配置如表 3 所示:
表 3:各模型在基准评测中所使用的并行策略配置详情。
| 模型名称 | MoE 层数 | 占用卡数 | 数据并行(DP) | 流水线并行(PP) | 专家并行(EP) | 单卡专家数 | 激活 Top-$k$ |
|---|---|---|---|---|---|---|---|
| Qwen3-30B-A3B [7] | 48 | 64 | 4 | 1 | 16 | 8 | Top-8 |
| GLM4.5-106B-A12B [65] | 45 | 64 | 1 | 2 | 32 | 4 | Top-8 |
| Qwen3-235B-A22B [7] | 94 | 160 | 1 | 5 | 32 | 4 | Top-8 |
所有模型每层均包含 128 个专家,且每个 token 路由激活 8 个专家。Attention 模块均启用 Data Parallel Attention [6] 以消除 TP 通信开销。所有对比实验中的全局训练批次大小(Global Batch Size)统一设为 1024,每个训练批次切分为 32 个微批次(单微批次 32 样本)。Rollout 阶段单样本最大生成长度为 8K tokens,使用 GRPO 算法每 Prompt 采样 4 条 Response(Temperature=1.0, Top-$p$=1.0)。
图 10:在混合数据集(Mixed Dataset)上,不同模型规模下的端到端训练吞吐表现(tokens/s)。
图 11:在 Qwen3-235B-A22B 超大模型上,不同业务领域数据集下的训练吞吐对比。
核心实验发现:
表 4:在 Qwen3-235B-A22B 模型上,逐层剥离各机制后的归一化吞吐提升拆解。
| 序号 | 方案组合 | 推理领域(Reasoning)归一化吞吐 | 推理提升增量 $\Delta$ | 混合领域(Mixed)归一化吞吐 | 混合提升增量 $\Delta$ |
|---|---|---|---|---|---|
| 1 | Megatron-LM(原生基线) | 1.00 | - | 1.00 | - |
| 2 | (1) + 批次间专家重排(Inter-batch Reordering) | 1.12 | +12% | 1.15 | +15% |
| 3 | (2) + 批次内动态副本(Intra-batch Replication) | 1.46 | +34% | 1.39 | +24% |
表 4 严密展示了 ReLibra 两层机制的独立贡献:
图 12:批次间专家重排机制消融分析(对比原生 EP、贪心 LPT 专家放置 ReLibra-Expert、贪心数据分配 ReLibra-Data 以及完整退火 ReLibra-Inter)。
图 13:批次内动态副本机制消融分析(对比原生基线、仅微批次贪心副本放置 ReLibra-Replication、仅微批次连续分流 ReLibra-Splitting 以及完整联合 MILP 启发式 ReLibra-Intra)。
深入机制微观消融:
图 14:不同硬件 EP 并行规模与不同领域数据集下,各 GPU 间 Rank-Level 负载偏斜度箱线图分布。
作者统计了所有参与运算的 GPU 卡间实际处理 token 数量的 Rank 级偏斜度(Rank-Level Skewness,即 $\max_g L_g / \mathrm{avg}_g L_g$)。如图 14 所示:
图 15:ReLibra 控制平面在独立 CPU 节点上的求解耗时(批次间重排 + 批次内副本)与 Post-Rollout 处理时间的重叠对比。
表 5:各模型在单张 GPU 上为层共享副本缓冲区分配的参数额外显存占用与绝对占比。
| 模型名称 | 层共享副本缓冲区显存容量($r=2$) | 单卡常驻模型参数显存 | 额外显存开销占比 |
|---|---|---|---|
| Qwen3-30B-A3B | 18 MiB | 6.26 GiB | 0.28% |
| GLM4.5-106B-A12B | 22 MiB | 16.64 GiB | 0.13% |
| Qwen3-235B-A22B | 72 MiB | 4.35 GiB | 1.62% |
图 16:ReLibra 动态副本同步开销(前向参数推送与反向梯度推送)与 Attention 算子执行时间的异步重叠评估。
在某些 Off-Policy 强化学习体系中,为了最大化样本利用效率,通常会引入对同一批经验数据的多次参数更新(Multiple Updates per Rollout [26, 27]),或者采取生成与训练完全解耦的异步执行架构 [66, 67]。在这些情况下,随着参数被更新,严格意义上的前向 Token 路由评分会与 Rollout 产生漂移。
然而,最近的产业级探索(如 DeepSeek-V3.2 [20]、Kimi [28] 以及郑等人针对大模型强化学习稳定性的研究 [29])明确指出:在 Off-Policy 模式下采用强制路由重放(Forced Routing Replay)是一项极其关键的算法稳定机制。它强制要求被更新的专家必须且仅能使用由自己亲自生成的 Token 样本执行训练更新,有效规避了策略漂移引发的模型崩溃。因此,ReLibra 的确定性路由重放假设与现代大规模 RL 算法演进方向高度契合。
MoE 训练系统优化:
现有的 MoE 负载均衡设计主要面向预训练阶段。这类工作(包括 FlexMoE [14]、SmartMoE [15]、PopFetcher [17]、EPLB [25]、Lina [68])普遍依赖长时序的历史激活观测去外推下一阶段的稀疏度。在遇到强化学习领域跳变和缺乏平衡损失约束时,其预测准确率急剧退化;同时预训练极其短促的单步时间预算(几十微秒)也封死了复杂优化的可能。在底层通信调度方面,一类工作专注于分层 All-to-All 架构(如 Tutel [69]、HetuMoE [70]、DeepEP [50]);一类聚焦通信与计算的算子级重叠(如 MegaScale-MoE [71]、Centauri [72]、MegaScale [73]、COMET [74]、FLUX [75]);另一类则探索专家靠近数据的摆放策略(如 Janus [76])。ReLibra 提出的时间尺度分层算法在数学建模上与这些底层通信库完全正交,并可直接无缝结合。
大语言模型强化学习系统优化:
目前的强化学习系统主要分为同步式(如 HybridFlow [77])与异步式(如 RollArt [78]、StreamRL [67])。许多工作致力于消除 Rollout 阶段的推理瓶颈,例如通过推测解码加速生成的 TLT [40] 和 Seer [41],或者通过阶段融合削减 GPU 空闲的 RLHFuse [42]。ReLibra 专注于攻克互补的训练瓶颈,直接解决 MoE 在训练反向中的微批次不均衡问题,构成了端到端系统拼图的关键一环。
本文针对大模型强化学习(RL)训练中 MoE 架构出现的微批次维度高度剧烈波动的负载不均衡难题,提出了 ReLibra 系统。ReLibra 巧妙地捕捉并利用了 RL 工作流中固有的“路由重放”机遇,彻底摒弃了传统预训练系统所依赖的脆弱历史预测假设,在训练启动前获取了微批次级确切路由。系统提出了时间尺度解耦设计理念,将跨节点宏观专家重排分配于批次间尺度,将节点内细粒度动态副本分配于微批次尺度,精准匹配了集群的分层互联网络带宽特性。通过轻量级的层共享显存复用缓冲区、高效的模拟退火和混合整数线性规划启发式算法,ReLibra 在不侵占任何训练关键路径的前置时间窗口内,以近乎忽略不计的系统开销达成了微批次级的精细负载均衡。端到端实验表明,ReLibra 取得了相比于 Megatron-LM 最高达 $1.6\times$ 的显著吞吐提升,逼近了理论极限性能。
下表系统整理了第 3、4、5 节方法细节与实现中所引用的核心参考文献、标号、原文上下文描述及学术索引:
| 引用标号 | 论文题目 / 会议 / 年份 / URL | 原文引用段落与核心观点描述 |
|---|---|---|
| [6] | DeepSeek-V3 Technical Report arXiv preprint, 2024. https://arxiv.org/abs/2412.19437 |
在 4.1 节形式化建模中用于支撑专家不切分独立放置的工业实践;在 4.2 节中用于支撑 Cross-Node EP 搭配 Data Parallel Attention 架构以减少 TP 通信的实践背景。 |
| [8] | Kimi k2: Open Agentic Intelligence arXiv preprint, 2025. https://arxiv.org/abs/2507.20534 |
在 4.1 节用于印证大规模细粒度 MoE 模型将单个专家完整映射在单个 GPU 上、不使用卡内切分的系统设计事实。 |
| [9] | MiniMax-M1: Scaling Test-Time Compute Efficiently with Lightning Attention arXiv preprint, 2025. https://arxiv.org/abs/2506.13585 |
在 4.1 节用于佐证近年前沿 MoE 架构普遍采用海量专家且单个专家独立部署的原则。 |
| [17] | PopFetcher: Towards Accelerated Mixture-of-Experts Training via Popularity Based Expert-Wise Prefetch USENIX ATC, 2025. |
在 4.1 节与 5 节中用于指出先验 MoE 预训练工作以 token 数量度量专家负载,以及传统方案仅在全局批次维度维持静态固定副本的局限性。 |
| [25] | Expert Parallelism Load Balancer (EPLB) GitHub Repository, 2025. https://github.com/deepseek-ai/eplb |
在 4.1 节指出其以 token 数量形式化负载;并在 5 节指出 EPLB 仅能以全局批次为周期进行粗粒度重排与副本放置,无法感知微批次内剧烈波动的缺陷。 |
| [30] | Optimization and Approximation in Deterministic Sequencing and Scheduling: A Survey Annals of Discrete Mathematics, 1979. |
在 2.3 节与方法论开篇用于证明简化状态下的 MoE 专家负载映射在数学上等价于同构机器完工时间最小化问题($P\|\|C_{\max}$),属于经典 NP-hard 问题。 |
| [44] | On the Uncapacitated Location Problem Annals of Discrete Mathematics, 1977. |
在 2.3 节用于论证当引入显存上限与副本决策时,该问题扩展为类似设施选址问题(Facility Location Problem)的高维 NP-hard 属性。 |
| [45] | Linear-Programming-Based Load Balancer (LPLB) GitHub Repository, 2025. https://github.com/deepseek-ai/LPLB |
在 4.1 节用于说明基于 token 计数建模计算与网络负载的标准做法;并在第 5 节对比其仅在微批次做单一副本限制的 token 分流。 |
| [46] | Alibaba HPN: A Data Center Network for Large Language Model Training ACM SIGCOMM, 2024. |
在 4.1 节中用于证明轨道优化网络拓扑(Rail-Optimized Topology)在大规模大模型训练集群中的通用性与代表性。 |
| [47] | Accelerating Mixture of Experts Training With Rail-Optimized InfiniBand Networking in Crusoe Cloud NVIDIA GTC, 2024. https://www.nvidia.com/en-us/on-demand/session/gtc24-s63014/ |
在 4.1 节用于论证在 GPU 数据中心部署 Rail-Optimized InfiniBand 架构对于支持 MoE 分布式训练通信模式的标准规范性。 |
| [48] | GPUDirect RDMA NVIDIA Developer Documentation, 2026. https://developer.nvidia.com/gpudirect |
在 4.1 节用于定义跨节点同轨道(Same-Rail)通信的技术载体,即跨机 GPU 通过底层 RDMA 网卡进行显存直通传输。 |
| [49] | Optimized Primitives for Inter-GPU Communication (NCCL) GitHub Repository, 2026. https://github.com/NVIDIA/nccl |
在 4.1 节用于描述跨节点跨轨道(Cross-Rail)通信中,利用本地 NVLink 中继再接入 RDMA 网卡的多跳流量行为。 |
| [50] | DeepEP: An Efficient Expert-Parallel Communication Library GitHub Repository, 2025. https://github.com/deepseek-ai/DeepEP |
在 4.1 节用于推导通信模型时间公式,证实其通过重叠/流水线化 NVLink 与 RDMA 实现全双工并发;并在第 6 节作为底层高效 All-to-All 算子库使用。 |
| [51]–[54] | DistServe (OSDI '24), MegaScaleInfer (SIGCOMM '25), DistTrain (SIGCOMM '25), Heddle (arXiv '26) | 在 4.1 节用于说明系统通过离线 Profile 基准测试与数学插值获取硬件算力浮点常数 $F$ 及有效网络通信带宽 $B_{\mathrm{nvlink}}, B_{\mathrm{rdma}}$ 的标准系统建模方法。 |
| [55] | Bounds on Multiprocessing Timing Anomalies SIAM Journal on Applied Mathematics, 1969. |
在 4.2 节中用于引出批次间专家重排算法的贪心初值起点——基于长作业优先(LPT)启发式放置规则。 |
| [56] | The SCIP Optimization Suite 9.0 arXiv preprint, 2024. |
在 5.2 节用于指出通用商业/开源求解器在应对多层多微批次的 MILP 时由于分支定界搜索树膨胀而难以在毫秒级预算内产出高质量可行解的工程缺陷。 |
| [57] | Slime: An LLM Post-Training Framework for RL Scaling GitHub Repository, 2025. https://github.com/THUDM/slime/ |
在第 6 节系统实现部分,指出 ReLibra 的强化学习多阶段工作流协调机制是基于开源框架 Slime 深度拓展而来。 |
| [58] | Megatron-LM GitHub Repository, 2025. https://github.com/NVIDIA/Megatron-LM |
在第 6 节中,指出数据平面的分布式 MoE 训练引擎、张量管理与反向传播调度是基于 Megatron-LM 深度修改与嵌入实现。 |
| [59] | SGLang: Efficient Execution of Structured Language Model Programs NeurIPS, 2024. |
在第 6 节中,指出 Rollout 采样引擎基于 SGLang 构建,用于驱动自回归解码并在内核中注入路由捕获器。 |
| [60] | Open Multi-Processing (OpenMP) OpenMP Architecture Review Board, 2026. https://www.openmp.org/ |
在第 6 节中,用于实现控制平面求解器在 128 核 CPU 节点上的跨层、跨微批次多线程高度并发。 |
| [61] | NVIDIA Hopper Architecture In-Depth NVIDIA Technical Blog, 2022. https://developer.nvidia.com/blog/nvidia-hopper-architecture-in-depth/ |
在第 6 节中,用于实现基于 Hopper 硬件特有张量内存加速器(TMA)的异步持久化参数与梯度搬运内核。 |