ReLibra: Routing-Replay-Guided Load Balancing for MoE Training in Reinforcement Learning

发表时间: 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 执行时间为目标：

$$T_{\mathrm{MoE}} = \max_{g \in \mathcal{G}} T_g^{\mathrm{comp}} + \max_{g \in \mathcal{G}} T_g^{\mathrm{comm}}$$


其中 $\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）：

$$\min_{y, z} T_{\mathrm{MoE}} \quad \text{s.t.} \quad \sum_{g \in \mathcal{G}(e)} y_{j,e,g} = 1, \quad y_{j,e,g} \le z_{e,g}, \quad \sum_{e} z_{e,g} \le r$$
并通过增量贪心启发式算法迭代选定瓶颈专家并求解连续线性规划以分配词元，避免了通用求解器的庞大开销。

效果如何 实验在包含 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%，副本同步和求解时间分别被计算和后处理完全掩盖。该方法的局限在于依赖采样与训练路由一致的重放先验；且专家复制受限于单节点内，跨节点的突发失衡仍需依赖批次间重排调节。

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）级不均衡特征：

  1. 领域狭窄与频繁切换：预训练使用海量且多样化的网络语料，而 RL 训练通常针对特定且狭窄的领域（如数学 [21]、代码生成 [22]、指令遵循 [23]）。不同领域的微批次会激活完全不同的热专家集合；当多种领域混合训练时，热专家集合会在微批次之间发生剧烈切换。
  2. 缺乏算法层面的平衡损失机制：预训练通常使用辅助平衡损失（auxiliary balancing losses）[1–3] 或路由偏置（routing bias）[6] 来防止专家欠训练；但在 RL 阶段，由于模型在预训练时已被塑形完成，为了不损害模型生成性能，这些人为约束机制会被彻底移除 [9, 21, 24]，导致 expert 负载分布呈现极高的偏斜度与动态不稳定性。

现有的 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）。

本文的主要创新点与系统贡献如下：

  • 首次揭示并量化了 MoE RL 训练中微批次级负载不均衡的独特动态性，并指出 Rollout 与 Training 阶段的同参、同数据特性所带来的“路由重放”确定性优化窗口。
  • 提出了时间尺度解耦（Timescale Decomposition）的层次化架构：将宏观的跨节点“专家重排（Expert Reordering）”与微观的节点内“专家副本（Expert Replication）”分别映射到批次间（inter-batch）和批次内（intra-batch / micro-batch）的时间尺度上，精确匹配数据中心分层网络带宽（跨节点低带宽 RDMA 与节点内高带宽 NVLink）。

  • 设计了高效求解算法与无额外开销的轻量系统架构：

    • 在批次间，采用基于交换的模拟退火算法（Swap-based Simulated Annealing），结合目标平滑（Log-Sum-Exp）与数据局部性优化，在整个训练批次跨节点均衡计算与通信。
    • 在批次内，设计了层共享副本缓冲区（Layer-Shared Replica Buffer），在 GPU 上仅分配极少的常驻副本槽位并跨所有 Transformer 层复用，将内存占用降至忽略不计（0.13%–1.62%），并使用轻量级持续异步 TMA kernel 将副本同步完全隐藏在 Attention 计算中。
    • 形式化并求解混合整数线性规划（MILP），利用增量贪心启发式算法在毫秒级内同时完成副本放置与细粒度 token 分配。

  • 全面的端到端评估：在涵盖 30B 到 235B 参数规模的 MoE 模型（Qwen3、GLM-4.5 等）及多领域 RL 任务上，ReLibra 相比 Megatron-LM 实现了高达 $1.6\times$ 的训练吞吐提升，相比即使注入先验上帝负载（Oracle Loads）的先进基线 EPLB 仍提升达 $1.2\times$，且吞吐达到理想绝对均衡基线（Balanced）的 $90\% - 94\%$。

2 背景知识与核心观察

