ForeMoE：利用路由预见性实现 RL 后训练中 MoE 微步级负载均衡

笔记日期： 2026-06-13 笔记作者： Zhongzhu Zhou 论文标题： Harnessing Routing Foresight for Micro-step-level MoE Load Balancing in RL Post-training 作者： Yuming Zhou, Haoyang Li, Sheng Lin, Yanfeng Zhao, Tong Zhao, Xupeng Miao, Jie Jiang, Fangcheng Fu, Bin Cui arXiv： 2606.11867v1，2026-06-11 状态： arXiv 预印本（北京大学 / 上海交通大学 / 腾讯）

一句话总结

ForeMoE 是首篇系统性研究 MoE 模型强化学习后训练中专家负载不均衡问题的工作，它发现 RL 后训练的负载波动发生在微步（micro-step）层面而非步骤层面，并利用 RL 特有的”路由回放”（router replay）结构——rollout 阶段记录的路由决策在后续阶段被精确复用——将负载均衡从在线预测问题转化为离线规划问题，最终在 64 张 GPU 上实现最高 1.45× 端到端加速。

前置知识

在深入技术细节之前，需要先建立以下背景知识。

Transformer 与 FFN 层

标准 Transformer 由多头注意力层和前馈网络（FFN）层交替堆叠组成。对于隐藏向量 $x \in \mathbb{R}^d$，FFN 的计算为：

$$\text{FFN}(x) = W_2 \cdot \sigma(W_1 x + b_1) + b_2$$

FFN 参数量约占模型总参数的三分之二，且每个 token 都要经过同样的 FFN 计算——随着模型规模增大，这一计算代价迅速攀升。

专家混合模型（MoE）架构

MoE 架构用 $E$ 个并行专家网络替代稠密 FFN，加上一个轻量级路由器 $G$：

$$\text{MoE}(x) = \sum_{i \in \text{Top-}K(G(x))} g_i(x) \cdot f_i(x)$$

路由器机制：路由器是一个线性层，产生 logits $r = W_{\text{router}} x \in \mathbb{R}^E$，然后对 top-$K$ 个 logits 做 softmax 得到门控权重：

$$g_i(x) = \frac{\exp(r_i)}{\sum_{j \in \text{Top-}K} \exp(r_j)}, \quad i \in \text{Top-}K(r)$$

以 $K=2$，$E=64$ 为例，每个 token 只激活 64 个专家中的 2 个，理论计算量降至稠密模型的 3.1%。当前主流大型 MoE 模型包括：

• DeepSeek-V3/R1：671B 参数，每 token 激活 37B；每层 256 个专家，top-8 路由
• Qwen3-30B-A3B：30B 参数，激活约 3B；每层 64 个专家，top-2 路由（ForeMoE 实验使用）

专家并行（EP）与 All-to-All 通信

分布式训练中，$E$ 个专家被分配到 $R$ 张 GPU 上（称为”rank”），每张 GPU 负责 $E/R$ 个专家。每次 MoE 层的前向传播都需要两轮 All-to-All 集合通信：

• Dispatch（分发）：每张 GPU 将 token 发送到持有其对应专家的 GPU
• Combine（收集）：每张 GPU 将专家计算结果返回给 token 的源 GPU

EP 示意（E=4, R=2, K=1）：

  Rank 0 的 token: [T0→E2, T1→E0, T2→E3, T3→E0]
  Rank 1 的 token: [T4→E1, T5→E3, T6→E0, T7→E2]

  Dispatch All-to-All：
    Rank 0 → Rank 0：T1, T3（E0 在 Rank 0）
    Rank 0 → Rank 1：T0, T2（E2, E3 在 Rank 1）
    Rank 1 → Rank 0：T6    （E0 在 Rank 0）
    Rank 1 → Rank 1：T4, T5, T7（E1, E3, E2 在 Rank 1）
  
  并行计算→ Combine All-to-All（逆向）

All-to-All 的延迟与负载最重的 rank 处理时间成正比。如果某个 rank 接收的 token 是平均值的 2 倍，整个 All-to-All 就慢 2 倍。

负载不均衡率定义为：

$$\rho^{(m)} = \frac{\max_r L_r^{(m)}}{\bar{L}^{(m)}}, \quad \bar{L}^{(m)} = \frac{B_m \cdot K}{R}$$

其中 $B_m$ 是微步的 token 数，$R$ 是 rank 数。$\rho = 1$ 表示完美均衡，$\rho = 2$ 表示最慢 rank 处理了 2 倍于平均的 token。

RL 后训练流程

RL 后训练（使用 PPO、GRPO、DAPO 等算法）将预训练模型对齐到人类偏好或提升推理能力。标准流程包含三个顺序阶段：

阶段 1 — Rollout（采样）：当前策略模型 $\pi_\theta$ 对一批 prompt 进行自回归解码，生成回复。所有 token 在所有层的路由决策被完整记录，形成路由表 $\mathcal{R}$。

阶段 2 — Recompute（重计算）：训练框架（Megatron/FSDP）对生成的 token 重新做前向传播，计算对数概率 $\log \pi_\theta(y|x)$。这是必要的，因为推理引擎与训练框架使用不同的并行策略，不能直接复用推理阶段的计算结果。由于显存限制，序列被分成 $M$ 个微步（micro-step）分批处理。

阶段 3 — Policy Update（策略更新）：计算策略梯度（如 GRPO 目标），同样分 $M$ 个微步积累梯度，最后一次性更新模型参数。

路由回放（Router Replay）：关键约束——Recompute 和 Policy Update 阶段必须使用与 Rollout 完全相同的 token-to-expert 路由决策，以保证重要性采样的正确性。这意味着：路由表 $\mathcal{R}$ 在 Rollout 结束后就完全固定，后续所有微步的专家分配都提前已知。