2.1 背景知识

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. Rollout 阶段：系统接收一批 Prompts，使用模型以自回归解码（Auto-regressive Decoding）[37, 38] 方式为每个 Prompt 生成一个或多个候选回复（Response）。生成完成后，系统利用外部奖励模型或基于确定规则的验证函数进行后置打分（如 GRPO [24] 或 PPO [39] 的 Critic 模型）。
  2. Post-Rollout 处理与 Training 阶段：在打分阶段，系统可能还会对参考模型（Reference Model）进行前向推理以评估 KL 散度以稳定策略更新。随后，Prompts 与生成的 Responses 及奖励标量被组装为全局训练批次（Global Training Batch）。由于激活显存（Activation Memory）受限，全局批次被进一步切分为若干微批次（Micro-batches）。每个微批次在集群上独立执行前向传播和反向传播并累加梯度。当所有微批次计算完毕后，累积梯度被用于更新参数，新参数同步回 Rollout 实例以启动下一轮循环。

2.2 强化学习中的负载不均衡特征分析（Key Observations）

论文选取了包括推理、代码、指令遵循以及混合领域在内的标准基准（详见下表 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 到 15 倍之间剧烈震荡），不同领域展现出完全不同的激活拓扑。在混合训练场景下，不同微批次激活的专家子集跳跃频繁，导致静态部署方案瞬间失效。


图 5：相邻微批次间负载最高的前 $k$ 个专家的平均重合比例（Intersection Ratio）。数值越低代表热专家切换越剧烈。

观察二：相邻微批次间热专家交集比例普遍低于 50%。为了严格量化时间维度上的不均衡跳变，作者定义了第 $t$ 个与第 $t+1$ 个微批次间 Top-$k$ 最热专家的重合比例（Intersection Ratio）：

$$\frac{|H_t^k \cap H_{t+1}^k|}{k}$$


其中 $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 与跨机架滞后。

2.3 优化机遇与核心挑战

核心机遇：路由重放（Routing Replay）与充裕的决策窗口。

  1. 完全确定的未来路由：在 RL 训练中，Rollout 阶段生成的样本直接输入给后续的 Training 阶段。由于两者执行的 Prompt 和生成的 Token 序列完全一致，并且基于同一组 MoE 权重，其 Router 产生的分值与挑选的 Top-$k$ 专家是完全确定且一模一样的。训练系统在正式启动反向传播计算前，便已经知晓了整个 Training Batch 每一个微批次的精确路由分布，无需进行任何概率猜测或历史滑动平均估计。在工业界近期实践中（如 DeepSeek-V3.2 [20]、DAPO [21]、Kimi [28] 等），为了保证算法收敛稳定性，也主动采用路由重放以确保专家使用生成对应 token 的原样本进行更新。
  2. 宽松的时间窗口：与预训练阶段仅有几十微秒的运行时决策窗口不同，RL 的 Rollout 结束后存在一个显著的时间间隙——即 Post-Rollout Processing（包括外部打分、KL 散度评估等），其耗时通常在数秒到数十秒量级。这为系统执行稍微复杂的全局优化算法提供了充足且不会占用关键路径（Critical Path）的时间预算。

核心挑战：微批次级组合优化的 NP-Hard 属性。
尽管路由信息已知，但为数十个层、数十个微批次实时规划最优的负载均衡方案在计算上极其棘手。系统可采用的手段包括：

  • 专家重排（Expert Reordering）：在 EP 组内对专家的物理驻留 GPU 进行置换。
  • 专家副本（Expert Replication）：为过载的专家在空闲 GPU 上动态创建计算副本。

在显存容量上限、通信带宽拓扑约束下，微批次级的联合规划构成了一个超大规模的组合优化问题。文献已证实，即使完全不考虑副本、仅考虑无显存约束的同构机器完工时间最小化（Makespan Minimization on Identical Machines, $P||C_{\max}$），该问题也是 NP-hard 的 [30]。若加上副本选择与多节点容量约束，问题等价于设施选址问题（Facility Location Problem）[44] 的变种，在微秒/毫秒级时限内求取全局精确最优解是不可能的。

核心设计原则：时间尺度解耦（Timescale Decomposition）。
针对跨节点分布式集群的不对称拓扑（节点内 NVLink 具备数百 GB/s 带宽，而跨节点 RDMA 链路带宽要低得多），ReLibra 提出了时间尺度的层次化解耦：

  • 专家重排（跨节点）放置在批次间时间尺度（Inter-batch Timescale）：在跨节点网络上频繁重排专家需要进行昂贵的参数交换并重构通信路由拓扑。因此，将专家重排作为一种粗粒度的全局平衡手段，基于整个 Training Batch 聚合的累计负载执行一次，将开销分摊到整个训练批次。
  • 专家副本（节点内）放置在批次内时间尺度（Intra-batch Timescale）：专家副本不改变已有常驻专家的分布，仅在热点出现时按需扩容。ReLibra 将副本动态复制严格限制在节点内部，利用极高带宽的 NVLink 在微批次粒度对瞬态负载尖峰进行弹性吸收与本地自适应。

3 ReLibra 系统架构设计概览


图 7：ReLibra 系统的总体分层架构，包含控制平面（求解器）与数据平面（路由捕获器与执行器）。

ReLibra 的整体架构划分为控制平面（Control Plane）与数据平面（Data Plane），如图 7 所示：

3.1 控制平面（Control Plane）

控制平面统筹全局 RL 工作流状态：

  • 路由收集与重放接入：在 Rollout 阶段收集每个样本在各生成步的 Token-to-Expert 路由矩阵。

  • 解耦负载均衡求解器（Load-Balancing Solver）：将复杂的端到端优化拆解为两阶段求解流水线：

    1. 批次间专家重排求解器：依据全局训练批次的聚合路由，求解跨节点的全局专家重排与数据局部性优化方案。
    2. 批次内微批次级副本求解器：依据重排后的拓扑，对各微批次求解节点内局部的专家副本放置与微观 Token 分流比例。

  • 非关键路径重叠：控制平面的计算任务被完全调度在 Post-Rollout 处理阶段（如奖励计算、KL 散度评估期间）异步并行执行，其端到端求解耗时完全脱离训练的关键执行路径。

3.2 数据平面（Data Plane）

数据平面部署于分布式 GPU 训练集群中，由两个核心模块协同驱动：

  • 路由捕获器（Routing Capturer）：紧密嵌入在 Rollout 推理引擎（如 SGLang）的自回归生成循环中。在每步 Token 解码时，该组件捕获 Router 输出的高维 Top-$k$ 专家索引并写入极小的 GPU 显存缓冲区，随后异步搬运至 Host 端聚合为样本级路由决策表并推送到控制平面。
  • 负载均衡执行器（Load-Balancing Executor）：在训练阶段拉取控制平面计算好的调度计划并严格执行：
    • 在每个 Global Training Batch 开启前，驱动全集群执行批次间专家重排，重塑跨节点布局。
    • 在批次内各微批次的前向与反向流水线中，通过层共享副本缓冲区（Layer-Shared Replica Buffer）动态管理热专家权重：在前向传播前通过 SPMD 模式由原主 GPU 主动 Push 参数到同节点的目标 GPU 缓冲区；在反向传播完成后，再将副本累积的梯度同步回拉并合并到主 GPU。整个过程与 Attention 计算异步重叠。

4 批次间专家重排算法细节

4.1 形式化问题建模

批次间专家重排致力于依据整个训练批次聚合的路由数据，最小化整个 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_g^{\mathrm{comp}} = \sum_{e \in E_g} l_e = \sum_{e \in E_g} \sum_{j \in \mathcal{G}} x_{j,e}$$