图 1 对比了预训练和 RL 后训练的负载特征差异：

graph LR
    subgraph PreTrain["预训练（步骤级）"]
        PT1["步骤开始前：\n路由未知\n→ 依赖历史 EMA 预测"]
        PT2["步骤执行：\n大批次平均效应\n→ 负载相对稳定"]
    end
    subgraph RLTrain["RL 后训练（微步级）"]
        RL1["Rollout：\n记录路由表 R\n→ 后续路由完全已知"]
        RL2["Recompute/Policy Update：\n微步批次小\n→ 负载高频波动"]
    end
    PT1 --> PT2
    RL1 --> RL2

图 1：预训练（步骤级统计稳定，路由不可知）与 RL 后训练（微步级高频波动，但路由提前已知）的对比。这一结构差异是 ForeMoE 设计的核心出发点。

核心问题：RL 后训练中的微步级负载不均衡

两层特征：步骤级稳定，微步级高波动

ForeMoE 在 Qwen3-30B-A3B 上的实测揭示了 RL 后训练负载的两个特征：

步骤级：稳定但倾斜。 RL 后训练聚焦于数学、代码等集中领域，预训练已形成专家特化——某些专家专门处理数学推理 token，另一些处理代码相关 token。由于任务领域集中，这些”热门专家”在每个训练步骤都持续获得大量 token，步骤级负载分布高度稳定但严重倾斜。

微步级：高度可变且倾斜。 每个训练步骤被分成 $M$ 个微步，每个微步只包含少量序列。以 $B_m = 32$ 条序列、$K=2$、$E=64$ 为例，每个微步只有约 $32 \times 512 \times 2 = 32{,}768$ 个 token 路由分配。在如此小的批次下，统计平均效应失效，专家 token 数服从泊松分布：

$$c_{e,m} \sim \text{Poisson}(B_m \cdot K \cdot p_e)$$

变异系数（CV = 标准差/均值）为 $1/\sqrt{B_m K p_e}$，对某个 $p_e = 0.05$ 的专家，CV ≈ 14%。对应的每个 rank 的负载不均衡率可轻易达到 $\rho = 1.5\sim2.3$。

为什么预训练方法失效

1. EMA 预测的是步骤级均值：$\hat{c}_e \approx \bar{c}_e$，完全无法捕捉微步级偏差 $\delta_{e,m} = c_{e,m} - \bar{c}_e$
2. 规划时间预算极短：预训练负载均衡触发频率低（每 $T$ 步一次），有足够时间做优化；微步级均衡需要在微步之间（约 10ms 内）完成规划
3. 专家迁移频率过高：预训练中专家迁移代价可以在多个步骤内摊销；微步级均衡需要每个微步都可能重新布局专家，迁移效率要求极高

关键机遇：路由可预知性

RL 后训练的路由回放约束反而成为了机遇。Rollout 结束后，系统已经知道 Recompute 和 Policy Update 中每个微步的精确 token 分配：

$$\mathcal{R} = \{(t, \ell, e_{t,\ell,k}) : t \in \text{all tokens}, \ell \in \text{all layers}, k \in [K]\}$$

这将负载均衡从”在线预测”转化为”离线规划”——不再需要猜测未来，只需针对已知的精确未来负载计算最优方案。

ForeMoE 系统设计

整体架构

ForeMoE 由两个核心模块组成：四阶段规划器（Four-Stage Planner）和专家迁移引擎（Expert Transfer Engine）。

flowchart TB
    subgraph Input["输入"]
        R["路由表 R（Rollout 后）"]
    end
    subgraph Planner["四阶段规划器"]
        S1["阶段 1：基础布局优化器 BPO\n（每步执行一次）\n输入：步骤级均值负载\n输出：基础布局 P*"]
        S2["阶段 2：微步偏差调整器 MDA\n（每微步执行一次）\n输入：精确路由 R[m] + P*\n输出：调整方案 A[m]"]
        S3["阶段 3：Token 重分配验证\n检查路由回放一致性"]
        S4["阶段 4：负载验证\n计算调整后 ρ"]
    end
    subgraph Transfer["专家迁移引擎"]
        T1["Recompute 阶段：CPU 辅助路径\nGPU→CPU DRAM→对端 CPU→GPU\nPCIe DMA，并行于 GPU 计算"]
        T2["Policy Update 阶段：GPU Direct\nNVLink 直连迁移\n~600 GB/s，低延迟"]
    end
    R --> S1
    R --> S2
    S1 --> S2
    S2 --> S3
    S3 --> S4
    S4 --> T1
    S4 --> T2

图 2：ForeMoE 系统架构。路由表驱动整个规划流程，规划结果由专家迁移引擎执行，两种硬件路径分别服务于不同的 RL 训练阶段。

问题形式化

输入：

• $E$ 个专家，$R$ 个 rank，$M$ 个微步
• 每个专家的显存代价 $m_e$，每个 rank 的显存预算 $C$
• 来自路由表的精确负载：$c_{e,m}$（专家 $e$ 在微步 $m$ 的 token 数）

决策变量：

• $x_{e,r} \in \{0,1\}$：专家 $e$ 的主副本放在 rank $r$
• $y_{e,r} \in \{0,1\}$：专家 $e$ 在 rank $r$ 有副本

目标——最小化所有微步的最大负载（makespan）：

$$\min \quad \max_{r \in [R],\; m \in [M]} \sum_{e=1}^{E} z_{e,r,m} \cdot c_{e,m}$$

约束条件：

$$\text{(Memory)} \quad \sum_e (x_{e,r} + y_{e,r}) \cdot m_e \leq C \quad \forall r$$ $$\text{(Coverage)} \quad \sum_r x_{e,r} \geq 1 \quad \forall e$$ $$\text{(Routing consistency)} \quad z_{e,r,m} \leq x_{e,r} + y_{e,r} \quad \forall e,r,m$$