其中 $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}\}$：

  1. $\mathrm{loc}$ 本地内存直拷：源 GPU 与目标专家位于同一 GPU。
  2. $\mathrm{nv}$ 节点内 NVLink 传输：源 GPU 与目标专家位于同一节点的不同 GPU。
  3. $\mathrm{sr}$ 跨节点同轨道（Same-Rail）RDMA 传输：源 GPU 与目标专家位于不同节点，但处于相同的 Rail（即相同的 GPU 槽位索引），通过 GPUDirect RDMA [48] 直连，无需跨越节点内 NVLink。
  4. $\mathrm{cr}$ 跨节点跨轨道（Cross-Rail）中继传输：源 GPU 与目标专家处于不同节点且 Rail 索引不同；数据必须先通过本地 NVLink 中继到本地同 Rail 的 GPU，再通过 RDMA 传输到目的节点 [49, 50]。

据此，在 Dispatch 阶段，GPU $g$ 注入到 NVLink 上的发送（tx）与接收（rx）通信负载精确形式化为：

$$C_g^{\mathrm{nvlink, tx}} = \sum_{e: \tau(g,e) \in \{\mathrm{nv}, \mathrm{cr}\}} x_{g,e}$$

$$C_g^{\mathrm{nvlink, rx}} = \sum_{e \in E_g} \sum_{j: \tau(j,e) \in \{\mathrm{nv}, \mathrm{cr}\}} x_{j,e}$$
其中 $\mathrm{nv}$ 和 $\mathrm{cr}$ 流量分别构成了直接传输与中继中转的本地网络压力。同理，跨节点 RDMA 链路的负载 $C_g^{\mathrm{rdma, tx}}$ 与 $C_g^{\mathrm{rdma, rx}}$ 可以将上述集合中的 $\{\mathrm{nv}, \mathrm{cr}\}$ 替换为 $\{\mathrm{sr}, \mathrm{cr}\}$ 精确推导。而在 Combine 阶段，由于只是将 token 激活值从目标专家回传至源 GPU，各通信链路上的数据流动路径完全呈对称翻转。

MoE 模块端到端执行时间估算建模。
设 Transformer 隐藏层维度为 $h$，前馈网络中间层维度为 $h'$。每个 token 经过单一专家的 FFN 计算需要执行 3 次 GEMM 操作，对应浮点运算量约为 $6hh'$ FLOPs（忽略极小开销的逐元素激活与门控权重计算）。若单卡有效计算算力为 $F$ FLOPS，则 GPU $g$ 上的专家计算时间为：