NP 困难性（§5 详细证明）：通过归约自装箱问题（Bin Packing）可以证明该优化问题是 NP 困难的。在 Qwen3-30B-A3B 的设置下（64 个专家，8 个 rank），枚举所有可能布局需要评估 $8^{64} \approx 10^{57}$ 种方案——完全不可行。

四阶段规划器

关键分解原理：将每个微步的负载分解为步骤级均值与微步级偏差之和：

$$L_r^{(m)} = \bar{L}_r + \delta_r^{(m)}$$

由于步骤级负载 $\bar{L}_r$ 稳定，可以用一个”基础布局”来最小化平均情形；而偏差 $\delta_r^{(m)}$ 虽然波动大，但搜索空间小（只需调整几个专家），可以快速规划。

阶段 1 — 基础布局优化器（BPO）

每步执行一次，利用路由表中的步骤级均值负载 $\{\bar{c}_e\}$ 计算稳定基础布局：

算法 1：基础布局优化器（BPO）
────────────────────────────────────────────────────────────────────
输入：步骤级均值负载 {c̄_e}，上一步布局 P_prev（热启动），
      显存预算 C，复制预算 δC
输出：基础布局 P*

1. P = P_prev  （热启动：充分利用步骤级负载稳定性）
2. 计算每个 rank 的当前负载：load(r, P) = Σ_{e on r} c̄_e / R
3. 理想负载：ideal = total_tokens * K / R

4. 【专家迁移阶段】
   按 |load(host(e), P) - ideal| 降序排列专家
   对每个专家 e：
       r_cur = 当前所在 rank
       对每个候选 rank r'（按负载升序）：
           若显存允许 且 load(r',P) + c̄_e < load(r_cur,P)：
               将 e 迁移至 r'：更新 P, load(r_cur), load(r')
               Break

5. 【专家复制阶段】
   找出过载 rank O = {r : load(r,P) > (1+ε)·ideal}
   对每个 r_o ∈ O（按负载降序）：
       e* = r_o 上负载最重的专家
       找显存充裕且负载最低的 rank r_u：
           在 r_u 上复制 e*，两端各承担约 c̄_{e*}/2

6. 返回 P*
────────────────────────────────────────────────────────────────────

阶段 2 — 微步偏差调整器（MDA）

每个微步执行一次，根据当前微步的精确路由 $R[m]$ 计算偏差修正方案：

算法 2：微步偏差调整器（MDA）
────────────────────────────────────────────────────────────────────
输入：基础布局 P*，当前微步路由 R[m]（精确 token 数 c_{e,m}），
      偏差阈值 ε
输出：调整方案 A[m] = [(专家, 源 rank, 目标 rank), ...]

1. 基于 P* 和 R[m] 计算各 rank 实际负载 L_r
2. 过载 rank：O = {r : L_r > (1+ε)·ideal_m}
3. 欠载 rank：U = {r : L_r < (1-ε)·ideal_m}
4. 若 O 为空：A[m] = {} 无需调整，直接返回
5. 对每对 (r_o, r_u)（按负载差降序）：
       e* = r_o 上在本微步 token 数最多的专家（且不超调）
       若 c_{e*,m} ≤ 过载量 × 1.5：
           A[m] 中添加 (e*, r_o, r_u)
           更新 L_{r_o} -= c_{e*,m}，L_{r_u} += c_{e*,m}
6. 返回 A[m]
────────────────────────────────────────────────────────────────────

阶段 3 — Token 重分配验证：验证调整方案 A[m] 与路由回放约束一致——专家迁移后，该专家的 token 通过 All-to-All 被重定向到新 rank，逻辑路由（哪个 token 去哪个专家）不变。

阶段 4 — 负载验证：计算调整后的不均衡率 $\rho$，若超过阈值则回退到基础布局（接受次优均衡以避免迁移开销过大）。

专家迁移引擎

规划完成后，需要在微步计算开始前完成专家权重的物理迁移。核心创新是为不同 RL 阶段选择不同硬件路径：

graph TB
    subgraph Recompute["Recompute 阶段 — CPU 辅助路径"]
        direction LR
        GPU0["GPU 0\n（源专家权重）"]
        CPU0["CPU DRAM 0\n（中转缓冲）"]
        CPU1["CPU DRAM 1\n（中转缓冲）"]
        GPU1["GPU 1\n（目标 rank）"]
        GPU0 -- "D2H DMA\nPCIe ~32 GB/s" --> CPU0
        CPU0 -- "高速网络\nRDMA" --> CPU1
        CPU1 -- "H2D DMA\nPCIe ~32 GB/s" --> GPU1
    end
    subgraph PolicyUpdate["Policy Update 阶段 — GPU Direct"]
        direction LR
        GPU2["GPU 0\n（源专家权重）"]
        GPU3["GPU 1\n（目标 rank）"]
        GPU2 -- "NVLink 4\n~600 GB/s\n直连" --> GPU3
    end

图 3：专家迁移引擎的双路径设计。Recompute 阶段用 CPU 辅助路径（经由 CPU DRAM 中转），与 GPU 计算并行执行；Policy Update 阶段用 GPU Direct NVLink，带宽高出 PCIe 约 20 倍。

为什么 Recompute 用 CPU 路径？ Recompute 阶段 GPU 正在执行 MoE 前向传播计算。CPU 辅助 DMA 传输通过 PCIe 总线并行进行，不打断 GPU 计算。以 Qwen3-30B-A3B 为例，每个专家权重约 12.5 MB，PCIe 传输时间约 0.4ms，而微步计算约 80–120ms——传输完全隐藏在计算窗口内。

为什么 Policy Update 用 GPU Direct？ Policy Update 涉及频繁的梯度同步，通信在关键路径上。NVLink 4 的 ~600 GB/s 带宽比 PCIe 快约 20 倍，12.5 MB 专家权重传输只需约 20µs，可轻松嵌入算子间隙。