$$T_g^{\mathrm{comp}} = \frac{6hh' L_g^{\mathrm{comp}}}{F}$$


若 NVLink 与 RDMA 的实测有效双工通信带宽分别为 $B_{\mathrm{nvlink}}$ 与 $B_{\mathrm{rdma}}$，由于通信链路全双工特性（tx 与 rx 并发）以及先进 All-to-All 算子（如 DeepEP [50]）对 NVLink 与 RDMA 传输流水线的重叠掩盖设计，GPU $g$ 的通信时间建模为：

$$T_g^{\mathrm{comm}} = \max_{\Phi \in \{\mathrm{tx}, \mathrm{rx}\}} \left\{ \max \left( \frac{C_g^{\mathrm{nvlink}, \Phi}}{B_{\mathrm{nvlink}}}, \frac{C_g^{\mathrm{rdma}, \Phi}}{B_{\mathrm{rdma}}} \right) \right\}$$
系统中的物理常数 $F, B_{\mathrm{nvlink}}, B_{\mathrm{rdma}}$ 通过标准的基准测试与插值 Profile 获得 [51–54]。由于跨卡同步屏障的存在，分布式 EP 组总的 MoE 模块执行完工时间受限于木桶效应中最慢的落后者 GPU：
$$T_{\mathrm{MoE}} = \max_{g \in \mathcal{G}} T_g^{\mathrm{comp}} + \max_{g \in \mathcal{G}} T_g^{\mathrm{comm}}$$

4.2 基于交换的模拟退火搜索算法（Swap-based Simulated Annealing）

由于 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$ 算子：

$$\mathrm{LSE}_\beta(z_1, \dots, z_n) = \frac{1}{\beta} \log \sum_{k=1}^n e^{\beta z_k}$$


当缩放因子 $\beta \to \infty$ 时，$\mathrm{LSE}_\beta \to \max$。平滑后的连续优化目标函数定义为：

$$\widetilde{T}_{\mathrm{MoE}} = \mathrm{LSE}_\beta \left( \widetilde{T}_g^{\mathrm{comp}}, g \in \mathcal{G} \right) + \mathrm{LSE}_\beta \left( \widetilde{T}_g^{\mathrm{comm}}, g \in \mathcal{G} \right)$$
其中通信时间的内部双工 $\max$ 算子同样经过 LSE 平滑。该平滑函数赋予了次瓶颈 GPU 改善时以明确的负梯度奖励。通过超参数对比，ReLibra 确定选取 $\beta = 20$ 作为生产默认参数。

数据局部性二次优化（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 绑定映射表。

5 批次内专家动态副本算法细节

5.1 层共享副本缓冲区设计（Layer-Shared Replica Buffer）

在微批次尺度上动态创建热专家副本面临着显存容量与同步延迟的双重严苛约束。


图 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 通信策略：

  • 前向传播（Forward）：拥有热专家的源卡主动利用 NVLink 将权重推送到同节点目标卡的共享副本缓冲区中。该同步操作与当前微批次的前置 Attention 计算高度并行。
  • 反向传播（Backward）：目标卡上的副本完成反向微批次计算后，将其产生的专家梯度通过 NVLink 推回至原宿主 GPU 进行累加。同样，梯度回传被完全重叠在反向阶段的 Attention 反向计算之内。

5.2 微批次级副本放置与分流联合规划算法

混合整数线性规划（MILP）严格数学建模。
针对某一层特定微批次，ReLibra 将专家副本选择、放置节点与 Token 切分比例联合形式化为一个标准的混合整数线性规划问题。
对于给定的微批次，设输入矩阵为 $\{x_{j,e}\}$。设 $\mathcal{G}(e) \subseteq \mathcal{G}$ 表示与专家 $e$ 的宿主卡同处一个物理节点内的所有候选 GPU 集合。
引入决策变量：

  • $z_{e,g} \in \{0, 1\}$：二值整数决策变量。表示是否允许在 GPU $g$ 上放置专家 $e$ 的副本。对于专家 $e$ 原生驻留的本地 GPU $g_{\mathrm{home}}$，强制恒等于 1；
  • $y_{j,e,g} \in [0, 1]$：连续分流决策变量。表示源自 GPU $j$、路由给专家 $e$ 的 token 集合中，最终实际分配给 GPU $g$ 上的实例（原始实体或副本）处理的具体比例。

综上，该微批次的 MoE 执行完工时间优化问题定义为：

$$\min_{y_{j,e,g} \in [0,1], z_{e,g} \in \{0,1\}} T_{\mathrm{MoE}} = T^{\mathrm{comp}}(y, z) + T^{\mathrm{comm}}(y, z)$$


受限于以下物理约束条件：

$$\sum_{g \in \mathcal{G}(e)} y_{j,e,g} = 1, \quad \forall j, e \text{ with } x_{j,e} > 0$$
$$y_{j,e,g} \le z_{e,g}, \quad \forall j, \forall e, \forall g \in \mathcal{G}(e)$$
$$z_{e,g} = 1, \quad \forall g \in \mathcal{G}, \forall e \in E_g$$
$$\sum_{e \notin E_g} z_{e,g} \le r, \quad \forall g \in \mathcal{G}$$

约束逻辑推导解析：

  • 约束 (8)：守恒完备性约束。保证对任意非空的路由请求 $(j,e)$，分流比例之和严格等于 1，确保原始 Token 数量 $x_{j,e}$ 被完全守恒切分为各 GPU 上的分流项 $x_{j,e} \cdot y_{j,e,g}$。
  • 约束 (9)：逻辑绑定约束。强制约束只有当 GPU $g$ 上实际放置了专家 $e$ 的副本（即 $z_{e,g} = 1$）或该专家原生就在该卡上时，分配给它的 Token 比例 $y_{j,e,g}$ 才能大于 0；否则只能为 0。
  • 约束 (10)：主实体固定约束。锁定原生物理位置的二值变量为 1。
  • 约束 (11)：层共享显存槽位硬约束。保证任意单一 GPU 为外来专家开辟的临时副本总数不得超过共享缓冲区容量上限 $r$。

由于目标函数中的最大化算子可通过引入辅助变量进行标准线性化，整个模型在数学上是一个紧凑的 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

6 系统实现细节

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] 进行多核多线程并行化编排：

  • 对于批次间专家重排，OpenMP 线程池按“层数 $\times$ 随机种子数”将退火任务均匀并行分发；
  • 对于批次内微批次级副本求解，各线程独立接管单一“层数 $\times$ 微批次”维度的贪心 LP 求解过程。整个求解完全并发无锁化运行。