流水线重叠：迁移引擎在微步 $m$ 执行计算的同时，异步准备微步 $m+1$ 所需的专家布局：

t = 0ms:     微步 m 开始计算
t = 0ms:     异步：MDA 规划微步 m+1 的调整方案（CPU 上 ~5ms）
t = 5ms:     专家迁移开始（CPU 路径或 GPU Direct，取决于阶段）
t = ~85ms:   微步 m 计算结束
t = ~85ms:   微步 m+1 的迁移已完成
t = 85ms:    微步 m+1 使用新布局开始计算

常见情况下，迁移完全被计算遮蔽，有效开销为零。

NP 困难性证明（摘要）

定理：判断是否存在不均衡率 $\leq \rho^*$ 的布局方案，是 NP 完全问题。

证明思路（归约自装箱问题）：

给定装箱问题实例：物品 $\{s_1, \ldots, s_n\}$，箱容量 $W$，目标 $R$ 个箱。构造负载均衡实例：

• $E = n$ 个专家，专家 $e$ 的显存代价 $m_e = s_e$
• $R$ 个 rank，每个显存预算 $C = W$
• 单微步（$M=1$），专家 $e$ 的 token 数 $c_e = s_e$
• 目标不均衡率 $\rho^* = 1$（完美均衡）

存在完美均衡布局当且仅当物品恰好能被等量装入 $R$ 个容量 $W$ 的箱中，即等价于装箱问题可行性。由于装箱问题是 NP 完全的，负载均衡决策问题也是 NP 完全的。$\square$

实践含义：对于 64 专家、8 rank 的设置，枚举所有布局需要 $8^{64} \approx 10^{57}$ 次——ForeMoE 的层次化分解（BPO+MDA）将其转化为 $O(E \cdot R)$ 级别的贪心步骤，在 5ms 内完成规划。

实验结果

实验设置

• 模型：Qwen3-30B-A3B（30B 参数，激活约 3B；每层 64 专家；top-2 路由；8 个专家并行 rank）
• 集群：64 张 H100 级 GPU，NVLink 4 互连，Megatron-LM 训练框架，vLLM rollout 引擎
• RL 算法：GRPO，全局批次 256 条序列，8 个微步/步骤
• 数据集：DAPO-Math-17k（数学推理）、CodeForces（竞赛编程）
• 基线：无负载均衡（No LB）、步骤级 EMA 均衡（EMA LB）、全知先知（Oracle）

端到端加速

方法	DAPO-Math	CodeForces	GPU 利用率	峰值不均衡率
No LB	1.00×	1.00×	62%	2.3×
EMA LB	1.08×	1.06×	67%	1.88×
ForeMoE	1.45×	1.38×	88%	1.21×
Oracle	1.52×	1.46×	92%	1.08×

ForeMoE 达到了 Oracle 性能的 95%（1.45× vs 1.52×），同时 GPU 利用率从 62% 提升到 88%。EMA 基线仅有 8% 提升，进一步验证了步骤级统计对微步级不均衡毫无作用。

图 4 展示了微步级不均衡率的分布对比：

微步不均衡率分布对比（DAPO-Math-17k, 64 GPU）:

  不均衡范围    No LB      EMA LB   ForeMoE
  ρ < 1.1      ██░░░░░░   ████░░░░  ████████████████████████ 61%
  1.1-1.3      ████░░░░   ████████  ████████████ 32%
  1.3-2.0      ████████████████░░   ████████████████ 47%   ███ 6%
  ρ ≥ 2.0      █████████████░░░░   ██████ 15%   ░ 0%

  No LB:    3% / 12% / 52% / 33%
  EMA LB:   9% / 21% / 55% / 15%
  ForeMoE: 61% / 32% /  6% /  0%

图 4：ForeMoE 将 93% 的微步的不均衡率控制在 1.3× 以下，而无负载均衡时只有 15% 的微步满足此条件（高于 2.0× 的极端不均衡占 33%）。

消融实验

配置	DAPO-Math 加速	CodeForces 加速
仅 BPO（无 MDA）	1.21×	1.17×
仅 MDA（无 BPO）	1.18×	1.15×
BPO + MDA（完整规划器）	1.45×	1.38×
仅 CPU 路径	1.31×	1.28×
仅 GPU Direct	1.28×	1.24×
完整 ForeMoE（双路径）	1.45×	1.38×

BPO 和 MDA 互补：BPO 降低步骤级基线偏斜，MDA 在此基础上精细调整每个微步；两种迁移路径也互补，CPU 路径服务于 Recompute，GPU Direct 服务于 Policy Update。

图 5 可视化完整系统的迁移时间线：

gantt
    dateFormat s
    title 迁移与计算重叠时间线（Recompute 阶段）
    section 微步 m
    GPU 前向计算          :done,   c1, 0, 3s
    section 微步 m+1 准备
    MDA 异步规划（CPU）   :active, t1, 0, 0.3s
    专家权重迁移（PCIe）   :        t2, 0.3, 2.7s
    section 微步 m+1
    GPU 前向计算          :        c2, 3, 3s

图 5：迁移完全被前一个微步的计算遮蔽，有效额外延迟为零。

规划器延迟分析

操作	执行时间	占微步时间比例
BPO（每步一次）	~50ms	步骤时间 ~1000ms → 5%
MDA（每微步一次）	~5ms	微步时间 ~100ms → 5%
CPU 路径迁移（12.5MB 专家）	~0.4ms	完全遮蔽
GPU Direct 迁移（12.5MB 专家）	~0.02ms	完全遮蔽

边界条件分析

ForeMoE 效果最好的场景：

• 微步批次小（每微步序列数 < 64）
• 任务领域集中（数学、代码，专家特化明显）
• 专家数量多、rank 数量合理（灵活性高）
• 显存有余量用于专家复制
• 高带宽互连（NVLink 4 或 PCIe 4.0+）