基于持久化异步 TMA 内核的显存缓冲区通信实现。
针对 NVIDIA Hopper 架构的硬件特质，ReLibra 使用张量内存加速器（Tensor Memory Accelerator, TMA）[61] 编写了持久化异步数据传输 CUDA 内核。为了防止参数搬运抢占 Attention 计算的算力核心，该持久化内核被严格硬编码限制在极少量的流式多处理器（Streaming Multiprocessors, SMs）上执行。绝大多数 SMs 得以全速投入到计算密集的 FlashAttention 内核中，真正实现了零竞争、零抖动的通信-计算异步重叠。

7 实验环境配置

7.1 硬件配置

  • 测试集群规模：20 个物理计算节点（Total 160 卡）。
  • 计算加速芯片：每节点配备 8 张 NVIDIA Hopper 架构高端 GPU，单节点配备 1.5 TB Host 物理内存。

  • 互联网络拓扑：

    • 跨节点通信：轨道优化（Rail-Optimized）的 InfiniBand RDMA 网络，每节点聚合双向网络带宽为 $8 \times 400\ \mathrm{Gbps}$。
    • 节点内卡间互联：第三代高带宽 NVLink，单向/双向带宽高达 $900\ \mathrm{GB/s}$。

  • 独立求解节点：控制平面求解器部署于一台独立的 CPU-only 算力节点上，配备 128 核 Intel Xeon 高性能 CPU，由 Ray [64] 统一管理任务分发。

7.2 软件与依赖环境

  • 操作系统与系统驱动：Linux OS，NVIDIA GPU Driver 535.161.08。
  • 底层依赖与算子库：CUDA 12.6, PyTorch 2.8.0, FlashAttention-3 [62], TransformerEngine 2.7.0 [63], DeepEP 1.1.0 [50], OpenMP [60]。
  • 框架支撑：Slime [57]（RL 编排）、Megatron-LM [58]（训练后端）、SGLang [59]（Rollout 后端）。

7.3 评测模型与混合并行切分策略

实验选用了来自不同机构、横跨三个数量级规模的先进生产级 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）。

7.4 对比基线系统（Baselines）

  1. Megatron-LM [58]：业界标准的最优原生开源框架，深度集成了 DeepEP 高性能 All-to-All 算子以及 GroupedGEMM 计算内核。
  2. PopFetcher [17]：基于历史相邻层激活模式预测当前批次热专家并进行预取副本的系统，在实验中经过滑动窗口精细微调。
  3. $\mathbf{EPLB^+}$：文献提出的基于历史启发式的静态专家平衡系统 EPLB [25]，此处被赋予了“上帝视角”的 Oracle Loads（直接输入通过路由重放获知的真实负载），在全局批次粒度执行贪心重排与副本放置。
  4. $\mathbf{LPLB^+}$：文献提出的基于线性规划的动态 token 分流系统 LPLB [45]，同样注入了真实 Oracle Loads，每微批次求解 LP 分流，但算法硬性限制每个专家至多拥有 1 个副本。
  5. Balanced（理想上限）：强制将所有 token 均匀随机散布给所有专家的人为构造最优基线。该基线无法保留模型真实路由逻辑，仅作为硬件无不均衡条件下的性能理论上限。

8 实验结果与详细分析