ForeMoE 效果较弱的场景：

• 微步批次大（统计平均效应恢复）
• 任务领域多样（路由分布均匀，基线不均衡本就不严重）
• 显存极度紧张（无法复制专家）
• 多节点训练（NVLink 不可用，仅 PCIe 或 InfiniBand）

硬性前提条件：RL 训练框架必须实现路由回放（Router Replay），如 VERL、AceMo 等。若系统不强制路由一致性，路由表 $\mathcal{R}$ 的预见价值消失。

批判性分析：不足与可改进之处

主要不足

W1 — 仅在单一模型规模下测试。 所有实验使用 Qwen3-30B-A3B。DeepSeek-V3 有 256 个专家/层（而非 64 个），671B 参数，专家并行度更高。在更大模型上，微步批次相对专家数的比例、显存容量、迁移代价均有本质差异，1.45× 的加速能否保持完全未知。

W2 — 未报告训练质量指标。 论文只报告了墙钟时间加速，没有提供奖励曲线、下游基准（MATH500、HumanEval、AIME）的精度对比，或收敛效率（每样本）的比较。专家迁移是否会引入浮点误差影响梯度质量，完全没有验证。这是一个严重的遗漏。

W3 — Oracle 差距未解释。 ForeMoE 仅达到 Oracle 性能的 95%（1.45× vs 1.52×）。差距来自 MDA 贪心次优、显存限制复制、迁移延迟残余还是层间依赖？论文没有分析，令读者无法判断如何进一步改进。

W4 — 未报告专家复制的显存开销。 对于显存接近上限的大规模训练，每个复制专家约 12.5 MB 的额外开销不可忽视，但论文完全没有提及复制预算设置和显存代价。

W5 — BPO 热启动假设的稳定性。 BPO 使用上一步布局作为热启动，依赖步骤级负载的稳定性。若 RL 早期阶段模型快速学习导致路由偏好剧烈变化，热启动反而会是坏的起点。论文未分析训练初期的热启动质量。

作者淡化的局限

L1 — 路由回放并非 RL 的普遍特性。 论文将路由回放描述为 RL 后训练的固有属性，但实际上这是特定实现（VERL、AceMo）的设计选择，旨在解决推理-训练并行化不匹配问题。某些 PPO 变体，尤其是使用独立参考模型的实现，不要求路由一致性，因此 ForeMoE 的核心机遇不存在。

L2 — 只测试了集中领域数据集。 DAPO-Math-17k 和 CodeForces 都是”集中领域”任务，论文的动机也明确假设了集中领域会导致专家特化和步骤级负载偏斜。对于通用指令跟随、多语言对齐等多样性更高的 RL 任务，路由分布更均匀，基线不均衡本就较低，ForeMoE 的收益可能大幅缩水——但论文没有给出任何此类实验。

L3 — 没有多节点结果。 64 张 GPU 应为单机箱或单 Pod 配置（使用 NVLink）。真实生产级训练横跨数百节点，节点间无 NVLink，CPU 辅助路径也变为通过 InfiniBand 跨节点传输。在这种场景下迁移延迟能否保持在计算窗口内，是重要的未知数。

改进建议

I1 — 增加多尺度模型实验：增加至少一个 70B+ MoE 模型（如 Mixtral-8x22B 或 DeepSeek-V2 规模），验证加速效果和规划开销的可扩展性。

I2 — 加入训练质量对比：报告奖励曲线和下游任务（MATH500、HumanEval）的精度，证明 ForeMoE 与 baseline 收敛到相同质量的策略。

I3 — 多节点版本：实现基于 InfiniBand RDMA 的跨节点专家迁移路径，并测量是否仍能保持计算遮蔽。

I4 — 多样领域实验：在通用指令数据集上运行 RL 实验，量化路由分布均匀时 ForeMoE 的实际收益（或收益缺失）。

I5 — 显存-加速 Pareto 曲线：展示不同复制预算下的加速比与显存开销的 Pareto 曲线，给工程师实际部署参考。

可复现性

• 论文未提供开源代码仓库
• 关键超参数（MDA 偏差阈值 $\varepsilon$、复制预算 $\delta C$、BPO 热启动更新频率）分散在正文中未集中列出
• 实验使用 H100 + NVLink 4；不同硬件（A100 + NVLink 3，多节点 IB）需要重新测量迁移带宽
• 路由表显存开销：约 200 MB（30B 模型，512 token 回复），长上下文 RL（8K token）会增加 16 倍

总结

ForeMoE 的贡献是清晰而有价值的：它识别出 MoE RL 后训练中一个被忽视的性能瓶颈，给出了问题的形式化刻画、NP 困难性证明，以及一个利用特定结构（路由回放）实现实用解的完整系统设计。

1.45× 的端到端加速是切实可见的——相当于将 4 小时的训练减少将近 1.5 小时。层次化规划器（BPO+MDA 分解）和双路径迁移引擎（CPU 辅助 + GPU Direct）的设计是硬件感知的，工程上也具有参考价值。

主要短板在于评估范围：单一模型规模、两个集中域数据集、无训练质量对比、无多节点测试、无代码开源。随着 MoE RL 后训练（DeepSeek-R1 式、Kimi K2.5 式）成为前沿模型开发的标配，这一方向的重要性只会增加，ForeMoE 为后续研究奠定了扎实基础。

附录：关键符号速查表