8.1 整体吞吐性能对比


图 10：在混合数据集（Mixed Dataset）上，不同模型规模下的端到端训练吞吐表现（tokens/s）。


图 11：在 Qwen3-235B-A22B 超大模型上，不同业务领域数据集下的训练吞吐对比。

核心实验发现：

  • 大幅领先行业已有方案：如图 10 与图 11 所示，ReLibra 在各种配置下均大幅超越现有系统。相较于原生 Megatron-LM，吞吐提升幅度达到 $1.21\times$ 至 $1.58\times$（在 Qwen3-235B 的 Coding 数据集上接近 $1.6\times$）；相较于预取机制的 PopFetcher，性能提升 $1.17\times$ 至 $1.27\times$；相较于拥有上帝视角的 $\mathrm{EPLB^+}$ 与 $\mathrm{LPLB^+}$，加速比依旧稳定保持在 $1.06\times$ 至 $1.21\times$。
  • 逼近理论均衡上限：在所有测试集上，ReLibra 达到了理想全均衡基线（Balanced）$90\% - 94\%$ 的吞吐。这证实其时间尺度解耦架构成功将微批次维度的绝大多数失衡气泡压榨殆尽。
  • 深入洞察 $\mathrm{EPLB^+}$ 与 $\mathrm{LPLB^+}$ 的内在局限：$\mathrm{LPLB^+}$ 在指令遵循和混合领域表现优于 $\mathrm{EPLB^+}$，但在数学推理与代码任务中却反向落后于 $\mathrm{EPLB^+}$。其根源在于：$\mathrm{LPLB^+}$ 为限制求解时间，硬编码限制每个热专家至多只能有 1 个副本，在面对高度极化偏斜的长尾尖峰时分流能力不足；而 $\mathrm{EPLB^+}$ 缺乏微批次级适应。ReLibra 借助预规划将微批次联合规划完全移出关键路径，同时突破了单副本限制与静态规划缺陷。

8.2 逐层递进消融实验（Ablation Study）

表 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\%$（推理）和 $15\%$（混合），说明宏观拉平各物理节点负载消除了最严重的跨机架通信和算力拥塞；
  • 在此拓扑基础上叠加微批次级动态副本，吞吐进一步大幅跳跃 $34\%$ 和 $24\%$，这印证了解耦设计的有效性：批次间重排为节点内副本吸收微观脉冲提供了良好的先决基底。


图 12：批次间专家重排机制消融分析（对比原生 EP、贪心 LPT 专家放置 ReLibra-Expert、贪心数据分配 ReLibra-Data 以及完整退火 ReLibra-Inter）。


图 13：批次内动态副本机制消融分析（对比原生基线、仅微批次贪心副本放置 ReLibra-Replication、仅微批次连续分流 ReLibra-Splitting 以及完整联合 MILP 启发式 ReLibra-Intra）。

深入机制微观消融：

  1. 批次间专家重排剖析（图 12）：相比于两类贪心基线（仅用 LPT 贪心放置专家的 ReLibra-Expert，以及仅根据样本长度做局部性贪心分配的 ReLibra-Data），ReLibra-Inter 实现了 $1.05\times - 1.16\times$ 的超额加速。这证明单纯使用贪心算法无法兼顾“计算均衡”与“网络拓扑”，而平滑模拟退火与数据局部性二次退火能找到兼顾两者的全局折中点。
  2. 批次内专家动态副本剖析（图 13）：对比仅做微批次离散副本放置的 ReLibra-Replication 与仅做无副本限制连续 Token 分流的 ReLibra-Splitting，ReLibra-Intra 实现了 $1.12\times - 1.24\times$ 和 $1.09\times - 1.12\times$ 的显著提升。这表明将离散放置与连续分流视为孤立问题均不可行，MILP 联合建模与迭代贪心分流是实现高精度平抑波动的核心。

8.3 不均衡抑制能力案例研究（Case Study）


图 14：不同硬件 EP 并行规模与不同领域数据集下，各 GPU 间 Rank-Level 负载偏斜度箱线图分布。