符号	含义
$E$	每层专家总数
$R$	专家并行 rank 数（GPU 数）
$K$	Top-K 路由，每个 token 激活专家数
$M$	每步微步数
$B_m$	微步批次大小（token 数）
$c_{e,m}$	微步 $m$ 中专家 $e$ 收到的 token 数（来自路由表 $\mathcal{R}$）
$\bar{c}_e$	步骤级均值：$\frac{1}{M}\sum_m c_{e,m}$
$\delta_{e,m}$	微步级偏差：$c_{e,m} - \bar{c}_e$
$L_r^{(m)}$	微步 $m$ 中 rank $r$ 的实际负载
$\rho^{(m)}$	微步 $m$ 的不均衡率：$\max_r L_r^{(m)} / \bar{L}^{(m)}$
$\mathcal{R}$	Rollout 阶段记录的路由表
$P^*$	BPO 输出的基础布局
$A[m]$	MDA 输出的微步 $m$ 调整方案
$x_{e,r}$	专家 $e$ 主副本是否在 rank $r$
$y_{e,r}$	专家 $e$ 是否在 rank $r$ 有复制副本

延伸阅读

• DeepSeek-R1（arXiv 2501.12948）：大规模 MoE RL 后训练的代表性工作，ForeMoE 的最直接目标应用场景
• GRPO（Shao et al., 2024）：ForeMoE 实验所用 RL 算法，Group Relative Policy Optimization
• DAPO（arXiv 2503.14476）：开源大规模 GRPO 系统，构建在 ForeMoE 所针对的 RL 训练流水线之上
• Switch Transformer（arXiv 2101.03961）：大规模 MoE 训练奠基性工作，专家并行与负载均衡背景
• MegaScale（arXiv 2402.15627）：生产级 LLM 大规模训练基础设施，ForeMoE 的系统部署背景
• VAPO（arXiv 2504.05118）：长链式推理 RL，长回复会放大微步路由表规模，让 ForeMoE 的价值更突出
• Mooncake（arXiv 2407.00079）：KV Cache 中心化 LLM 服务架构，展示 KV 缓存拆解如何改变推理流水线
• SGLang（arXiv 2312.07104）：结构化 LLM 执行引擎，RadixAttention 前缀缓存，可作为 Rollout 引擎与 ForeMoE 集成

深度分析：负载不均衡的数学建模

为什么微步级不均衡比步骤级严重得多

设全步 $N \cdot L_{\text{resp}}$ 个 token，专家 $e$ 的长期路由概率为 $p_e$，$K=2$，共 $E=64$ 个专家，$R=8$ 个 rank（每 rank 8 个专家）。

步骤级：全步 token 数 $N_{\text{total}} = 256 \times 512 \times 2 = 262{,}144$，专家 $e$ 的期望 token 数：

$$\mathbb{E}[c_e^{\text{step}}] = N_{\text{total}} \cdot p_e$$

对热门专家 $p_e = 0.04$，$\mathbb{E}[c_e^{\text{step}}] = 10{,}486$，标准差约 $\sqrt{N_{\text{total}} \cdot p_e (1-p_e)} \approx 100$，相对波动 CV = 1%——极低。

微步级：每微步 token 数 $N_m = N_{\text{total}} / M = 32{,}768$，同一专家：

$$\mathbb{E}[c_{e,m}] = 32{,}768 \times 0.04 = 1{,}311$$

标准差 $\approx \sqrt{32{,}768 \times 0.04 \times 0.96} \approx 35$，CV = 2.7%。

但 rank 级负载（8 个专家合计）：每个 rank 上专家的 $P_r = \sum_{e \in r} p_e$，若某 rank 集中了热门专家，$P_r$ 可能达到 $0.04 \times 3 = 0.12$（其他 rank 为 $0.085$，均值 $1/R = 0.125$）。此时最重 rank 的期望负载比平均高约 4%，但由于 $N_m \cdot K \cdot P_r$ 的泊松分布标准差相对均值为 $\frac{1}{\sqrt{N_m K P_r}} \approx \frac{1}{\sqrt{394}} \approx 5\%$，加上倾斜，实际 $\max_r L_r^{(m)}$ 远高于均值。

图 6 展示了不同批次大小下最大不均衡率的理论期望值：

xychart-beta
    title "批次大小与最大不均衡率（理论）"
    x-axis ["256", "512", "1024", "2048", "4096"]
    y-axis "E[ρ_max]" 1 --> 2.5
    line [2.1, 1.75, 1.52, 1.35, 1.20]

图 6：随着每微步批次增大，统计平均效应增强，最大不均衡率趋近于 1。ForeMoE 的实验设置每微步约 512–1024 个序列对应的 token（约 32K token），处于不均衡最严重的区间。

GRPO 目标函数与路由回放的关联

GRPO（Group Relative Policy Optimization）的训练目标为：

$$\mathcal{L}_\text{GRPO}(\theta) = \mathbb{E}_{q \sim P(Q),\; \{o_i\} \sim \pi_{\theta_\text{old}}(\cdot|q)} \left[ \frac{1}{G} \sum_{i=1}^G \min\left( r_i(\theta) \hat{A}_i,\; \text{clip}(r_i(\theta), 1-\epsilon, 1+\epsilon) \hat{A}_i \right) \right]$$

其中 $r_i(\theta) = \frac{\pi_\theta(o_i|q)}{\pi_{\theta_\text{old}}(o_i|q)}$ 是重要性权重，$\hat{A}_i$ 是相对优势估计。

计算 $\pi_\theta(o_i|q)$ 需要 Recompute 阶段重新前向传播——且为了保证与 rollout 对应的 $\pi_{\theta_\text{old}}(o_i|q)$ 使用相同的计算图，必须重放相同的路由决策。

ForeMoE 的正确性保证：专家迁移仅改变”哪张 GPU 执行哪个专家”，不改变”哪个 token 被路由到哪个专家”。因此路由回放约束（$e_{t,\ell}$ 不变）严格满足，重要性权重 $r_i(\theta)$ 的计算在数值上等价于无迁移情况，梯度正确。

工程视角：ForeMoE 与现有系统的集成

与 VERL 的集成