作者统计了所有参与运算的 GPU 卡间实际处理 token 数量的 Rank 级偏斜度（Rank-Level Skewness，即 $\max_g L_g / \mathrm{avg}_g L_g$）。如图 14 所示：

  • 随并行度扩展的表现（图 14a）：在原生 Megatron-LM 下，随着 EP 规模从 8 扩大到 64，单卡承载的专家总数减少，卡间失衡迅速恶化，最大偏斜度一路恶化飙升至 6.0 左右；而 ReLibra 在 EP=8、16、32、64 下的平均偏斜度分别压制在惊人的 1.00、1.03、1.06、1.08，几乎保持绝对平衡。
  • 跨领域的通用鲁棒性（图 14b）：在 235B 模型的全领域测试中，面对原生系统在数学与代码场景下高达 5.8 的剧烈失衡，ReLibra 的偏斜度箱线图稳稳贴合在 1.02 至 1.07 之间，展现出跨任务模式的稳定性。

8.4 系统固有开销分析（System Overhead）


图 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 算子执行时间的异步重叠评估。

  • 求解时间开销（图 15）：实验数据显示，得益于 C++ OpenMP 增量更新与启发式剪枝，即使面对层数最多的 Qwen3-235B（94 层），总求解耗时亦仅为约 17 秒，而对应的 Post-Rollout 阶段耗时超过 20 秒。所有求解用时被 $100\%$ 隐藏在非关键路径之后，对训练启动造成零停顿。
  • 批次间专家重排网络开销：在 Global Batch Size 为 1024 的环境下，跨节点交换专家参数的耗时仅占整个批次训练执行总耗时的 $1.4\% - 2.1\%$，平摊后的时间代价极低。
  • 动态副本显存开销（表 5）：在每张卡配置 2 个副本槽位（$r=2$）时，各模型的参数显存额外占用仅为 18 MiB 至 72 MiB，在卡内总模型参数显存中的占比处于 $0.13\%$ 至 $1.62\%$ 之间。反向梯度显存与参数显存完全一致，未引发任何 OOM 风险。
  • 动态副本通信与 Attention 掩盖重叠（图 16）：前向参数推送耗时仅 1–3 ms，远低于 Attention 前向计算耗时（3–6 ms）；反向梯度推送耗时约 2–5 ms，同样远低于反向 Attention 计算耗时（5–16 ms）。由于持久化 TMA 内核限制在极少 SM 运行，Attention 算子在掩盖同步前后的执行耗时几乎没有变化，实现了近乎完美的硬件并发。

9 补充细节：Off-Policy 强化学习拓展讨论与相关工作

9.1 Off-Policy 强化学习下的适用性讨论

在某些 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 算法演进方向高度契合。

9.2 相关工作深入对比

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 在训练反向中的微批次不均衡问题，构成了端到端系统拼图的关键一环。

10 结论

本文针对大模型强化学习（RL）训练中 MoE 架构出现的微批次维度高度剧烈波动的负载不均衡难题，提出了 ReLibra 系统。ReLibra 巧妙地捕捉并利用了 RL 工作流中固有的“路由重放”机遇，彻底摒弃了传统预训练系统所依赖的脆弱历史预测假设，在训练启动前获取了微批次级确切路由。系统提出了时间尺度解耦设计理念，将跨节点宏观专家重排分配于批次间尺度，将节点内细粒度动态副本分配于微批次尺度，精准匹配了集群的分层互联网络带宽特性。通过轻量级的层共享显存复用缓冲区、高效的模拟退火和混合整数线性规划启发式算法，ReLibra 在不侵占任何训练关键路径的前置时间窗口内，以近乎忽略不计的系统开销达成了微批次级的精细负载均衡。端到端实验表明，ReLibra 取得了相比于 Megatron-LM 最高达 $1.6\times$ 的显著吞吐提升，逼近了理论极限性能。

11 方法细节引用文献汇总

下表系统整理了第 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）的异步持久化参数与梯度搬运内核。