VERL（Versatile Reinforcement Learning）是目前支持 MoE RL 后训练的主要开源框架之一，内置路由回放。ForeMoE 可以作为插件集成，关键改动点：

VERL 训练循环（伪代码）：

原始流程：
  Y, R = vllm_rollout(prompts, θ)
  log_probs = megatron_recompute(Y, R, placement=default)
  loss.backward()

ForeMoE 增强流程：
  Y, R = vllm_rollout(prompts, θ)
  P* = BPO(R.step_mean(), P_prev)   ← 新增
  for m in range(M):
    A_m = MDA(P*, R[m])              ← 新增
    transfer_engine.cpu_path(A_m)    ← 新增
    log_probs_m = megatron_recompute_microstep(Y_m, R[m], P*)
  ...
  P_prev = P*                        ← 新增（热启动）

主要工程代价：

1. 路由表存储：约 200MB（30B 模型，512 token 回复）
2. BPO CPU 计算：约 50ms/步，完全异步，不在关键路径
3. MDA CPU 计算：约 5ms/微步，完全被 GPU 计算遮蔽
4. 修改 All-to-All dispatch 内核，支持按 rank 重定向（需约 1000 行 CUDA 代码）

与预训练负载均衡的对比总结

负载均衡方法二维比较（规划精度 × 执行频率）

执  高 |              ForeMoE-Full ●     Oracle ●
行  |
频  中 |   FasterMoE ●    ForeMoE-BPO ●
率  |   LAER-MoE ●
   低 |              EPLB ●
      +--------------------------------------------→ 规划精度
        低            中            高

方法评分（满分 10 分）
┌──────────────┬────────────┬────────────┬──────────────────────────┐
│ 方法          │ 规划精度    │ 执行频率    │ 特点                      │
├──────────────┼────────────┼────────────┼──────────────────────────┤
│ FasterMoE    │ 3/10       │ 2/10       │ 预训练设计，无 RL 路由回放   │
│ LAER-MoE     │ 3.5/10     │ 2.5/10     │ 负载感知路由，无预见性      │
│ EPLB         │ 4.5/10     │ 3/10       │ 专家并行负载均衡，步级粒度  │
│ ForeMoE-BPO  │ 7/10       │ 3/10       │ 高精度基础规划，缺微步调整  │
│ ForeMoE-Full │ 8.5/10     │ 8.5/10     │ 层次化规划，接近理想区域    │
│ Oracle       │ 10/10      │ 10/10      │ 离线最优解（不可实时实现）   │
└──────────────┴────────────┴────────────┴──────────────────────────┘

图 7：各方法在”规划精度”（来自路由预见性的信息质量）和”执行频率”（响应微步级波动的能力）上的对比。ForeMoE 通过路由回放实现了高精度规划，通过层次化分解实现了高频执行，接近理想区域。

如果我来改进这项工作

从研究者的视角，ForeMoE 的思路可以在以下几个方向推进：

方向 1：Layer-aware 层级规划

目前 ForeMoE 对每一层 MoE 层使用相同的布局，但实际上不同 MoE 层的路由模式不同（不同层专注于不同语义特征）。可以为每层独立维护一套专家布局，代价是规划和迁移开销增加 $n_\text{layers}$ 倍，可能需要进一步的层级重叠策略。

方向 2：路由表压缩

当前路由表 $\mathcal{R}$ 的存储开销为 $O(N \cdot n_\text{layers} \cdot K)$，对于 8K token 长上下文 RL 约 3.2GB。可以利用步骤级稳定性：只存储微步偏差 $\delta_{e,m}$，其标准差远小于均值，可以用低精度（int8）量化存储，节省 4× 空间。

方向 3：多节点 RDMA 迁移

在多节点配置中，节点间带宽降至约 200 GB/s（InfiniBand HDR 或 RoCE 200G）。12.5 MB 专家权重迁移时间约 0.5ms——在 100ms 计算窗口内仍可重叠，但需要新的 RDMA 感知传输实现和跨节点内存管理。

方向 4：预测性路由学习

ForeMoE 的前提是路由回放，但若 RL 框架未实现该约束，可否通过学习一个轻量级路由预测器，从 rollout 的部分路由信息预测完整的后续路由？这将使 ForeMoE 的思路适用于不支持路由回放的框架，代价是引入预测误差。

附录 B：与相关领域工作的关联

与操作系统进程迁移的类比

ForeMoE 的专家迁移在概念上与操作系统中的进程迁移（Process Migration）高度类似：

OS 进程迁移	ForeMoE 专家迁移
进程（状态）	专家权重
CPU 核	GPU rank
进程调度器	Four-Stage Planner
工作负载预测	路由回放（精确预知）
迁移开销	专家权重传输时间
负载均衡目标	最小化 All-to-All makespan

关键区别：OS 调度器无法精确预知未来工作负载，而 ForeMoE 通过路由回放实现了精确预知——相当于一个”全知”的调度器。

与 Mooncake 的对比

Mooncake（FAST 2025）通过 KV Cache 解耦（Prefill/Decode 分离）解决推理服务的负载均衡。两者的核心区别：

• Mooncake 针对推理服务的 KV Cache 分布
• ForeMoE 针对 RL 训练的 MoE 专家负载
• 两者都利用了工作负载的结构性预见性（Mooncake 利用 Prefix Cache 的复用性，ForeMoE 利用路由回放的确定性）

结语

读完 ForeMoE，我印象最深的是”把约束变成机遇”的系统思维方式：路由回放本来是一个让工程师头疼的正确性约束，ForeMoE 却从中挖掘出了精确的负载预测信息，彻底改变了问题性质。

这种思维模式在系统研究中很有价值：下次遇到系统约束，值得多想一步——这个约束是否暴露了什么额外的结构信息，可以被利用？

在工作的局限性上，我认为最关键的缺失是训练质量验证。1.45× 的加速令人印象深刻，但”更快收敛到同样质量的策略”与”更快但收敛到次优策略”对实际价值的影响天差地别。期待后续工作填补这一空白。

附录 C：逐步理解四阶段规划器的完整运行示例

以一个简化例子帮助理解整个流程：$E=4$ 个专家，$R=2$ 个 rank，$K=1$，$M=2$ 个微步，每微步 8 个 token。

Rollout 阶段（路由表记录）：

Rollout 完成后，路由表 R 记录：
  微步 1（token 0-7）：
    token 0 → 专家 0，token 1 → 专家 2，token 2 → 专家 0
    token 3 → 专家 1，token 4 → 专家 0，token 5 → 专家 3
    token 6 → 专家 2，token 7 → 专家 0
    → c_{0,1}=4, c_{1,1}=1, c_{2,1}=2, c_{3,1}=1
  
  微步 2（token 8-15）：
    token 8 → 专家 1，token 9 → 专家 1，token 10 → 专家 3
    token 11 → 专家 3，token 12 → 专家 1，token 13 → 专家 3
    token 14 → 专家 1，token 15 → 专家 3
    → c_{0,2}=0, c_{1,2}=4, c_{2,2}=0, c_{3,2}=4

初始布局（默认，无负载均衡）：

Rank 0：专家 0, 1    Rank 1：专家 2, 3

微步 1 负载：
  Rank 0：c_{0,1} + c_{1,1} = 4 + 1 = 5
  Rank 1：c_{2,1} + c_{3,1} = 2 + 1 = 3
  均值 = 4，不均衡率 ρ = 5/4 = 1.25

微步 2 负载：
  Rank 0：c_{0,2} + c_{1,2} = 0 + 4 = 4
  Rank 1：c_{2,2} + c_{3,2} = 0 + 4 = 4
  均值 = 4，不均衡率 ρ = 1.0 ✓（完美均衡）

BPO 阶段（步骤级均值 $\bar{c}_e$）：

c̄_0 = (4+0)/2 = 2.0，c̄_1 = (1+4)/2 = 2.5
c̄_2 = (2+0)/2 = 1.0，c̄_3 = (1+4)/2 = 2.5

默认布局步骤级负载：
  Rank 0：c̄_0 + c̄_1 = 2.0 + 2.5 = 4.5
  Rank 1：c̄_2 + c̄_3 = 1.0 + 2.5 = 3.5
  均值 = 4.0，不均衡率 = 4.5/4.0 = 1.125

BPO 调整：将专家 1（c̄=2.5，在 Rank 0）迁至 Rank 1
新布局：Rank 0：{0, 2}，Rank 1：{1, 3}
新步骤级负载：
  Rank 0：c̄_0 + c̄_2 = 2.0 + 1.0 = 3.0
  Rank 1：c̄_1 + c̄_3 = 2.5 + 2.5 = 5.0
  更差了！BPO 应选另一个调整：将专家 3 迁至 Rank 0

最终 BPO 布局：Rank 0：{0, 3}，Rank 1：{1, 2}
新步骤级负载：
  Rank 0：2.0 + 2.5 = 4.5 → 还是 1.125
  
实际上均衡的最优布局是 Rank 0：{1, 2}，Rank 1：{0, 3}：
  Rank 0：2.5 + 1.0 = 3.5，Rank 1：2.0 + 2.5 = 4.5 → 1.125
  
这个简单例子说明步骤级负载本身无法完美均衡，
MDA 需要进一步微调以处理微步级波动。

MDA 阶段（微步 1）：

使用 P* = {Rank 0: {0,1}, Rank 1: {2,3}}（BPO 后布局）
微步 1 实际负载：
  Rank 0：c_{0,1} + c_{1,1} = 4 + 1 = 5
  Rank 1：c_{2,1} + c_{3,1} = 2 + 1 = 3
  均值 = 4.0，ρ = 1.25 （Rank 0 过载，Rank 1 欠载）

MDA 识别：专家 0 在 Rank 0，本微步 c_{0,1}=4（最重）
  迁移专家 0 到 Rank 1：
    Rank 0：c_{1,1} = 1
    Rank 1：c_{0,1} + c_{2,1} + c_{3,1} = 4+2+1 = 7 → 过载！
  
  改为复制专家 0（在 Rank 1 也保留一份）：
    Rank 0 处理 token{2,4,7}（3 个 token），Rank 1 处理 token{0}（1 个）
    Rank 0：3 + 1 = 4，Rank 1：1 + 2 + 1 = 4
    不均衡率 ρ = 1.0 ✓ 完美均衡

这个例子展示了 BPO（处理步骤级不均衡）+ MDA（处理微步级偏差）+ 专家复制（解决超载而不引入转移过度）的协同作用。

附录 D：与 DeepSeek-R1 训练场景的对应

ForeMoE 论文发表时，DeepSeek-R1 的 RL 训练细节已经公开（arXiv 2501.12948）。将 ForeMoE 与 DeepSeek-R1 训练对应来看：

DeepSeek-R1 训练特征	ForeMoE 的对应价值
MoE 架构（256 专家/层）	EP 度高，负载均衡灵活性大
GRPO 算法，长链式推理（8K+ token）	路由表大，微步多，ForeMoE 价值最大化
集中于数学/代码推理	集中域，专家特化明显，步骤级负载倾斜
大规模训练（数千 GPU）	现有结果（64 GPU）需扩展至多节点
路由回放（推测启用）	ForeMoE 的核心前提满足

如果 ForeMoE 在 DeepSeek-R1 规模的训练上有效（256 专家，8K token 推理），按论文的加速比（1.38–1.45×），相当于将本来需要 2 个月的 RL 训练压缩到约 6 周——从研究效率的角度，这是非常可观的节省。当然，这需要多节点扩展和更大模型的验证。