arXiv CSopen access

Harnessing Routing Foresight for Micro-step-level MoE load balancing in RL Post-training

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† ‡ Shanghai Jiao Tong University

§ Tencent One MoE Layer

Input

Mixture-of-Experts (MoE) and reinforcement learning (RL) post-training now dominate large language model (LLM) development, yet expert load imbalance remains a critical challenge. Existing load-balancing systems target pre-training by relying on historical step-level statistics. However, these methods fail under the unique workload dynamics of RL posttraining: the step-level load is stable, but the tiny batch sizes processed during micro-steps cause severe, high-frequency load fluctuations. We introduce ForeMoE, a micro-step-level load balancing system for MoE RL post-training. Instead of relying on historical statistics, ForeMoE exploits the multi-stage RL pipeline (rollout, recompute, policy update) by using foreseeable routing information from the rollout stage to proactively guide load balancing in the remaining stages. To support frequent per-micro-step reconfiguration, ForeMoE employs a hierarchical planner that decomposes the NP-hard load balancing problem into tractable sub-components, alongside a transfer engine that leverages complementary hardware paths (CPU-assisted and GPU-direct) for overlapped expert transfer. Evaluations on 64 GPUs demonstrate that ForeMoE achieves up to a 1.45× speedup over state-of-the-art RL post-training systems.

1

Router

4 Experts, EP = 2

All-to-All Dispatch

arXiv:2606.11867v1 [cs.DC] 10 Jun 2026

Abstract

E0 E1 E2

E3

All-to-All Combine

† Peking University

Output

Tokens:

Figure 1. MoE layer with expert parallelism (EP). Load imbalance arises because different tokens are routed to different experts.

popular experts, leaving others underutilized. This significantly degrades overall training performance. Prior efforts to alleviate expert load imbalance largely center on the pre-training phase [2, 13, 30, 39, 56, 67, 68, 70]. As Figure 2 shows, the primary mechanisms use expert relocation and replication to reconfigure expert placement, subsequently reassigning tokens to balance the load. However, effective reconfiguration requires precise knowledge of load patterns. Since routing behavior cannot be known in advance during pre-training (Figure 3(a)), existing systems typically estimate the current load based on historical statistics from previous steps [13, 30, 39, 56, 67, 68]. Motivation: micro-step-level load balancing in RL. While these historical prediction methods work well for pre-training, we find they fall short in RL post-training. As shown in Figure 3(b), RL post-training operates in three sequential stages, including rollout, recompute, and policy update. As detailed in §3, RL post-training targets concentrated task domains (e.g., mathematics or coding) with model experts that have already specialized during pre-training. Consequently, the expert routing pattern converges and remains highly stable across steps. However, a stark instability emerges at the micro-step level: The recompute and policy update stages process data in numerous micro-steps (Figure 3(b)), each containing few samples. This relatively small batch size per micro-step destroys the statistical averaging effect seen at the step level, causing severe, high-frequency fluctuations in expert load. Consequently, load balancing in RL post-training is inherently a micro-step-level problem. Pre-training systems relying on step-level historical statistics simply cannot keep pace with these fine-grained variations. Furthermore, the latency budget available for planning and the possibility to amortize expert transfer costs largely vanish when reconfiguration must be performed at micro-step granularity.

Introduction

The Mixture-of-Experts (MoE) [24] architecture has emerged as the dominant paradigm for scaling large language models (LLMs), powering recent models such as the DeepSeek [5, 6], Qwen [44, 64], and Kimi [57, 58] series. Concurrently, reinforcement learning (RL) post-training, using algorithms like PPO [47], GRPO [49], and DAPO [66], has proven essential for aligning LLMs with human preferences and unlocking advanced reasoning capabilities. As RL post-training is increasingly applied to massive MoE models (e.g., DeepSeekR1 [5], Kimi K2.5 [57]), optimizing the efficiency of these training systems has become a critical challenge. However, as illustrated in Figure 1, MoE training inherently suffers from load imbalance. The routing mechanism often assigns a disproportionate number of tokens to a few *Equal contribution. Contact: Yuming Zhou ([email protected]), Haoyang Li ([email protected]), Fangcheng Fu ([email protected]) and Bin Cui ([email protected]). 1

Challenges. Therefore, achieving effective micro-step-level load balancing in RL post-training requires addressing three fundamental challenges: (C1) What signal should guide load balancing, if historical step-level statistics fail? (C2) How can we design an algorithm capable of solving the complex loadbalancing problem at high, per-micro-step frequencies? (C3) How can we architect a system that supports the rapid expert transfers necessary for such frequent reconfigurations? Our Solutions. Fortunately, the characteristics of the RL post-training provide opportunities to address each of these challenges. To address (C1), we recognize that the routing decisions made during the rollout stage are directly reused in the subsequent recompute and policy update stages (Figure 3(b)) [32, 73]. This structural trait allows us to use rollout routing information as a precise, foreseeable signal for proactive load balancing. To address (C2), we leverage the observation that expert load is stable across steps. We decompose the original NP-hard load balancing problem into two parts: computing a stable base placement at the step level, and performing complementary adjustments at the micro-step level. This drastically reduces algorithmic overhead without sacrificing solving quality. Finally, to address (C3), we observe that different RL stages naturally favor different expert transfer paths. CPU-assisted transfers over PCIe are optimal for the recompute stage, while GPU-direct transfers are more practical for the policy update stage (detailed in §6.1). Exploiting these complementary paths enables rapid expert transfers required for micro-step reconfigurations. To this end, we introduce ForeMoE, a novel system that tackles the underexplored challenge of load imbalance in MoE RL post-training by enabling micro-step-level reconfiguration. ForeMoE harnesses foreseeable routing information from the rollout stage to orchestrate precise load balancing. At its core, ForeMoE features a Four-Stage Planner that breaks the NP-hard load balancing problem into tractable sub-components, and an Expert Transfer Engine that achieves fully overlapped reconfiguration by exploiting complementary transfer paths for different RL stages. In summary, our paper makes the following contributions: • Our paper presents the first study of load imbalance characteristics in MoE RL post-training. Unlike pre-training, RL post-training exhibits highly stable step-level load distributions but severe micro-step-level fluctuations, highlighting the need for micro-step-level load balancing. • We propose ForeMoE, which leverages foreseeable routing information for proactive micro-step load balancing. ForeMoE introduces algorithmic innovations to decompose the NP-hard load balancing problem hierarchically, alongside systemic innovations that utilize complementary hardware paths to achieve rapid expert transfer. • Evaluations on MoE models on 64 GPUs demonstrate that ForeMoE achieves up to a 1.45× end-to-end speedup compared to state-of-the-art RL post-training systems.

Machine 0 GPU 0 GPU 1

E0 150 E1 50

E2 400

Rel: Relocated Expert, Rep: Replicated Expert Machine 1 Machine 0 GPU 0 GPU 1 GPU 0 GPU 1 E4 E6 E0 50 150 150 E2 E7 275 E1 E5 50 50 300 E2’ E4 125 50

E3 150

200

550 350 200 (a) Original Placement Max load: 550, Imbalance: 1.69

Rep

Rel

Machine 1 GPU 2 GPU 3

E3 150

E6 150 E7 50 E5’ 125

Rel

Rep

E5 175

325 325 325 325 (b) After Relocation + Replication Max load: 325, Imbalance: 1.00

Figure 2. Expert relocation and replication mechanisms. (a) Original placement with severe load imbalance. (b) After jointly applying expert relocation (e.g., E3 and E4) and expert replication (e.g., E2’ and E5’), perfect load balance is achieved.

2

Preliminaries

In this section, we provide background on the MoE architecture, examine its load balancing issues, and present the RL post-training pipeline for MoE models. 2.1

MoE Architecture

The MoE [3, 24, 35] architecture replaces the standard feedforward network (FFN) in each Transformer layer with a set of 𝐸 parallel expert networks and a lightweight router module. As shown in Figure 1, for each input token, the router produces a probability distribution over all experts and selects the top-𝐾 experts to process that token. The final output is the weighted sum of the selected experts’ outputs, where the weights are the corresponding router probabilities. As shown in Figure 1, Expert Parallelism (EP) [24] is used to partition experts across devices. Processing a MoE layer involves two phases of All-to-All collective communication: a dispatch phase that routes each token from its source rank to the rank hosting its assigned expert, and a combine phase that returns the expert outputs back to the source ranks. 2.2

MoE Load Balancing

As depicted in Figure 1, expert routing is often highly skewed, with a small subset of experts receiving a disproportionate number of tokens, leading to severe load imbalance across ranks. To address this issue, many prior works [2, 13, 30, 39, 56, 67, 68] have studied load balancing for MoE pre-training. As shown in Figure 2, their approaches can be broadly categorized into two mechanisms: expert relocation, which adjusts the expert-to-rank mapping, and expert replication, which places additional copies of heavily loaded experts. Both expert relocation and expert replication rely on accurate estimates of expert loads. However, as illustrated in 2

model

Training fwd + bwd

M

No Routing Info Beforehand (a) Pre-training

prompts Rollout

micro-step 1 log probs

+ responses Recompute

M

M

micro-step 1 Policy Update fwd + bwd

fwd only

loss

M accum grad

Routing Info

(b) RL Post-training

Figure 3. Comparison of (a) a pre-training step and (b) an RL post-training step for MoE models. (a) During MoE pre-training, routing

Token Share (%)

information is unavailable before execution, requiring existing approaches to predict routing behavior from historical statistics. (b) An RL post-training step consists of rollout, recompute, and policy update stages. The rollout stage generates responses and records routing information, which is subsequently reused by the recompute and policy update stages, making the routing behavior foreseeable.

Figure 3(a), routing information is unavailable before a pretraining step begins executing. Consequently, existing approaches predict the current step’s expert loads from historical statistics observed in previous steps and use the resulting estimates to guide expert relocation and replication. While expert load balancing has been well-studied for pre-training, it remains largely underexplored in RL posttraining. RL workloads exhibit distinct load dynamics that require re-balancing at a much finer timescale, making pretraining methods fall short. We will elaborate on this in §3.

Step Micro-step (a) DAPO-Math-17k

Micro-step (b) CodeForces

Figure 4. Expert load characteristics during RL post-training on Qwen3-30B-A3B. Step-level expert load remains stable but skewed, while micro-step-level expert load exhibits substantial fluctuations.

2.3 RL Post-Training for MoE Models As shown in Figure 3(b), RL post-training typically follows three sequential stages: Rollout. The model (i.e., the current policy) generates responses for a batch of prompts through auto-regressive decoding. This stage typically uses a high-throughput inference framework like vLLM [23] or SGLang [74]. Recompute. To address the inherent mismatch between rollout and training [17, 29, 73], systems recompute log probabilities via a forward pass on the training framework (e.g., FSDP [72] or Megatron [38, 55]). These log probabilities will be used to correct the importance sampling [28, 65, 71] and become part of the loss in the subsequent policy update stage. Due to memory limits, sequences are split into micro-batches and processed sequentially in micro-steps. Policy Update. The model is then optimized using a policy gradient objective (e.g., GRPO [49]). For memory efficiency, this stage also operates over sequential micro-steps to accumulate gradients before applying the final model update. Throughout these three stages, routing consistency must be preserved (i.e., router replay) [32, 73]. Consequently, the routing information generated during rollout is recorded and reused throughout the following stages. As a result, unlike pre-training, the routing information for recompute and policy update stages is foreseeable prior to their execution.

3

Step

where different experts become specialized in distinct linguistic patterns [22] or knowledge domains [61]. In contrast, RL post-training is typically performed on more concentrated tasks (e.g., mathematics [33, 60] or coding [20, 62]) over a model whose experts have already been specialized during pre-training. Figure 4 illustrates two distinctive characteristics of expert load during RL post-training on the DAPOMath-17k [8] and CodeForces [42] datasets, respectively. Step-level: stable but skewed. Because expert specialization has already been established before RL post-training, the overall routing distribution changes little across training steps. Consequently, when aggregated over a sufficiently large number of samples, expert loads exhibit only limited step-level variation. Meanwhile, the concentrated task domain tends to activate a subset of experts more frequently than others, causing some experts to consistently receive substantially more tokens. As a result, the long-term expert load distribution remains skewed. Micro-step-level: highly variable and skewed. Within a training step, however, each micro-step contains only a small number of samples whose routing behaviors may differ significantly. The reduced sample size weakens the statistical averaging effect observed at the step level, causing individual micro-steps to deviate substantially from the long-term load distribution. Consequently, expert loads can fluctuate dramatically across micro-steps, resulting in load distributions that are both highly variable and skewed.

Observations and Current Limitations

During pre-training, MoE models gradually develop expert specialization by learning from massive and diverse corpora, 3

Current limitations. These observations suggest that load balancing in RL post-training should primarily target the micro-step timescale. However, existing MoE load-balancing techniques are predominantly designed for pre-training workloads [2, 13, 39, 56, 67, 68, 70] and fall short for three reasons. First, they rely on historical step-level statistics and therefore cannot react to rapid load fluctuations at micro-step granularity. Second, re-balancing is typically triggered only after significant load shifts, allowing ample time for planning, whereas micro-step-level adaptation operates under much stricter latency budgets. Third, expert transfers are relatively infrequent in pre-training and their costs can be amortized across many steps, while micro-step-level reconfiguration requires substantially higher transfer efficiency.

4

Four-stage Planner (§8)

Routing Info

Base Expert Placement

Rollout Rollout Worker 0 Routing Collector

Rollout Worker 1 Routing Collector

Per-micro-step Adjustments Expert Relocation

Expert Replication

Reconfig Plan (per microstep)

Recompute & Policy Update Training Worker 0 150 50 E0 E1

Token Assignment

Training Worker 1 50 350 Token E2 E3 Expert

Expert Transfer Engine (§6)

Planning-Execution Overlap Plan Plan Plan Recompute Policy Update

E0 50

E2

E3’

50

200

E1 E3 E0’ 50 150 100

Figure 5. Overview of ForeMoE. The Rollout Collector on each rollout worker collects routing information and feeds it to the Fourstage Planner (§8). For each micro-step, the planner determines the optimal expert placement and token assignment. The Expert Transfer Engine (§6) then reconfigures expert placement as needed.

Challenges and Opportunities

Therefore, achieving load balancing at the micro-step timescale in RL post-training requires addressing three key challenges: • The step-level historical statistics widely used during pretraining are no longer informative. What information should guide load-balancing decisions instead? • From an algorithmic perspective, solving the optimal loadbalancing plan is NP-hard (see §7), yet decisions must be made at micro-step timescales. How can we design the algorithm that achieves near-optimal load-balancing quality while meeting stringent real-time requirements? • From a systems perspective, micro-step-level balancing requires frequent reconfiguration throughout different micro-steps. How can experts be transferred efficiently enough to keep pace with such rapid reconfigurations? To address these challenges, we identify three key opportunities, respectively, which together make micro-step-level load balancing both practical and effective. Opportunity 1: Exploiting foreseeable routing information as a planning signal. As shown in Figure 3(b), by the time recompute or policy update begins, the routing behavior has already been determined during rollout and is fully observable. This exposes exact per-token, per-layer, and per-micro-step routing information, allowing the system to proactively plan load-balancing actions. Opportunity 2: Exploiting the stability of step-level expert load. While expert load exhibits noticeable microstep-level fluctuations, these fluctuations are not arbitrary: they remain centered around the aggregate step-level expert load distribution. This structure enables us to separate shortterm load-balancing decisions from long-term ones. Specifically, a stable expert placement can be derived to match the aggregate expert load distribution over training steps, while micro-step-level balancing focuses only on correcting localized deviations from this baseline. As a result, the full problem can be split into a base expert placement solving and a complementary adjustment per-micro-step.

This decomposition preserves solving quality while substantially reducing the solving overhead. Opportunity 3: Exploiting complementary expert transfer paths for different RL stages. Frequent micro-step load balancing requires equally frequent expert transfers. To make such reconfiguration practical, it is important to exploit efficient hardware paths. We identify two complementary transfer paths: (1) CPU-assisted transfer, where experts are fetched from CPU memory, and (2) GPU-direct transfer, which is widely adopted in prior work [2, 13, 39, 56, 67], where experts are obtained via GPU-to-GPU direct transfer. Importantly, we find that different RL stages impose fundamentally different requirements on expert transfers. The recompute stage is forward-only and therefore only involves transferring expert parameters. In contrast, the policy update stage includes both forward and backward computation, and must additionally handle gradient transfers and gradient accumulation across micro-steps. These differences fundamentally affect how expert transfers should be performed, creating an opportunity to employ the most suitable transfer path for each stage (detailed in §6.1).

5

Overview of Our Solutions

Guided by the above opportunities, as shown in Figure 5, we design ForeMoE, a system that leverages foreseeable routing for micro-step-level MoE load balancing in RL post-training. To exploit routing information, ForeMoE introduces a Routing Collector on each rollout worker. During the rollout stage, the collector records the router’s top-𝐾 expert selections for every token at every MoE layer, which provide the view of routing and serve as the foundation for planning. To exploit the stability of step-level expert load, ForeMoE introduces a Four-stage Planner that takes the routing information as input. The planner decomposes the original NP-hard optimization problem into a sequence of tractable 4

6.1 Microstep k

E0

E1

E2

E0 Microstep k+1

E0

E2 E3’

E1

E3

E2

E3

E1

E3 E0’

(a) CPU-assisted Expert Transfer

E0

E1

E2

E3

CPU-assisted expert transfer. Figure 6(a) illustrates the CPU-assisted path. Each machine keeps a full master copy of its layers’ experts in pinned CPU memory, while only the experts currently in use occupy GPU memory. When the planner produces a per-micro-step plan, the engine reads which experts each rank needs for the upcoming micro-step, identifies those not already resident on the GPU, and prefetches them from CPU over PCIe. GPU-direct expert transfer. Figure 6(b) illustrates the GPU-direct path, which is widely adopted in prior MoE pre-training systems [2, 13, 39, 56, 67]. Unlike prefetching from CPU over PCIe, relocation and replication are realized through GPU-to-GPU direct transfers. Design rationale. In ForeMoE, we employ the CPU-assisted path during the recompute stage and the GPU-direct path during the policy update stage. The rationale behind this design choice is as follows. Both transfer paths support communication–computation overlap by issuing expert transfers ahead of their actual use. However, their key difference lies in where the expert copies reside and which hardware channel carries the transfers. In the CPU-assisted design, the complete set of experts is maintained in CPU memory, allowing any GPU to easily fetch any expert. In contrast, the GPU-direct design can easily access only experts that already reside on GPUs within the same machine. Otherwise, if an expert resides on another machine, inter-machine transfers are triggered, which are substantially more expensive than intra-machine transfers over NVLink, making the overhead difficult to hide (as evaluated in §10.3). Consequently, we restrict the GPU-direct path to be planned only within a machine. Under this restriction, CPU-assisted transfer is theoretically preferable because it provides unrestricted access to the full expert pool, thus exposing a much larger placement space to the planner. The recompute stage therefore adopts the CPU-assisted path, as it only involves expert parameter transfers. The policy update stage, however, must additionally handle gradient transfers. Unlike parameters, gradients must be accumulated across micro-steps. Such frequent updates would require continuous synchronization between CPU and GPU if gradients were maintained on the CPU, introducing substantial overhead. We therefore adopt the GPU-direct path for the policy update stage. Further discussion and a feasibility analysis of using a CPU-assisted variant for the policy update stage are provided in Appendix B.

GPU-to-GPU E0

E2 E3’

E1

E3 E0’

(b) GPU-direct Expert Transfer

Figure 6. Two expert transfer paths. (a) CPU-assisted path: each machine maintains a full copy of expert weights in pinned CPU memory. GPUs prefetch the experts required by the next microstep via PCIe. (b) GPU-direct path: expert weights reside on GPUs. Reconfiguration is performed through GPU-to-GPU transfers.

subproblems and sequentially performs (1) base expert placement, (2) expert relocation, (3) expert replication, and (4) token assignment. While base expert placement is performed only once to match the aggregate expert load distribution across steps, the remaining three subproblems are solved for each micro-step to accommodate the micro-step-level load variations. This hierarchical planning process substantially reduces the solving overhead while preserving solution quality. Notably, it generates separate reconfiguration plans for the recompute and policy update stages, specifying expert placement and token assignment for each micro-step. To exploit the complementary expert transfer paths for different RL stages, ForeMoE provides a unified Expert Transfer Engine that supports both CPU-assisted and GPU-direct transfers. To minimize reconfiguration overhead, ForeMoE employs fine-grained communication–computation overlap to hide transfer costs, enabling efficient reconfiguration. As shown at the bottom of Figure 5, throughout this workflow, the Four-stage Planner executes on CPUs concurrently with GPU training. Furthermore, planning tasks for different layers and micro-steps are independent, enabling extensive parallelism. Consequently, the entire planning process can be fully overlapped with the training pipeline, incurring effectively zero critical-path overhead. In the following, we first present the Expert Transfer Engine (§6) for efficient reconfiguration. We then introduce the theoretical formulation of the load-balancing problem (§7), followed by the Four-stage Planner (§8), which decomposes the original NP-hard problem into tractable subproblems.

6

Complementary Expert Transfer Paths

Expert Transfer Engine 6.2

This section first presents two paths underlying the Expert Transfer Engine: a CPU-assisted path and a GPU-direct path. We then discuss the rationale for using different transfer paths for different RL stages (§6.1). Building on these paths, we describe how they enable micro-step-level reconfiguration (§6.2) and further analyze the overhead (§6.3).

Micro-step-level Reconfiguration

Prior MoE pre-training systems typically perform reconfiguration at the step level [13, 39, 56, 67]. The low frequency makes serial expert transfers practical, as the overhead is small relative to step time. Micro-step-level reconfiguration, by contrast, demands much more frequent transfers, so the 5

Per-layer Transfer Plan 0

Micro-step 0 Forward

Plan 1

Micro-step 0 Backward

Micro-step 1 Forward

No Transfer Needed Plan 0

Micro-step 2 Forward

Stream 1

MoE Compute

Stream 2 Transfer layer l

Plan 2

Restore Plan 0 Micro-step 0 Forward

Per-layer Transfer

Micro-step 0 Backward

Micro-step 1 Forward

Micro-step 1 Backward

Stream 1

No Transfer Needed

Copy-Out

MoE Compute

Stream 2 Transfer layer l

Plan 1

backward layer l+1 MoE Attention Combine All-to-All

layer l-1 MoE Combine

Copy-Out

layer l MoE Dispatch

MoE Compute

Copy-In forward

layer l Attention

All-to-All

MoE Combine

MoE Dispatch

MoE Compute

MoE Combine

Copy-In

(b) Three-stage Transfer Procedure

(a) Micro-step-level Reconfiguration

Figure 7. (a) Micro-step-level reconfiguration for the policy update stage, which includes both forward and backward passes. The forward and backward passes of the same micro-step use the same reconfiguration plan. (b) During both the forward and backward passes, per-layer expert transfer is overlapped with the execution on the main stream through a three-stage procedure.

overhead can no longer be easily amortized. Hiding this overhead thus becomes critical to efficient reconfiguration. Per-layer expert transfer. To hide this overhead, we design a per-layer transfer procedure. The key idea is to dedicate a separate CUDA stream to transfer the experts required by the next layer. By initiating the transfer ahead of time, communication on the auxiliary stream can overlap with the computation of the current layer on the main stream. For the recompute stage, which consists only of forward pass and involves parameter transfers alone, the CPU-assisted path enables experts to be prefetched from CPU memory directly. As a result, expert transfers can be easily overlapped. The policy update stage is more challenging. As illustrated in Figure 7(a), it consists of both forward and backward passes, requiring distinct overlap mechanisms for the two passes. Furthermore, expert transfer in this stage involves both parameters and gradients. To improve bandwidth utilization, we should aggregate them into a single transfer. Based on these considerations, we design the three-stage procedure shown in Figure 7(b):

policy updates) for all unexecuted micro-steps. In the recompute stage, a plan is consumed and discarded immediately after the micro-batch’s forward pass. The policy update stage, however, requires the same plan to serve both the forward and backward passes of a micro-step. As shown in Figure 7(a), under pipeline schedules such as 1F1B [36, 37], the backward pass may execute long after the forward pass, when the GPU placement has already been reconfigured for other microsteps. Thus, the plan must be retained until the backward pass completes, allowing it to restore the original forwardtime placement by replaying the old plan. 6.3

Overhead Analysis

Though the reconfiguration happens per-micro-step, both CPU-assisted and GPU-direct transfer (within a machine) incur minimal overhead in terms of memory and latency. (1) Regarding memory, the CPU-assisted path introduces no additional GPU memory consumption, while the GPU-direct path requires only minimal GPU memory for temporary send/receive buffers. These buffers are released immediately after communication and therefore do not contribute to peak memory usage. (2) Regarding latency, empirical evaluations in §10.5 demonstrate that the overhead can be fully overlapped for both paths. Appendix A further derives the conditions for overlapping per-micro-step, per-layer expert transfers, showing that the required per-rank sequence length is only a few thousand tokens for both paths, well below that of typical RL post-training workloads.

• Copy-out. Pack the parameters and gradients that need to leave the current GPU into a contiguous send buffer. • All-to-All swap. Exchange packed buffers. • Copy-in. Unpack the parameters and gradients from the receive buffer into the target memory on the GPU. For the backward pass in particular, for each replicated expert, we designate one replica as the main expert, and accumulate the partial gradients from all replicas into its gradient so that the optimizer applies a single update.

7 For both the forward and backward passes, the three stages of layer 𝑙’s transfer are scheduled against the surrounding training operations as shown in Figure 7(b): Copy-out overlaps with the MoE combine, All-to-All swap overlaps with the Attention, and Copy-in overlaps with the MoE dispatch. Reconfiguration plan management. While the planner continues to generate reconfiguration plans, the Expert Transfer Engine maintains the plans (for both recomputation and

Theoretical Modeling

This section presents the theoretical formulation of our microstep-level load balancing problem. We first develop a time model to estimate the execution time of a single MoE layer, capturing both computation and communication costs (§7.1). Based on this model, we show that, given the routing information of a micro-step, jointly optimizing expert placement and token assignment can be formulated as a Mixed-Integer Linear Program (MILP) (§7.2). 6

Table 1. Notation used in our problem formulation. Symbol 𝐸, 𝑃, 𝑀 𝑁𝑏 , 𝑁𝑟 , 𝑁𝑠 𝑤𝑠,𝑒 𝑥𝑒,𝑗 𝑟𝑠,𝑒,𝑗 𝐿𝑟 𝐶𝑖,𝑗 slots(𝑟 | 𝑚) ranks(𝑚) machine(𝑟 )

stage (one forward), 𝑛 1 = 1 and 𝑛 2 = 2 (one MoE dispatch and one MoE combine). For the policy update stage (forward plus backward), 𝑛 1 = 3 and 𝑛 2 = 4, since the backward pass triggers an additional gradient computation and additional MoE dispatch and MoE combine of gradients. This time model isolates the two metrics the placement problem must minimize: the maximum rank load 𝐿max and the maximum inter-machine traffic 𝐶 max . Both depend on how experts are placed across slots and how tokens are assigned to those slots.

Description Number of experts, EP ranks, and machines Base, redundant and total slots per rank Token volume from source rank 𝑠 to expert 𝑒 Binary, 1 if expert 𝑒 is assigned to slot 𝑗 Fraction of 𝑤𝑠,𝑒 routed to slot 𝑗 Total token load processed by rank 𝑟 Traffic volume from machine 𝑖 to machine 𝑗 Slots owned by rank 𝑟 or machine 𝑚 Ranks owned by machine 𝑚 Machine owning rank 𝑟

7.2

We can formulate the micro-step-level expert placement and token assignment problem for minimizing the per-layer 𝑇MoE as a Mixed-Integer Linear Program (MILP). The inputs consist of the per-source, per-expert token volumes 𝑤𝑠,𝑒 for 𝑠 ∈ [𝑃] and 𝑒 ∈ [𝐸], together with the constants 𝑃, 𝑀, 𝑁𝑏 , 𝑁𝑟 , and 𝑁𝑠 . The decision variables are: • 𝑥𝑒,𝑗 ∈ {0, 1}: placement variable. 𝑥𝑒,𝑗 = 1 if and only if expert 𝑒 is placed in slot 𝑗. • 𝑟𝑠,𝑒,𝑗 ∈ [0, 1]: assignment variable. 𝑟𝑠,𝑒,𝑗 denotes the fraction of token volume 𝑤𝑠,𝑒 that rank 𝑠 routes to slot 𝑗.

Terminology. Table 1 summarizes the notation used throughout the following sections. We elaborate on two key terms that appear repeatedly in our formulation: • Rank. An EP rank is a single device (i.e., GPU) within the EP group. There are 𝑃 ranks in total, distributed evenly across 𝑀 machines, with 𝑃/𝑀 ranks per machine. • Slot. Each rank is configured with 𝑁𝑠 slots, where every slot can hold the parameters as well as the gradients of one expert. Of the 𝑁𝑠 slots on a rank, 𝑁𝑏 = 𝐸/𝑃 are base slots (enough to cover every expert in the model exactly once if there were no replication), and the remaining 𝑁𝑟 are redundant slots available for replicated experts (i.e., 𝑁𝑟 = 𝑁𝑠 − 𝑁𝑏 ). Across the whole EP group there are 𝑃 · 𝑁𝑠 slots in total, indexed by 𝑗 ∈ {1, 2, . . . , 𝑃 · 𝑁𝑠 }. 7.1

Derived metrics. The rank load 𝐿𝑟 aggregates all tokens that are assigned to any slot owned by rank 𝑟 : 𝐿𝑟 =

∑︁

𝐶𝑖,𝑗 =

𝑤𝑠,𝑒 · 𝑟𝑠,𝑒,𝑗 .

(4)

∑︁

∑︁

𝑤𝑠,𝑒 ·𝑟𝑠,𝑒,𝑗 ′ , 𝑖 ≠ 𝑗 . (5)

𝑠 ∈ ranks (machine𝑖 ) 𝑒 ∈ [𝐸 ] 𝑗 ′ ∈ slots (machine 𝑗 )

Constraints. A valid expert placement and token assignment must satisfy the following four constraints, corresponding to expert slot capacity, expert coverage, token conservation, and token assignment feasibility, respectively (see Appendix C for a detailed explanation of each constraint):

(1)

Í𝐸

where 𝐿max = max𝑟 ∈ [𝑃 ] 𝐿𝑟 is the maximum token load, 𝐾1 is the coefficient, and 𝐵 1 is a fixed overhead per layer. Communication time. Communication is dominated by the most heavily loaded inter-machine directional link. In contrast, intra-machine traffic traverses the much faster NVLink fabric and is therefore not a bottleneck. Accordingly, the communication time can be modeled as:

Í𝑃 ·𝑁𝑠

𝑒=1 𝑥𝑒,𝑗 = 1,

𝑗=1

Í𝑃 ·𝑁𝑠 𝑗=1

∀𝑗 ∈ [𝑃 · 𝑁𝑠 ],

(6)

𝑥𝑒,𝑗 ≥ 1, ∀𝑒 ∈ [𝐸],

(7)

𝑟𝑠,𝑒,𝑗 = 1, ∀𝑠 ∈ [𝑃], 𝑒 ∈ [𝐸],

(8)

𝑟𝑠,𝑒,𝑗 ≤ 𝑥𝑒,𝑗 , ∀𝑠 ∈ [𝑃], 𝑒 ∈ [𝐸], 𝑗 ∈ [𝑃 · 𝑁𝑠 ].

(9)

Objective. The objective is to minimize 𝑇MoE as defined in Eq. (3). This objective is inherently nonlinear because 𝐿max and 𝐶 max are formulated as max functions. To address this, we apply the standard epigraph trick [45], replacing each maximum with an auxiliary variable. This yields a proper Mixed-Integer Linear Program (MILP). Complexity. This MILP has 𝐸 · 𝑃 · 𝑁𝑠 binary variables 𝑥𝑒,𝑗 and 𝐸 · 𝑃 2 · 𝑁𝑠 continuous variables 𝑟𝑠,𝑒,𝑗 , coupled by the constraint 𝑟𝑠,𝑒,𝑗 ≤ 𝑥𝑒,𝑗 . It is NP-hard in general, and must be solved for every (micro-step, layer) pair. Solving it within the planner’s per-micro-step time budget is therefore infeasible.

(2)

where 𝐶 max = max𝑖≠𝑗 𝐶𝑖,𝑗 is the maximum inter-machine traffic, 𝐾2 is the coefficient, and 𝐵 2 is a fixed latency. Total layer time. Summing the two components, 𝑇MoE = 𝑛 1 · (𝐾1 · 𝐿max + 𝐵 1 ) + 𝑛 2 · (𝐾2 · 𝐶 max + 𝐵 2 ),

∑︁

The inter-machine traffic 𝐶𝑖,𝑗 aggregates all tokens that flow from ranks on machine 𝑖 to slots owned by machine 𝑗:

We model the execution time of a single MoE layer as the sum of computation and communication time. Both are bounded by the slowest rank, as All-to-All barriers force every rank to wait for the slowest one. Computation time. The computation time is set by the rank with the largest token load:

𝑇comm = 𝐾2 · 𝐶 max + 𝐵 2,

∑︁ ∑︁

𝑠 ∈ [𝑃 ] 𝑒 ∈ [𝐸 ] 𝑗 ∈ slots (𝑟 )

MoE Layer Time Model

𝑇comp = 𝐾1 · 𝐿max + 𝐵 1,

Joint MILP Formulation

(3)

where 𝑛 1 and 𝑛 2 are the numbers of computation and communication rounds invoked by one layer pass. For the recompute 7

8

Algorithm 1 Stage 1: Base Expert Placement

Four-stage Planner

The complexity of the original MILP problem motivates the design of our Four-stage Planner, which decomposes the per(micro-step, layer) MILP into four sequential stages: base expert placement (Stage 1), expert relocation (Stage 2), expert replication (Stage 3) and token assignment (Stage 4). This four-stage decomposition directly accommodates the temporal patterns identified in §3. Specifically, Stage 1 captures the step-level stability: since the per-step aggregate load remains nearly constant across consecutive training steps, a single base mapping computed from this aggregate is reusable over many steps (§8.1). In contrast, Stages 2–4 absorb the micro-step-level volatility: each micro-step’s empirical load deviates from the aggregate, and these stages adjust both expert placement and token assignment to fit the per-micro-step load variation (§8.2).

8.1

Stage 1: Base Expert Placement

Stage 1 assigns every expert to a single base slot using only Í (𝑖 ) (𝑖 ) the per-step aggregate load 𝑤¯ 𝑠,𝑒 = 𝑖 𝑤𝑠,𝑒 , where 𝑤𝑠,𝑒 denotes the load of expert 𝑒 from source rank 𝑠 in the 𝑖-th micro-step. The output 𝐴base is a placement that fills the 𝑁𝑏 base slots on every rank. The remaining 𝑁𝑟 redundant slots per rank remain empty and are filled during Stage 3. We adopt the hierarchical greedy procedure shown in Algorithm 1. It first determines a machine-level placement that specifies the hosting machine of each expert, and then refines it into a rank-level placement that distributes each machine’s assigned experts across its local ranks. Machine-level placement. Experts are sorted by descendÍ ing aggregate load 𝑤¯ 𝑒 = 𝑠 𝑤¯ 𝑠,𝑒 and placed greedily. For each expert, we select the machine that minimizes a joint cost combining computation load and cross-machine traffic. In Algorithm 1, ML[𝑚] tracks machine 𝑚’s accumulated computation load, and MC[𝑚] tracks its accumulated inbound cross-machine traffic. Placing expert 𝑒 on machine ¯ 𝑚 Í would raise ML[𝑚] by 𝑤𝑒 and raise MC[𝑚] by Δ𝑚,𝑒 = ¯ 𝑠,𝑒 , the volume of tokens arriving from other 𝑠 : machine(𝑠 )≠𝑚 𝑤 machines to reach 𝑒. Weighting the two increments by the time-model coefficients of §7.1 gives the scoring function score(𝑚, 𝑒). We place 𝑒 on machine 𝑚 ∗ = arg min𝑚 score(𝑚, 𝑒) and update ML[𝑚 ∗ ], MC[𝑚 ∗ ] accordingly. Rank-level placement. Within each machine, we then distribute its placed experts across the local ranks using the Longest Processing Time (LPT) [11] heuristic: In Algorithm 1, RL[𝑟 ] tracks rank 𝑟 ’s accumulated computation load. Experts are processed in descending order of load and each is placed on the currently least-loaded local rank. This step balances per-rank computation load within each machine without affecting any cross-machine traffic decided during the machine-level placement. 8

Require: Experts 𝐸, EP size 𝑃, machines 𝑀, Per-step aggregate load matrix 𝑊¯ of size 𝑃 × 𝐸 Ensure: Base expert placement 𝐴base : [𝐸] → [𝑃] Í ¯ 𝑒 ← 𝑠 𝑤¯ 𝑠,𝑒 for all 𝑒 1: 𝑤 ¯𝑒 2: Sort experts in descending order of 𝑤 3: ML[𝑚] ← 0, MC[𝑚] ← 0 for 𝑚 ∈ [𝑀] 4: // Machine-level placement 5: for each expert 𝑒 in sorted order do 6: for 𝑚 = 1 to 𝑀 do Í 7: Δ𝑚,𝑒 ← 𝑠:machine(𝑠 )≠𝑚 𝑤¯ 𝑠,𝑒 8: score(𝑚, 𝑒) ← 𝑛 1 𝐾1 (ML[𝑚] + 𝑤¯ 𝑒 ) + 𝑛 2 𝐾2 (MC[𝑚] + Δ𝑚,𝑒 ) 9: end for 10: 𝑚 ∗ ← arg min𝑚 score(𝑚, 𝑒), place expert 𝑒 in machine 𝑚 ∗ 11: ML[𝑚 ∗ ] += 𝑤¯ 𝑒 ; MC[𝑚 ∗ ] += Δ𝑚∗ ,𝑒 12: end for 13: // Rank-level placement 14: for each machine 𝑚 do 15: RL[𝑟 ] ← 0 for 𝑟 ∈ ranks(𝑚) 16: Sort experts on 𝑚 by descending 𝑤¯ 𝑒 17: for each expert 𝑒 on 𝑚 do 18: 𝑟 ∗ ← arg min𝑟 ∈ranks(𝑚) RL[𝑟 ] 19: 𝐴base [𝑒] ← 𝑟 ∗ ; RL[𝑟 ∗ ] += 𝑤¯ 𝑒 20: end for 21: end for 22: return 𝐴base

8.2

Stages 2–4: Per-micro-step Adjustments

Once the base expert placement 𝐴base is determined, each micro-step only needs to adjust the expert placement and determine the token assignment around this base. Algorithm 2 outlines this procedure. Stage 2: expert relocation. The base placement is no longer (𝑖 ) optimal because each micro-step’s load matrix 𝑤𝑠,𝑒 deviates from the aggregate 𝑤¯ 𝑠,𝑒 . To correct this, we first leverage expert relocation by repeatedly swapping pairs of experts. At each round, the algorithm selects the most heavily loaded rank as the swap source, pairs it against every other rank as a swap target, and evaluates a small top-𝐾× top-𝐾 window of candidate expert pairs per target. The swap that yields the largest reduction in the layer time objective is committed. Restricting the source to the bottleneck rank and bounding the per-target candidate set by 𝐾 keeps each round’s cost at 𝑂 (𝑃 · 𝐾 2 ). This loop terminates when no swap improves the objective or the maximum round count is reached. Stage 3: expert replication. With Stage 2’s placement fixed, 𝑃 · 𝑁𝑟 redundant slots remain. We allocate them one at a time. At each step, for every (expert, rank) candidate, we estimate the objective reduction obtained by adding that replica using a locality-aware token assignment heuristic: tokens prefer same-machine replicas to avoid cross-machine traffic, and any leftover tokens are routed to minimize per-rank load. The candidate yielding the largest reduction is selected. The

Table 2. Models, parallelism configurations, and GPUs used for

Algorithm 2 Stages 2-4: Per-micro-step Adjustments

training (the remainder are used for rollout). DP, TP, and EP denote the data-, tensor-, and expert-parallelism degrees.

Require: Base expert placement 𝐴base , per-micro-step loads {𝑊 (𝑖 ) }𝑖=1..𝑁 , redundant slots per rank 𝑁𝑟 , max relocation rounds 𝑇 Ensure: Per-micro-step expert placements {𝐴 (𝑖 ) } and token assignments {𝑟 (𝑖 ) } 1: for each micro-step 𝑖 = 1..𝑁 in parallel do 2: 𝐴 (𝑖 ) ← 𝐴base Init loads RL[𝑟 ]=𝐿𝑟 and traffic LT[𝑖, 𝑗]=𝐶𝑖,𝑗 from 𝐴 (𝑖 ) ,𝑊 (𝑖 ) 3: 4: // Stage 2: expert relocation via swaps 5: for 𝑟𝑜𝑢𝑛𝑑 = 1 to 𝑇 do 6: ℎ ← arg max𝑟 RL[𝑟 ] ⊲ bottleneck rank 7: For each 𝑟𝑙 ≠ ℎ, evaluate swapping a top-𝐾 heaviest expert on ℎ with a top-𝐾 lightest expert on 𝑟𝑙 8: (𝑒ℎ , 𝑒𝑙 , 𝑟𝑙 )★ ← swap minimizing Δ ⊲ Δ: change in layer-time objective 9: if Δ ≥ 0 then break 10: end if 11: Apply swap; update RL, LT 12: end for 13: // Stage 3: expert replication via redundant slots 14: for 𝑘 = 1 to 𝑃 · 𝑁𝑟 do 15: (𝑒 ★, 𝑟 ★, Δ) ← replica candidate with largest estimated drop (locality-aware assignment) 16: if Δ ≥ 0 then break 17: end if 18: Place a replica of 𝑒 ★ on 𝑟 ★ 19: end for 20: // Stage 4: token assignment via LP 21: Solve the LP with {𝑥𝑒,𝑗 } fixed by Stages 2-3 for {𝑟 (𝑖 ) } 22: end for 23: return {𝐴 (𝑖 ) }, {𝑟 (𝑖 ) }

Model

9

∀𝑠 ∈ [𝑃], 𝑒 ∈ [𝐸], ∀𝑠, 𝑒, 𝑗,

𝐿𝑟 ≤ 𝐿 ,

∀𝑟 ∈ [𝑃],

𝐶𝑖,𝑗 ≤ 𝐶 ∗,

∀𝑖 ≠ 𝑗,

128 128 256

⟨4, 4, 16⟩ ⟨8, 4, 32⟩ ⟨16, 2, 32⟩

16 (Total: 64) 32 (Total: 64) 32 (Total: 64)

Implementation

We implement ForeMoE on top of veRL [52], with approximately 10K lines of Python code. We extend Megatron [38, 55] to support dynamic expert placement for training and use vLLM [23] for rollout. Planner. The Four-stage Planner is implemented as a standalone Ray [34] actor. It leverages CPU cores across the entire cluster to plan different (micro-step, layer) instances in parallel, enabling scalable decision-making. As a result, planning is fully decoupled from GPU training and can be completely overlapped with the training process. Token assignment across replicated experts. Standard EP implementations dispatch each token to exactly one destination rank. With expert replication, a token may have multiple dispatch options. We implement a custom token dispatch kernel in Triton [59] that extends the standard Allto-All dispatch to support further token assignment among replicas. This Triton implementation avoids modifying the underlying NCCL [16, 41] All-to-All and maintains compatibility with existing EP communication infrastructure.

min

∗

# Train GPUs

since experts with a single placement have a deterministic token assignment and can be eliminated, (2) the constraint matrix is constructed in sparse form via vectorized operations, and (3) because (micro-step, layer) instances are independent, LPs are solved in parallel across CPU cores. Algorithm 2 forms the basis for both the recompute stage and the policy update stage. In the policy update stage, however, expert relocation (Stage 2) and expert replication (Stage 3) are subject to certain restrictions: both swapping and replica addition must consider whether the reconfiguration can be performed via the GPU-direct path within a machine. By contrast, the recompute stage faces no such limitations, as it can draw on the entire expert pool for relocation and replication through the CPU-assisted transfer path. We refer interested readers to Appendix D for the adapted algorithm under the restrictions during the policy update stage.

𝑟, 𝐿 ∗ , 𝐶 ∗

0 ≤ 𝑟𝑠,𝑒,𝑗 ≤ 𝑥𝑒,𝑗 ,

⟨𝐷𝑃,𝑇 𝑃, 𝐸𝑃⟩

Qwen3-30B-A3B Qwen3-30B-A3B Qwen3.5-35B-A3B

loop terminates when all redundant slots are filled or no candidate improves the objective. Stage 4: token assignment. Once the expert placement is all fixed, the binary variables 𝑥𝑒,𝑗 become constants and the MILP of §7.2 collapses to a linear program (LP) over the continuous token assignment variables 𝑟𝑠,𝑒,𝑗 . Concretely, given the placement {𝑥𝑒,𝑗 } produced by Stages 2-3, the LP for one (micro-step, layer) instance is 𝑛 1 𝐾1 𝐿 ∗ + 𝑛 2 𝐾2𝐶 ∗ Í s.t. 𝑗 𝑟 𝑠,𝑒,𝑗 = 1,

# Experts

(10)

where 𝐿𝑟 and 𝐶𝑖,𝑗 are the rank load and cross-machine traffic defined in §7.2. We solve Eq. (10) with the HiGHS [18] solver to obtain optimal token assignment. Three implementation optimizations keep the per-micro-step LP cheap: (1) only replicated experts contribute non-trivial decision variables,

10

Evaluation

Our evaluation is organized around four goals. We first quantify the end-to-end speedup that ForeMoE achieves over existing systems across different models, EP sizes, and datasets 9

Oracle

1.0 0× 1.0 8× 1.2 6× 1.3 2×

1.4 1× 1.5 2×

0×

1.1

1.0 0× 1.0 7× 1.2 4× 1.3 4×

1.4 5× 1.5 7×

8×

1.1

ForeMoE

8×

veRL+EPLB 1.0

veRL 1.0 0× 1.0 8× 1.2 0× 1.2 7×

1.0

0×

Policy update

1.0 0× 1.0 6× 1.2 2× 1.3 0×

Time (s)

300 200 100 0

Recompute

(a) Qwen3-30B, EP16 (b) Qwen3-30B, EP32 (c) Qwen3.5-35B, EP32 (d) Qwen3-30B, EP16 (e) Qwen3-30B, EP32 (f) Qwen3.5-35B, EP32 DAPO-Math-17k

DAPO-Math-17k

DAPO-Math-17k

CodeForces

CodeForces

CodeForces

Figure 8. End-to-end per-step latency across six configurations (a)–(f). Each bar stacks the recompute stage over the policy update stage. Numbers above the bars denote the end-to-end speedup over veRL.

Table 3. ForeMoE speedup over baselines in different RL stages

(§10.2). We then decompose this speedup to isolate the contributions of individual planning stages and transfer-path choices (§10.3). Next, we study how effectively ForeMoE reshapes rank load and inter-machine traffic, the two metrics targeted by our time model (§10.4). Finally, we verify that both planning and expert transfer overheads can be overlapped with training, even when scaling up (§10.5). 10.1

(“Rec.” denotes recompute and “Upd.” denotes policy update).

Config. (a) (b) (c) (d) (e) (f)

Experimental Setup

Testbed and measurements. All experiments are conducted on a cluster of 8 machines, each equipped with 8 NVIDIA H20 GPUs connected via NVLink. Following the predominant design in modern RL post-training systems [10, 12, 26, 31, 63, 75], we adopt a disaggregated setup where rollout and training (recompute and policy update) are placed on separate GPUs. In this setting, we observe that the rollout stage can actually be fully overlapped with training, and thus the end-to-end per-step latency reduces to the sum of recompute and policy update latencies. We measure this combined training latency to quantify ForeMoE’s speedup.1 Models and datasets. We evaluate two MoE models, Qwen330B-A3B and Qwen3.5-35B-A3B. For Qwen3-30B-A3B we test EP = 16 and EP = 32, and for Qwen3.5-35B-A3B we test EP = 32, yielding three parallelism configurations in total. Table 2 summarizes the settings of each. We run each configuration on two RL datasets: DAPO-Math-17k [8] and CodeForces [42]. Every training step draws 512 samples with a prompt length of 2K and a response length of 8K. Each rank is provisioned with 𝑁𝑟 = 2 redundant expert slots. Baselines. We compare ForeMoE with three baselines:

veRL

veRL+EPLB

Orcale

Rec.

Upd.

Rec.

Upd.

Rec.

Upd.

1.23× 1.55× 1.35× 1.24× 1.47× 1.35×

1.22× 1.40× 1.20× 1.19× 1.38× 1.22×

1.20× 1.31× 1.28× 1.15× 1.26× 1.25×

1.13× 1.19× 1.11× 1.11× 1.16× 1.13×

0.95× 0.96× 0.96× 0.96× 0.96× 0.97×

0.93× 0.90× 0.91× 0.94× 0.91× 0.95×

• veRL+EPLB. We augment veRL with EPLB [2], a representative step-level load balancer designed for MoE pretraining by DeepSeek-V3 [6]. EPLB relies on expert load statistics collected from the previous step and applies a greedy algorithm to determine expert relocation and replication for the current step. • Oracle. Oracle is a hypothetically balanced construct where every rank carries the same load, yielding a global lower bound on latency that is not physically realizable. 10.2

End-to-End Performance

Figure 8 reports the per-step latency across six configurations, while Table 3 further breaks down the speedups of ForeMoE in each RL stage relative to the baselines. Overall, ForeMoE accelerates per-step latency by 1.20– 1.45× over veRL and by 1.12–1.22× over the step-level load balancing baseline veRL+EPLB. In the recompute stage, ForeMoE reduces latency by 1.23–1.55× relative to veRL and by 1.15–1.31× relative to veRL+EPLB. In the policy update stage, the corresponding reductions are 1.19–1.40× and 1.11–1.19×, respectively. Furthermore, ForeMoE achieves 95–97% of the Oracle speedup in the recompute stage and 90–95% in the policy update stage. These results lead to three observations. First, while step-level load balancing can partially mitigate expert imbalance by addressing stable step-level skewness, a noticeable gap remains between veRL+EPLB and ForeMoE,

• veRL. veRL [52] is the state-of-the-art RL post-training system, with rollout–training disaggregation further enabled. It places experts sequentially across ranks in a fixed layout and performs no runtime load balancing. 1 ForeMoE’s design is orthogonal to the choice of RL architecture: it natively

supports both disaggregating and colocating rollout and training. We use the disaggregated setup here as it is the more widely adopted configuration and provides a clean, training-focused comparison. 10

Time (s)

200 150

1.00× 154.5

1.21× 127.6

1.26× 123.0

Recompute 1.28× 121.1

1.30× 119.1

100

Policy update 1.37× 112.6

Transfer path. Table 4 compares three transfer paths: CPUassisted, intra-machine GPU-direct, and unrestricted GPUdirect. The first two correspond to the paths adopted by ForeMoE during recompute and policy update, respectively. For the recompute stage, the CPU-assisted path achieves the best performance (30.0 s). Keeping a complete expert pool in CPU exposes the full placement space to the planner, enabling more effective load balancing, while PCIe expert transfers are largely hidden via communication–computation overlap. In contrast, intra-machine GPU-direct reaches 34.9 s, and unrestricted GPU-direct further degrades to 36.2 s. For the policy update stage, the CPU-assisted path is infeasible: maintaining gradient copies in host memory incurs prohibitive memory and synchronization overhead (see Appendix B). Between the two GPU-direct variants, the intramachine design again proves superior, achieving 76.8 s. Our evaluation of GPU-direct transfers reveals a clear trade-off. Although inter-machine transfers expand the placement space and yield a more balanced layout, realizing such layouts requires moving experts across the slower cross-machine link. Unlike other paths, this communication cannot be effectively overlapped, so the transfer overhead outweighs the modest load-balance gain. Thus, restricting GPU-direct transfers to intra-machine communication yields better performance, validating the design rationale in §6.1.

1.45× 106.8

50

B 0 L P T

Figure 9. B, L, P, and T denote base expert placement, expert relocation, expert replication, and token assignment, respectively. The end-to-end speedup over veRL is shown on top.

Table 4. Transfer path comparison. We report the per-step latency of each RL stage under different expert transfer paths. Transfer path CPU-assisted GPU-direct (intra-machine only) GPU-direct

Recompute

Policy update

30.0s 34.9s 36.2s

N/A 76.8s 81.8s

highlighting the importance of micro-step-level load balancing in RL post-training. Second, while both stages benefit from micro-step-level load balancing, the gains are more pronounced during the recompute stage. As discussed in §6.1, this is because the CPU-assisted transfer path provides access to the entire CPU-resident expert pool. This expands the placement space significantly, enabling ForeMoE to generate more effective load-balancing plans. Third, the small gap between ForeMoE and Oracle indicates limited remaining optimization headroom. In particular, the recompute stage exhibits less than a 5% gap across all configurations. This suggests that ForeMoE approaches the load-balancing upper bound in MoE RL post-training. 10.3

10.4

Case Study

We next examine how ForeMoE reshapes the two key metrics in its time model (§7.1, Eq. (3)): the compute imbalance ratio 𝐿max /𝐿¯ and the maximum inter-machine link traffic 𝐶 max . Figure 10 reports their distributions over 10 steps. For each step, every micro-step contributes one sample, and the resulting distribution is summarized as a box plot, with a line connecting the per-step medians. For veRL, two observations emerge. First, both metrics remain highly stable across steps: neither the box distributions nor their medians shift noticeably over time. This validates the step-level stability assumption underlying ForeMoE’s base placement (§8.1). Second, substantial variation exists at the micro-step level within each step. Under natural placement, the imbalance ratio spans from below 2.5 to nearly 5.8, despite a median of only 2.9. This pronounced volatility underscores the necessity of per-micro-step reconfiguration. ForeMoE substantially reduces both metrics. During the recompute stage, where the planner can leverage the full expert pool, the median imbalance ratio drops to approximately 1.02, while the median peak link traffic decreases from 40K to 18K tokens. During the policy update stage, despite being restricted to intra-machine transfers, the planner still reduces the median imbalance ratio to about 1.06 and the median peak link traffic to approximately 36K tokens.

Ablation Study

We dissect ForeMoE along two axes, the planning algorithm of the Four-stage Planner (§8) and the transfer path of the Expert Transfer Engine (§6), using Qwen3-30B-A3B with EP = 32 on the DAPO-Math-17k dataset as a representative. Planning decomposition. Figure 9 progressively enables the planner stages on top of veRL. The base mapping provides a large improvement (1.21×) by removing the coarse step-level load skew. Relocation and replication then progressively refine per-rank balance, and the final token assignment among replicated experts further reduces the remaining imbalance. Combined, these stages deliver a 1.45× end-to-end speedup. The results show that each stage contributes a nonoverlapping latency reduction, validating the effectiveness of the four-stage decomposition. 11

4 2 1

2

3

4

5

6

7

Training step

8

9 10

60 40 20 1

2

3

4

5

6

7

Training step

8

ForeMoE ForeMoE-opt

300

Time (s)

6

Max cross-machine link (K tokens)

Imbalance ratio

veRL ForeMoE (recompute) ForeMoE (policy update) (a) Compute imbalance (b) Max inter-machine traffic

9 10

+2.2%

200 +2.7% +3.1%

100 0

(a)

(b)

+2.0%

+3.3%

(c)

Recompute Policy update

(d)

Configuration

(e)

+1.4%

(f)

Figure 10. Per-step distribution of (a) the compute imbalance ratio and (b) the maximum inter-machine link traffic over 10 steps.

Figure 11. ForeMoE versus the idealized ForeMoE-opt. 10.5

Overhead Analysis and modified routing policies [25, 46, 76], which trade model quality for more uniform load. Systems approaches include runtime hot-expert replication [13], periodic relocation based on popularity [67], fine-grained replica adjustment [39], decoupling expert parameters from optimizer states for periteration redistribution [56], fully-sharded expert parallelism with frequent re-layout hidden by prefetching [30], and fastpaced token rescheduling [70]. These works all target pre-training, where load imbalance typically manifests at a coarse granularity. Hence, many methods reconfigure only at step boundaries [13, 39, 56, 67], failing to capture the micro-step-level fluctuations that dominate RL post-training. Moreover, pre-training cannot obtain routing information ahead of time, so many methods rely on historical statistics [13, 30, 39, 56, 67, 68]. Acquiring exact loads online would otherwise compress the time budget for decision and reconfiguration on the critical path [70]. In RL post-training, however, routing is foreseeable. ForeMoE exploits this property to enable micro-step-level optimization. Concurrent with our work, Relibra [21] also recognizes the importance of MoE load balancing in RL post-training. Similar to ForeMoE, it leverages rollout routing information and performs reconfiguration at micro-step granularity. However, Relibra focuses exclusively on the policy update stage and adopts only GPU-direct expert transfer. Moreover, while it supports expert replication, it does not support expert relocation as a micro-step-level reconfiguration mechanism. By contrast, ForeMoE supports both CPU-assisted and GPU-direct transfer paths, optimizes both the recompute and policy update stages, and jointly explores expert relocation and replication. Our ablation study (§10.3) demonstrates that these additional capabilities are all important contributors to ForeMoE’s performance gains, as exemplified by the gains from expert relocation in Figure 9 and the benefits of the CPU-assisted transfer path in Table 4. MoE efficiency optimization. Another line of work focuses on MoE efficiency optimization. Existing efforts either optimize communication, such as Tutel’s topology-aware 2D Hierarchical All-to-All [19] and DeepEP’s specialized All-toAll kernels [1], or overlap communication with computation

In this section, we analyze the runtime overhead introduced by ForeMoE’s planning (§8) and expert transfer (§6). Comparison with ForeMoE-opt. We first compare ForeMoE with ForeMoE-opt, an idealized variant that executes planning and expert transfer offline, excluding all overheads to isolate the pure step latency achievable by our algorithm. This serves as an upper bound on ForeMoE’s performance in the absence of planning and transfer costs. Figure 11 reports the per-step latency under different configurations. The latency gap between ForeMoE and ForeMoE-opt remains within 1.4–3.3%, indicating that ForeMoE successfully overlaps nearly all planning and transfer overhead with training. Scalability. We next examine how the planning and experttransfer overheads scale with cluster size. We fix EP = 16 and increase the GPU count by scaling along the data parallelism. Figure 12(a) compares the planning time with the corresponding stage time, aggregated across all layers and microsteps. As the cluster scales, micro-steps are distributed across a larger number of EP groups, reducing both the per-group workload and the corresponding planning time proportionally. Results show that planning remains a small fraction of training time at every scale, accounting for at most 22% of the recompute stage and 3% of the policy update stage. Such overhead can therefore be fully hidden through overlap. Figure 12(b) reports the per-layer expert transfer time, alongside the attention time it overlaps with. Because each EP group performs its relocation and replication independently, the transfer cost is invariant to the GPU count, holding at roughly 1.72 ms for the CPU-assisted path and 1.07 ms for the GPU-direct path (intra-machine) across all scales. Both remain well below the attention time of about 5.8 ms. This confirms that the transfer overhead stays fully overlapped regardless of the cluster scale.

11

Related Work

MoE load balancing. Mitigating expert load imbalance has been studied from both algorithmic and systems perspectives. Algorithmic methods include auxiliary balancing losses [7, 9, 50], expert capacity with token dropping [19, 24], 12

# GPUs

Planning Rec.

Upd.

16

7.95s

4.20s

32

4.65s

2.21s

64

3.25s

1.24s

(a) Planning Overhead

//dx.doi.org/10.1109/TKDE.2025.3554028. [4] Qiaoling Chen, Zijun Liu, Peng Sun, Shenggui Li, Guoteng Wang, Ziming Liu, Yonggang Wen, Siyuan Feng, and Tianwei Zhang. Respec: Towards optimizing speculative decoding in reinforcement learning systems, 2025. https://arxiv.org/abs/2510.26475. [5] DeepSeek-AI, Daya Guo, Dejian Yang, et al. Deepseek-r1: Incentivizing reasoning capability in llms via reinforcement learning, 2025. https: //arxiv.org/abs/2501.12948. [6] DeepSeek-AI, Aixin Liu, Bei Feng, et al. Deepseek-v3 technical report, 2025. https://arxiv.org/abs/2412.19437. [7] Nan Du, Yanping Huang, Andrew M Dai, Simon Tong, Dmitry Lepikhin, Yuanzhong Xu, Maxim Krikun, Yanqi Zhou, Adams Wei Yu, Orhan Firat, Barret Zoph, Liam Fedus, Maarten P Bosma, Zongwei Zhou, Tao Wang, Emma Wang, Kellie Webster, Marie Pellat, Kevin Robinson, Kathleen Meier-Hellstern, Toju Duke, Lucas Dixon, Kun Zhang, Quoc Le, Yonghui Wu, Zhifeng Chen, and Claire Cui. GLaM: Efficient scaling of language models with mixture-of-experts. In Proceedings of the 39th International Conference on Machine Learning, 2022. https://proceedings.mlr.press/v162/du22c.html. [8] Hugging Face. Dapo-math-17k dataset, 2025. https://huggingface.co/ datasets/BytedTsinghua-SIA/DAPO-Math-17k. [9] William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. Journal of Machine Learning Research, 2022. http://jmlr.org/papers/ v23/21-0998.html. [10] Wei Fu, Jiaxuan Gao, Xujie Shen, Chen Zhu, Zhiyu Mei, Chuyi He, Shusheng Xu, Guo Wei, Jun Mei, Jiashu Wang, Tongkai Yang, Binhang Yuan, and Yi Wu. Areal: A large-scale asynchronous reinforcement learning system for language reasoning, 2025. https://arxiv.org/abs/ 2505.24298. [11] R. L. Graham. Bounds on multiprocessing timing anomalies. SIAM J. Appl. Math., 1969. https://doi.org/10.1137/0117039. [12] Zhenyu Han, Ansheng You, Haibo Wang, Kui Luo, Guang Yang, Wenqi Shi, Menglong Chen, Sicheng Zhang, Zeshun Lan, Chunshi Deng, Huazhong Ji, Wenjie Liu, Yu Huang, Yixiang Zhang, Chenyi Pan, Jing Wang, Xin Huang, Chunsheng Li, and Jianping Wu. Asyncflow: An asynchronous streaming rl framework for efficient llm post-training, 2025. https://arxiv.org/abs/2507.01663. [13] Jiaao He, Jidong Zhai, Tiago Antunes, Haojie Wang, Fuwen Luo, Shangfeng Shi, and Qin Li. Fastermoe: modeling and optimizing training of large-scale dynamic pre-trained models. In Proceedings of the 27th ACM SIGPLAN Symposium on Principles and Practice of Parallel Programming, 2022. https://doi.org/10.1145/3503221.3508418. [14] Jingkai He, Tianjian Li, Erhu Feng, Dong Du, Qian Liu, Tao Liu, Yubin Xia, and Haibo Chen. History rhymes: Accelerating llm reinforcement learning with rhymerl, 2025. https://arxiv.org/abs/2508.18588. [15] Jian Hu, Xibin Wu, Wei Shen, Jason Klein Liu, Zilin Zhu, Weixun Wang, Songlin Jiang, Haoran Wang, Hao Chen, Bin Chen, Weikai Fang, Xianyu, Yu Cao, Haotian Xu, and Yiming Liu. Openrlhf: An easy-to-use, scalable and high-performance rlhf framework, 2025. https://arxiv.org/abs/2405.11143. [16] Zhiyi Hu, Siyuan Shen, Tommaso Bonato, Sylvain Jeaugey, Cedell Alexander, Eric Spada, James Dinan, Jeff Hammond, and Torsten Hoefler. Demystifying nccl: An in-depth analysis of gpu communication protocols and algorithms, 2025. https://arxiv.org/abs/2507.04786. [17] Wei Huang, Yi Ge, Shuai Yang, Yicheng Xiao, Huizi Mao, Yujun Lin, Hanrong Ye, Sifei Liu, Ka Chun Cheung, Hongxu Yin, Yao Lu, Xiaojuan Qi, Song Han, and Yukang Chen. Qerl: Beyond efficiency – quantization-enhanced reinforcement learning for llms, 2025. https: //arxiv.org/abs/2510.11696. [18] Q. Huangfu and J. A. J. Hall. Parallelizing the dual revised simplex method. Mathematical Programming Computation, 2018. https://doi. org/10.1007/s12532-017-0130-5.

(b) Expert Transfer Overhead

Figure 12. (a) Per-step planning time versus stage time. “Rec.” denotes recompute and “Upd.” denotes policy update. (b) Per-layer expert transfer time versus attention time.

through expert-workload chunking [53, 54]. Comet [69] further develops GPU kernel fusion via thread-block specialization. These techniques are orthogonal to ForeMoE, as they optimize MoE execution itself and can be directly combined with our load-balancing mechanisms. RL post-training systems. A growing body of work has explored efficient RL post-training for LLMs. Early systems such as OpenRLHF [15] adopt colocated architectures, while more recent designs embrace asynchronous and disaggregated execution [10, 12, 14, 26, 31, 40, 51, 63, 75] to decouple rollout from training, improving resource utilization. Other works accelerate rollout through speculative decoding [4, 14, 27, 43, 48] and low-precision quantization [17, 31]. While these systems address utilization and efficiency, they do not explicitly tackle MoE load balancing. Thus, their optimization objectives are largely orthogonal to ours, allowing ForeMoE to be seamlessly integrated with these approaches.

12

Conclusion

This paper presents the first study of load imbalance characteristics in MoE RL post-training. We show that, unlike MoE pre-training where imbalance primarily occurs at step granularity, RL post-training exhibits substantial micro-step-level load fluctuations, making fine-grained balancing necessary. Motivated by this observation, we design ForeMoE, a loadbalancing system that leverages foreseeable routing information from rollout as an accurate planning signal. ForeMoE enables micro-step-level load balancing through a four-stage planning process and rapid reconfiguration via complementary expert transfer paths. Experiments on a 64-GPU cluster demonstrate that ForeMoE achieves up to 1.45× end-to-end speedup over state-of-the-art MoE RL post-training systems.

References [1] Deepep: A high-performance communication library, 2025. https: //github.com/deepseek-ai/DeepEP. [2] Expert parallelism load balancer, 2025. https://github.com/deepseekai/EPLB. [3] Weilin Cai, Juyong Jiang, Fan Wang, Jing Tang, Sunghun Kim, and Jiayi Huang. A survey on mixture of experts in large language models. IEEE Transactions on Knowledge and Data Engineering, 2025. http: 13

[32] Wenhan Ma, Hailin Zhang, Liang Zhao, Yifan Song, Yudong Wang, Zhifang Sui, and Fuli Luo. Stabilizing moe reinforcement learning by aligning training and inference routers, 2025. https://arxiv.org/abs/ 2510.11370. [33] Prasanna Mayilvahanan, Ricardo Dominguez-Olmedo, Thaddäus Wiedemer, and Wieland Brendel. Math-beyond: A benchmark for rl to expand beyond the base model, 2025. https://arxiv.org/abs/2510.11653. [34] Philipp Moritz, Robert Nishihara, Stephanie Wang, Alexey Tumanov, Richard Liaw, Eric Liang, Melih Elibol, Zongheng Yang, William Paul, Michael I. Jordan, and Ion Stoica. Ray: a distributed framework for emerging ai applications. In Proceedings of the 13th USENIX Conference on Operating Systems Design and Implementation, 2018. https://dl.acm. org/doi/10.5555/3291168.3291210. [35] Siyuan Mu and Sen Lin. A comprehensive survey of mixture-ofexperts: Algorithms, theory, and applications, 2025. https://arxiv.org/ abs/2503.07137. [36] Deepak Narayanan, Aaron Harlap, Amar Phanishayee, Vivek Seshadri, Nikhil R. Devanur, Gregory R. Ganger, Phillip B. Gibbons, and Matei Zaharia. Pipedream: generalized pipeline parallelism for dnn training. In Proceedings of the 27th ACM Symposium on Operating Systems Principles, 2019. https://doi.org/10.1145/3341301.3359646. [37] Deepak Narayanan, Amar Phanishayee, Kaiyu Shi, Xie Chen, and Matei Zaharia. Memory-efficient pipeline-parallel dnn training. In Proceedings of the 38th International Conference on Machine Learning, 2021. https://proceedings.mlr.press/v139/narayanan21a.html. [38] Deepak Narayanan, Mohammad Shoeybi, Jared Casper, Patrick LeGresley, Mostofa Patwary, Vijay Korthikanti, Dmitri Vainbrand, Prethvi Kashinkunti, Julie Bernauer, Bryan Catanzaro, Amar Phanishayee, and Matei Zaharia. Efficient large-scale language model training on gpu clusters using megatron-lm. In Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis, 2021. https://doi.org/10.1145/3458817.3476209. [39] Xiaonan Nie, Xupeng Miao, Zilong Wang, Zichao Yang, Jilong Xue, Lingxiao Ma, Gang Cao, and Bin Cui. Flexmoe: Scaling large-scale sparse pre-trained model training via dynamic device placement. Proc. ACM Manag. Data, 2023. https://doi.org/10.1145/3588964. [40] Michael Noukhovitch, Shengyi Huang, Sophie Xhonneux, Arian Hosseini, Rishabh Agarwal, and Aaron Courville. Asynchronous rlhf: Faster and more efficient off-policy rl for language models, 2025. https://arxiv.org/abs/2410.18252. [41] NVIDIA. Nvidia collective communication library (nccl) documentation, 2025. https://docs.nvidia.com/deeplearning/nccl/user-guide/ docs/index.html. [42] Guilherme Penedo, Anton Lozhkov, Hynek Kydlíček, Loubna Ben Allal, Edward Beeching, Agustín Piqueres Lajarín, Quentin Gallouédec, Nathan Habib, Lewis Tunstall, and Leandro von Werra. Codeforces, 2025. https://huggingface.co/datasets/open-r1/codeforces. [43] Ruoyu Qin, Weiran He, Weixiao Huang, Yangkun Zhang, Yikai Zhao, Bo Pang, Xinran Xu, Yingdi Shan, Yongwei Wu, and Mingxing Zhang. Seer: Online context learning for fast synchronous llm reinforcement learning, 2025. https://arxiv.org/abs/2511.14617. [44] Qwen, :, An Yang, Baosong Yang, Beichen Zhang, et al. Qwen2.5 technical report, 2025. https://arxiv.org/abs/2412.15115. [45] R. Tyrell Rockafellar. Convex Analysis. Princeton University Press, 1970. http://www.jstor.org/stable/j.ctt14bs1ff. [46] Stephen Roller, Sainbayar Sukhbaatar, arthur szlam, and Jason Weston. Hash layers for large sparse models. In Advances in Neural Information Processing Systems, 2021. https://proceedings.neurips.cc/paper_files/ paper/2021/file/92bf5e6240737e0326ea59846a83e076-Paper.pdf. [47] John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms, 2017. https: //arxiv.org/abs/1707.06347. [48] Zelei Shao, Vikranth Srivatsa, Sanjana Srivastava, Qingyang Wu, Alpay Ariyak, Xiaoxia Wu, Ameen Patel, Jue Wang, Percy Liang, Tri

[19] Changho Hwang, Wei Cui, Yifan Xiong, Ziyue Yang, Ze Liu, Han Hu, Zilong Wang, Rafael Salas, Jithin Jose, Prabhat Ram, HoYuen Chau, Peng Cheng, Fan Yang, Mao Yang, and Yongqiang Xiong. Tutel: Adaptive mixture-of-experts at scale. In Proceedings of Machine Learning and Systems, 2023. https://proceedings.mlsys.org/paper_files/paper/ 2023/file/5616d34cf8ff73942cfd5aa922842556-Paper-mlsys2023.pdf. [20] Xue Jiang, Yihong Dong, Mengyang Liu, Hongyi Deng, Tian Wang, Yongding Tao, Rongyu Cao, Binhua Li, Zhi Jin, Wenpin Jiao, Fei Huang, Yongbin Li, and Ge Li. Coderl+: Improving code generation via reinforcement with execution semantics alignment, 2026. https://arxiv.org/abs/2510.18471. [21] Chao Jin, Xinming Wei, Yinmin Zhong, Chengxu Yang, Bingyang Wu, Ruidong Zhu, Zili Zhang, Yuliang Liu, and Xin Jin. Relibra: Routingreplay-guided load balancing for moe training in reinforcement learning, 2026. https://arxiv.org/abs/2605.08639. [22] Aditi Khandelwal, Marius Mosbach, Verna Dankers, Siva Reddy, and Golnoosh Farnadi. Towards democratizing LLMs: Investigating multilingual mixture-of-experts models. In Women in Machine Learning Workshop @ NeurIPS 2025, 2026. https://openreview.net/forum?id= Bwf4grCk3H. [23] Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying Sheng, Lianmin Zheng, Cody Hao Yu, Joseph Gonzalez, Hao Zhang, and Ion Stoica. Efficient memory management for large language model serving with pagedattention. In Proceedings of the 29th Symposium on Operating Systems Principles, 2023. https://doi.org/10.1145/3600006.3613165. [24] Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu, Dehao Chen, Orhan Firat, Yanping Huang, Maxim Krikun, Noam Shazeer, and Zhifeng Chen. {GS}hard: Scaling giant models with conditional computation and automatic sharding. In International Conference on Learning Representations, 2021. https://openreview.net/forum?id= qrwe7XHTmYb. [25] Mike Lewis, Shruti Bhosale, Tim Dettmers, Naman Goyal, and Luke Zettlemoyer. Base layers: Simplifying training of large, sparse models. In Proceedings of the 38th International Conference on Machine Learning, 2021. https://proceedings.mlr.press/v139/lewis21a.html. [26] Haoyang Li, Sheng Lin, Fangcheng Fu, Yuming Zhou, Xiaodong Ji, Yanfeng Zhao, Lefeng Wang, Jie Jiang, and Bin Cui. Unleashing efficient asynchronous rl post-training via staleness-constrained rollout coordination, 2026. https://arxiv.org/abs/2601.12784. [27] Bingshuai Liu, Ante Wang, Zijun Min, Liang Yao, Haibo Zhang, Yang Liu, Xu Han, Peng Li, Anxiang Zeng, and Jinsong Su. Spec-rl: Accelerating on-policy reinforcement learning with speculative rollouts, 2026. https://arxiv.org/abs/2509.23232. [28] Jiacai Liu, Yingru Li, Yuqian Fu, Jiawei Wang, Qian Liu, and Yu Shen. When speed kills stability: Demystifying rl collapse from the inferencetraining mismatch, 2025. https://yingru.notion.site/When-SpeedKills-Stability-Demystifying-RL-Collapse-from-the-InferenceTraining-Mismatch-271211a558b7808d8b12d403fd15edda. [29] Liyuan Liu, Feng Yao, Dinghuai Zhang, Chengyu Dong, Jingbo Shang, and Jianfeng Gao. Flashrl: 8bit rollouts, full power rl, 2025. https: //fengyao.notion.site/flash-rl. [30] Xinyi Liu, Yujie Wang, Fangcheng Fu, Xuefeng Xiao, Huixia Li, Jiashi Li, and Bin Cui. Laer-moe: Load-adaptive expert re-layout for efficient mixture-of-experts training. In Proceedings of the 31st ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 2, 2026. https://doi.org/10. 1145/3779212.3790180. [31] Han Lu, Zichen Liu, Shaopan Xiong, Yancheng He, Wei Gao, Yanan Wu, Weixun Wang, Jiashun Liu, Yang Li, Haizhou Zhao, Ju Huang, Siran Yang, Xiaoyang Li, Yijia Luo, Zihe Liu, Ling Pan, Junchi Yan, Wei Wang, Wenbo Su, Jiamang Wang, Lin Qu, and Bo Zheng. Part ii: Roll flash – accelerating rlvr and agentic training with asynchrony, 2025. https://arxiv.org/abs/2510.11345.

14

Dao, Ce Zhang, Yiying Zhang, Ben Athiwaratkun, Chenfeng Xu, and Junxiong Wang. Beat the long tail: Distribution-aware speculative decoding for rl training, 2025. https://arxiv.org/abs/2511.13841. [49] Zhihong Shao, Peiyi Wang, Qihao Zhu, Runxin Xu, Junxiao Song, Xiao Bi, Haowei Zhang, Mingchuan Zhang, Y. K. Li, Y. Wu, and Daya Guo. Deepseekmath: Pushing the limits of mathematical reasoning in open language models, 2024. https://arxiv.org/abs/2402.03300. [50] Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gated mixture-of-experts layer, 2017. https: //arxiv.org/abs/1701.06538. [51] Guangming Sheng, Yuxuan Tong, Borui Wan, Wang Zhang, Chaobo Jia, Xibin Wu, Yuqi Wu, Xiang Li, Chi Zhang, Yanghua Peng, Haibin Lin, Xin Liu, and Chuan Wu. Laminar: A scalable asynchronous rl post-training framework, 2025. https://arxiv.org/abs/2510.12633. [52] Guangming Sheng, Chi Zhang, Zilingfeng Ye, Xibin Wu, Wang Zhang, Ru Zhang, Yanghua Peng, Haibin Lin, and Chuan Wu. Hybridflow: A flexible and efficient rlhf framework. In Proceedings of the Twentieth European Conference on Computer Systems, 2025. http://dx.doi.org/10. 1145/3689031.3696075. [53] Shaohuai Shi, Xinglin Pan, Xiaowen Chu, and Bo Li. Pipemoe: Accelerating mixture-of-experts through adaptive pipelining. In IEEE INFOCOM 2023 - IEEE Conference on Computer Communications, 2023. https://doi.org/10.1109/INFOCOM53939.2023.10228874. [54] Shaohuai Shi, Xinglin Pan, Qiang Wang, Chengjian Liu, Xiaozhe Ren, Zhongzhe Hu, Yu Yang, Bo Li, and Xiaowen Chu. Schemoe: An extensible mixture-of-experts distributed training system with tasks scheduling. In Proceedings of the Nineteenth European Conference on Computer Systems, 2024. https://doi.org/10.1145/3627703.3650083. [55] Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catanzaro. Megatron-lm: Training multibillion parameter language models using model parallelism, 2020. https://arxiv.org/abs/1909.08053. [56] Athinagoras Skiadopoulos, Mark Zhao, Swapnil Gandhi, Thomas Norrie, Shrijeet Mukherjee, and Christos Kozyrakis. SYMI: Efficient Mixture-of-Experts training via model and optimizer state decoupling. In 23rd USENIX Symposium on Networked Systems Design and Implementation (NSDI 26), 2026. https://www.usenix.org/conference/nsdi26/ presentation/skiadopoulos. [57] Kimi Team, Tongtong Bai, Yifan Bai, et al. Kimi k2.5: Visual agentic intelligence, 2026. https://arxiv.org/abs/2602.02276. [58] Kimi Team, Yifan Bai, Yiping Bao, et al. Kimi k2: Open agentic intelligence, 2025. https://arxiv.org/abs/2507.20534. [59] Philippe Tillet, H. T. Kung, and David Cox. Triton: an intermediate language and compiler for tiled neural network computations. In Proceedings of the 3rd ACM SIGPLAN International Workshop on Machine Learning and Programming Languages, 2019. https://doi.org/10.1145/ 3315508.3329973. [60] Peng-Yuan Wang, Tian-Shuo Liu, Chenyang Wang, Ziniu Li, Yidi Wang, Shu Yan, Chengxing Jia, Xu-Hui Liu, Xinwei Chen, Jiacheng Xu, and Yang Yu. A survey on large language models for mathematical reasoning. ACM Comput. Surv., 2026. https://dl.acm.org/doi/10.1145/ 3786333. [61] Xi Wang, Soufiane Hayou, and Eric Nalisnick. The myth of expert specialization in moes: Why routing reflects geometry, not necessarily domain expertise, 2026. https://arxiv.org/abs/2604.09780. [62] Yuxiang Wei, Olivier Duchenne, Jade Copet, Quentin Carbonneaux, Lingming Zhang, Daniel Fried, Gabriel Synnaeve, Rishabh Singh, and Sida I. Wang. Swe-rl: Advancing llm reasoning via reinforcement learning on open software evolution, 2025. https://arxiv.org/abs/2502. 18449. [63] Bo Wu, Sid Wang, Yunhao Tang, Jia Ding, Eryk Helenowski, Liang Tan, Tengyu Xu, Tushar Gowda, Zhengxing Chen, Chen Zhu, Xiaocheng Tang, Yundi Qian, Beibei Zhu, and Rui Hou. Llamarl: A distributed

asynchronous reinforcement learning framework for efficient largescale llm training, 2025. https://arxiv.org/abs/2505.24034. [64] An Yang, Anfeng Li, Baosong Yang, et al. Qwen3 technical report, 2025. https://arxiv.org/abs/2505.09388. [65] Feng Yao, Liyuan Liu, Dinghuai Zhang, Chengyu Dong, Jingbo Shang, and Jianfeng Gao. Your efficient rl framework secretly brings you off-policy rl training, 2025. https://fengyao.notion.site/off-policy-rl. [66] Qiying Yu, Zheng Zhang, Ruofei Zhu, Yufeng Yuan, Xiaochen Zuo, Yu Yue, Weinan Dai, Tiantian Fan, Gaohong Liu, Lingjun Liu, Xin Liu, Haibin Lin, Zhiqi Lin, Bole Ma, Guangming Sheng, Yuxuan Tong, Chi Zhang, Mofan Zhang, Wang Zhang, Hang Zhu, Jinhua Zhu, Jiaze Chen, Jiangjie Chen, Chengyi Wang, Hongli Yu, Yuxuan Song, Xiangpeng Wei, Hao Zhou, Jingjing Liu, Wei-Ying Ma, Ya-Qin Zhang, Lin Yan, Mu Qiao, Yonghui Wu, and Mingxuan Wang. Dapo: An open-source llm reinforcement learning system at scale, 2025. https://arxiv.org/ abs/2503.14476. [67] Mingshu Zhai, Jiaao He, Zixuan Ma, Zan Zong, Runqing Zhang, and Jidong Zhai. SmartMoE: Efficiently training Sparsely-Activated models through combining offline and online parallelization. In 2023 USENIX Annual Technical Conference (USENIX ATC 23), 2023. https://www. usenix.org/conference/atc23/presentation/zhai. [68] Junyi Zhang, Chuanhu Ma, Xiong Wang, Yuntao Nie, Yuqing Li, Yuedong Xu, Xiaofei Liao, Bo Li, and Hai Jin. PopFetcher: Towards accelerated Mixture-of-Experts training via popularity based ExpertWise prefetch. In 2025 USENIX Annual Technical Conference (USENIX ATC 25), 2025. https://www.usenix.org/conference/atc25/presentation/ zhang-junyi. [69] Shulai Zhang, Ningxin Zheng, Haibin Lin, Ziheng Jiang, Wenlei Bao, Chengquan Jiang, Qi Hou, Weihao Cui, Size Zheng, Li-Wen Chang, Quan Chen, and Xin Liu. Comet: Finegrained computation-communication overlapping for mixture-ofexperts. In Proceedings of Machine Learning and Systems, 2025. https://proceedings.mlsys.org/paper_files/paper/2025/file/ e27ea0cd50b798ff8942caf9203f0992-Paper-Conference.pdf. [70] Chenqi Zhao, Wenfei Wu, Linhai Song, Yuchen Xu, and Yitao Yuan. Fine-grained moe load balancing with linear programming, 2026. https: //arxiv.org/abs/2511.16947. [71] Xin Zhao, Yongkang Liu, Kuan Xu, Jia Guo, Zihao Wang, Yan Sun, Xinyu Kong, Qianggang Cao, Liang Jiang, Zujie Wen, Zhiqiang Zhang, and Jun Zhou. Small leak can sink a great ship–boost rl training on moe with icepop!, 2025. https://ringtech.notion.site/icepop. [72] Yanli Zhao, Andrew Gu, Rohan Varma, Liang Luo, Chien-Chin Huang, Min Xu, Less Wright, Hamid Shojanazeri, Myle Ott, Sam Shleifer, Alban Desmaison, Can Balioglu, Pritam Damania, Bernard Nguyen, Geeta Chauhan, Yuchen Hao, Ajit Mathews, and Shen Li. Pytorch fsdp: Experiences on scaling fully sharded data parallel. Proc. VLDB Endow., 2023. https://doi.org/10.14778/3611540.3611569. [73] Chujie Zheng, Kai Dang, Bowen Yu, Mingze Li, Huiqiang Jiang, Junrong Lin, Yuqiong Liu, Hao Lin, Chencan Wu, Feng Hu, An Yang, Jingren Zhou, and Junyang Lin. Stabilizing reinforcement learning with llms: Formulation and practices, 2025. https://arxiv.org/abs/2512. 01374. [74] Lianmin Zheng, Liangsheng Yin, Zhiqiang Xie, Chuyue Sun, Jeff Huang, Cody Hao Yu, Shiyi Cao, Christos Kozyrakis, Ion Stoica, Joseph E. Gonzalez, Clark Barrett, and Ying Sheng. Sglang: efficient execution of structured language model programs. In Proceedings of the 38th International Conference on Neural Information Processing Systems, 2024. https://dl.acm.org/doi/10.5555/3737916.3739916. [75] Yinmin Zhong, Zili Zhang, Xiaoniu Song, Hanpeng Hu, Chao Jin, Bingyang Wu, Nuo Chen, Yukun Chen, Yu Zhou, Changyi Wan, Hongyu Zhou, Yimin Jiang, Yibo Zhu, and Daxin Jiang. Streamrl: Scalable, heterogeneous, and elastic rl for llms with disaggregated stream generation, 2025. https://arxiv.org/abs/2504.15930.

15

[76] Yanqi Zhou, Tao Lei, Hanxiao Liu, Nan Du, Yanping Huang, Vincent Zhao, Andrew M Dai, zhifeng Chen, Quoc V Le, and James Laudon. Mixture-of-experts with expert choice routing. In Advances in Neural Information Processing Systems, 2022. https://proceedings.neurips.cc/paper_files/paper/2022/file/ 2f00ecd787b432c1d36f3de9800728eb-Paper-Conference.pdf.

16

𝑇layer ≥ 𝑇attn + 𝑇moe , and overlap budgets are taken against this conservative estimate. The relative magnitudes of the FFN and attention terms are governed by

A Detailed Overhead Analysis of the Expert Transfer Engine This appendix provides a formal overhead analysis of the two Expert Transfer Engine designs described in §6. We first establish a per-layer compute time model that expresses the GPU’s available overlap budget in terms of model and workload variables (§A.1). We then analyze the CPU-assisted design (§A.2) and the GPU-direct design (§A.3), each with explicit overlap conditions and scaling discussions. §A.4 concludes with a concrete instantiation for the Qwen3-30B-A3B model on our H20 platform as numerical confirmation. For the analysis we consider one rank processing a single sequence of length 𝑛 in one micro-step. Throughout, let 𝐸 denote the number of experts per MoE layer, 𝐾 the number of experts activated per token (top-𝐾), ℎ the hidden dimension, ℎ 𝑓 𝑓 the FFN intermediate dimension of one expert, and 𝐿 the number of MoE layers held by each training worker (after pipeline parallelism partitioning, if any). Let 𝑆𝑒 denote the size of one expert in the model’s training precision and 𝑆𝑔 the size of one expert’s gradient tensor (in the precision used for gradient accumulation). For a SwiGLU expert with three weight matrices (gate, up, down) of size ℎ × ℎ 𝑓 𝑓 each, 𝑆𝑒 = 3 ℎ ℎ 𝑓 𝑓 ·𝑝 𝑤 bytes where 𝑝 𝑤 is the per-element parameter byte width, and 𝑆𝑔 has the same structure with the per-element gradient byte width 𝑝𝑔 . Let 𝑁𝑠 denote the number of experts each rank holds in the placement (base plus redundant), 𝐹 the rank’s effective compute throughput in FLOPS, 𝐵 PCIe the achievable host-to-device PCIe bandwidth per GPU, and 𝐵 NVLink the achievable per-GPU NVLink bandwidth. A.1

6 𝐾 ℎ𝑓 𝑓 𝑇moe ≈ , 𝑇attn 8ℎ + 2𝑛

which grows linearly in 𝐾 and ℎ 𝑓 𝑓 and shrinks as 𝑛 increases. A.2

We model the per-layer compute time of one rank as the sum of the attention block and the MoE block. Following the standard FLOPs accounting for transformer layers [38, 55], the per-rank forward FLOPs of one layer are: 8 𝑛 ℎ2 |{z} QKV+output proj.

FLOPsmoe =

2 𝑛2 ℎ |{z}

+

,

𝑇prefetch =

6 𝐾 𝑛 ℎ ℎ𝑓 𝑓 | {z }

𝑇moe =

(16)

Substituting Eqs. (13)–(14), with 𝑆𝑒 = 3 ℎ ℎ 𝑓 𝑓 𝑝 𝑤 , the inequality becomes a quadratic constraint on 𝑛:

(11)

+ 2 𝑛 ℎ 𝐸, | {z }

2 𝑛 2 + (8 ℎ + 6 𝐾 ℎ 𝑓 𝑓 + 2 𝐸) 𝑛 ≥

(12)

router

cpu

6 𝐾 𝑛 ℎ ℎ𝑓 𝑓 + 2 𝑛 ℎ 𝐸 . (13) 𝐹

The total compute time of one MoE layer on this rank is 𝑇layer = 𝑇attn + 𝑇moe + 𝑇a2a,

3 𝑁𝑠 ℎ 𝑓 𝑓 𝑝 𝑤 · 𝐹 , 𝐵 PCIe cpu

(17)

which is satisfied for all 𝑛 ≥ 𝑛 min , where 𝑛 min is the positive root of the corresponding quadratic equation. Discussions. Two observations follow from Eq. (17). First, the LHS is strictly increasing in 𝑛 while the RHS is independent of 𝑛, so the inequality is always satisfied for sufficpu ciently large sequences. The question is only whether 𝑛 min cpu is small enough to be of practical concern. Second, 𝑛 min depends on the hardware through the ratio 𝐹 /𝐵 PCIe , the GPU’s compute-density-per-byte. For modern accelerators this ratio is on the order of 103 FLOPs/byte, while the LHS coefficient 8ℎ + 6𝐾ℎ 𝑓 𝑓 + 2𝐸 is on the order of 104 to 105 for productioncpu scale MoE models, so 𝑛 min stays in the low thousands of tokens, which is small compared to a typical RL-post-training

where the factor of 6 in the FFN term reflects two FLOPs per multiply-add applied to the three weight matrices of a SwiGLU expert, and the FFN term assumes that the 𝐾 𝑛 token-expert pairs received by this rank under a balanced placement are processed locally. Dividing by the rank’s effective throughput 𝐹 yields: 8 𝑛 ℎ2 + 2 𝑛2 ℎ , 𝐹

𝑁𝑠 · 𝑆 𝑒 ≤ 𝑇layer . 𝐵 PCIe

causal attention

expert FFN (SwiGLU)

𝑇attn =

CPU-assisted Expert Transfer

We analyze the storage and transfer overheads of the CPUassisted design used in the recompute stage. CPU memory. Each training worker stores a master copy of every expert in every layer it owns, requiring 𝐿 · 𝐸 · 𝑆𝑒 bytes of pinned host memory. Because the recompute stage carries no gradient or optimizer state, this is the only expertrelated memory placed on CPU, and modern host memory comfortably exceeds 𝐿 · 𝐸 · 𝑆𝑒 for the model sizes targeted by ForeMoE. GPU memory. Each rank holds at most 𝑁𝑠 experts at a time. Incoming prefetches are written directly into the GPU memory regions that hold the rank’s expert parameters, so no separate buffer is required on the GPU. The per-rank GPU footprint is Θ(𝑁𝑠 · 𝑆𝑒 ), which matches the working set of a baseline EP layout with the same per-rank expert count. The CPU-assisted design therefore introduces no additional GPU memory overhead. Prefetch latency. In the worst case, all 𝑁𝑠 experts on a rank are replaced between two consecutive micro-steps, generating 𝑁𝑠 · 𝑆𝑒 bytes of host-to-device traffic per rank per layer. The prefetch is fully hidden whenever

Per-Layer Compute Time Model

FLOPsattn =

(15)

(14)

where 𝑇a2a accounts for the dispatch/combine All-to-All. For the overhead analyses below it suffices to lower-bound 17

micro-step. Larger ℎ, ℎ 𝑓 𝑓 , 𝐾, or 𝐸 all increase the LHS coefficpu cient and therefore shrink 𝑛 min further, so the bound holds increasingly comfortably for larger models. A.3

alone, while FP32 gradient accumulation roughly triples it. Both regimes remain well within the 𝑇attn budget on modern NVLink fabrics, as the next subsection illustrates.

GPU-direct Expert Transfer (Intra-machine) A.4

We analyze the storage and transfer overheads of the GPUdirect design used in the policy update stage. Because each swapped expert must take its gradient with it (the gradient produced under the previous placement must continue to be accumulated against the same parameter on its new rank), every transfer in this stage carries both the expert’s parameters and its gradients. GPU memory. Each rank still holds at most 𝑁𝑠 experts. The transfer pipeline additionally allocates one send buffer and one receive buffer per rank, each sized to hold a rank’s expert state (parameters and gradients of all 𝑁𝑠 experts on the rank), and shared across all MoE layers. The asymptotic GPU memory overhead is therefore Θ(𝑁𝑠 · (𝑆𝑒 + 𝑆𝑔 )) on top of the placement working set. Swap latency. Copy-out and Copy-in are intra-GPU memory operations whose cost is negligible compared to the MoE combine and dispatch they overlap with. The dominant stage is the NVLink swap, which carries up to 𝑁𝑠 · (𝑆𝑒 + 𝑆𝑔 ) bytes per rank per layer and is overlapped with Attention. The swap is therefore fully hidden whenever 𝑇swap =

𝑁𝑠 (𝑆𝑒 + 𝑆𝑔 ) ≤ 𝑇attn . 𝐵 NVLink

(18)

Writing 𝑆𝑒 + 𝑆𝑔 = 3 ℎ ℎ 𝑓 𝑓 (𝑝 𝑤 + 𝑝𝑔 ) and substituting Eq. (13), the condition becomes 2 𝑛2 + 8 ℎ 𝑛 ≥

3 𝑁𝑠 ℎ 𝑓 𝑓 (𝑝 𝑤 + 𝑝𝑔 ) · 𝐹 , 𝐵 NVLink

Concrete Instantiation: Qwen3-30B-A3B on H20

We instantiate the bounds with the Qwen3-30B-A3B [64] model and our H20-based deployment. The model has ℎ = 2048, ℎ 𝑓 𝑓 = 768, 𝐸 = 128, 𝐾 = 8, and 48 MoE layers. We assume BF16 weights (𝑝 𝑤 = 2) with FP32 gradient accumulation (𝑝𝑔 = 4), so 𝑆𝑒 = 6 ℎ ℎ 𝑓 𝑓 = 9.0 MiB and 𝑆𝑔 = 12 ℎ ℎ 𝑓 𝑓 = 18.0 MiB per expert. Per H20 GPU, 𝐹 ≈ 148 TFLOPS in BF16, 𝐵 PCIe ≈ 64 GB/s (PCIe Gen5 x16), and 𝐵 NVLink ≈ 450 GB/s. We take 𝑁𝑠 = 𝐸/𝑃 + 𝑁𝑟 = 128/16 + 2 = 10 as a representative working set (EP=16 layout with two redundant slots per rank), and 𝑛 = 8 K as a representative sequence length for long-context RL post-training. CPU storage. Per worker, 𝐿·𝐸·𝑆𝑒 = 48·128·9.0 MiB = 54 GiB of pinned host memory, well within the ≥ 1 TiB host capacity of our machines. Prefetch budget (recompute). Plugging the parameters cpu above into Eq. (17) gives 𝑛 min ≈ 1862 tokens. A single rank in our deployments routinely processes tens of thousands of tokens per micro-step, so this ∼ 1.9 K-token threshold is met by an order of magnitude or more. Equivalently, the worst-case prefetch volume of 𝑁𝑠 · 𝑆𝑒 = 90 MiB takes 1.47 ms over PCIe, while 𝑇layer for 𝑛 = 8 K tokens is ≈ 7.9 ms, leaving more than 6 ms of headroom. Swap budget (policy update). Plugging the parameters above into Eq. (19) gives 𝑛 nv min ≈ 2189 tokens, again well below the per-rank token count of a typical micro-step. Equivalently, the worst-case swap volume of 𝑁𝑠 · (𝑆𝑒 +𝑆𝑔 ) = 270 MiB takes 0.63 ms over NVLink, while 𝑇attn for 𝑛 = 8 K tokens is ≈ 3.7 ms, a 5.9× margin. These numerical margins confirm that the inequalities Eq. (16) and Eq. (18) hold by a comfortable factor in our deployment, and the scaling discussions of §A.2 and §A.3 indicate that the margins grow rather than shrink as models become deeper, wider, or longer-context.

(19)

nv which is satisfied for all 𝑛 ≥ 𝑛 nv min , where 𝑛 min is the positive root of the corresponding quadratic equation. GPU compute contention. NVLink collectives are realized through GPU streaming multiprocessors that issue the underlying memory transactions, which introduces a small amount of contention with the compute stream during the swap. This contention is bounded by the small fraction of total layer time that the NVLink transfer occupies, and we measure it to be negligible in practice. Discussions. Three observations follow from Eq. (19). First, the relevant compute-density-per-byte is 𝐹 /𝐵 NVLink , which on modern intra-machine fabrics is roughly 5–10× smaller than 𝐹 /𝐵 PCIe , so the high NVLink bandwidth dominates the gradient-inflated transfer volume. Second, the swap is hidden behind Attention alone (rather than the full 𝑇layer ), which is why Eq. (19) lacks the FFN-related terms. In exchange, the LHS still grows quadratically in 𝑛 via the attention compute, so 𝑛 nv min remains modest as long as the sequence length is not vanishingly small. Third, the linear factor (𝑝 𝑤 + 𝑝𝑔 ) exposes the cost of carrying gradients: training in pure BF16 (𝑝𝑔 = 𝑝 𝑤 = 2) doubles the swap volume relative to parameters

B

On a CPU-assisted Variant for the Policy Update Stage

The CPU-assisted design adopted for the recompute stage (§A.2) gains its expressiveness from holding the entire expert pool in pinned host memory and prefetching only what each micro-step requires. A natural question is whether the same approach extends to the policy update stage, where it would offer the same global placement flexibility. We have given this extension careful consideration. As we describe below, the bandwidth side admits a workable schedule (§B.1), but the memory side exposes a structural constraint specific to gradient transfer (§B.2) that ultimately drives our choice of the GPU-direct design as the default in ForeMoE. 18

B.1

memory differs from ordinary host memory in that it is nonpageable kernel-managed memory and competes directly with the OS file cache and other process allocations. Committing it at this scale degrades machine-wide throughput well before nominal RAM is exhausted, and risks OOM under memory-pressure spikes.

Bandwidth

Extending the recompute-stage CPU-assisted design to policy update introduces a second entity beyond the parameters: the main gradients, which are roughly twice the size of the parameters under BF16 weights and FP32 gradient accumulation (𝑝 𝑤 = 2, 𝑝𝑔 = 4, hence 𝑆𝑔 = 2 𝑆𝑒 ), and which must additionally be written back to CPU after each backward layer’s gradient accumulation completes. We allocate the parameter and gradient traffic across the two passes as follows. The forward pass is kept identical to the recompute case: during the forward of layer 𝐿, prefetch the parameters of layer 𝐿 + 1 over PCIe on a separate CUDA stream, overlapping the on-going forward computation. The backward pass then carries all of the gradient traffic: during the backward of layer 𝐿, the transfer stream issues, in pipelined order, (1) a prefetch of the gradients of layer 𝐿 − 1 so that they are on GPU before 𝐿 − 1’s backward begins, and (2) an offload of the just-accumulated gradients of 𝐿 to CPU. Per layer per rank, the resulting transfer volumes and overlap budgets are: • Forward: transfers 𝑁𝑠 𝑆𝑒 bytes, overlapped against 𝑇layer (identical to the recompute case). • Backward: transfers 2 𝑁𝑠 𝑆𝑔 = 4 𝑁𝑠 𝑆𝑒 bytes (one prefetch and one offload of the gradients), overlapped against backward compute of approximately 2𝑇layer . The transfer-to-compute ratio is 𝑁𝑠 𝑆𝑒 /(𝐵 PCIe 𝑇layer ) on the forward side and 2 𝑁𝑠 𝑆𝑒 /(𝐵 PCIe 𝑇layer ) on the backward side. Therefore backward is the binding constraint, with overlap budget approximately twice as tight as in the recompute case (e.g., on Qwen3-30B-A3B + H20 the binding margin shrinks from ∼ 5.4× to ∼ 2.7×). B.2

B.3

Outlook

The schedule of §B.1 tightens the overlap requirement above that of the recompute case, and the memory analysis of §B.2 shows that a per-worker gradient master copy at production model scale exceeds the practical pinned-memory budget of present-day machines. ForeMoE therefore adopts the GPUdirect design of §A.3 as its policy-update transfer engine. The CPU-assisted variant nonetheless remains an attractive direction in hardware regimes with abundant pinned host memory and lightweight cross-worker gradient-synchronization mechanisms, and we leave a careful exploration to future work.

C

MILP Constraint Details

This appendix expands on the four constraints used in the joint MILP formulation of §7.2, explaining the structural property each one enforces. Í Expert slot capacity ( 𝑒 𝑥𝑒,𝑗 = 1). A slot is a single physical position on a rank that can hold one expert tensor, so two experts cannot occupy the same slot. Í Expert coverage ( 𝑗 𝑥𝑒,𝑗 ≥ 1). Every expert in the model must be placed somewhere; otherwise its tokens have no destination. The inequality (rather than equality) allows replicas: an expert may be placed in more than one slot whenever the load profile justifies replication. Í Token conservation ( 𝑗 𝑟𝑠,𝑒,𝑗 = 1). For every (source rank, expert) pair, the entire token volume must be dispatched to some destination slot. This forbids dropping tokens. When expert 𝑒 has multiple replicas, this constraint allows source 𝑠’s tokens to be assigned among them in any proportions that sum to one. Token assignment feasibility (𝑟𝑠,𝑒,𝑗 ≤ 𝑥𝑒,𝑗 ). Tokens for expert 𝑒 may travel only to slots that actually hold 𝑒. If expert 𝑒 is not in slot 𝑗 (so 𝑥𝑒,𝑗 = 0), this constraint forces 𝑟𝑠,𝑒,𝑗 = 0, preventing any tokens from being routed to a slot that does not host the corresponding expert.

Memory: A Structural Constraint on Gradient Sharing

Parameters and gradients are not symmetric in how they consume memory resource. The parameters are read-only and identical across workers, so a single shared memory-mapped region per machine suffices, leaving the pinned-memory footprint unchanged from the recompute case (𝐿 · 𝐸 · 𝑆𝑒 per machine). The gradients, by contrast, are write-shared via expert replication: when an expert 𝑒 is replicated on multiple ranks of the same machine, each rank’s backward pass independently accumulates into "the gradient of 𝑒", so a shared backing buffer would invite concurrent writes that demand atomic accumulation or explicit cross-process locking, neither of which is cheap to implement in pinned host memory accessed over PCIe. The conservative alternative is to give every worker its own private gradient master copy, at 𝐿 ·𝐸 ·𝑆𝑔 pinned bytes per worker. For Qwen3-30B-A3B this is 48·128·18 MiB = 108 GiB per worker, which on an 8-GPU machine sums to ≈ 0.85 TiB. For Qwen3-235B-A22B without pipeline parallelism it climbs into the multi-TiB range per machine. Moreover, pinned

D

Policy Update Stage Algorithm

The policy update stage uses the GPU-direct exchange within a machine of §6.1, which carries both parameters and gradients per expert and confines relocation and replication to intra-machine NVLink transfers. In principle, Stages 2-4 of Algorithm 2 can be reused here almost unchanged, the only modification being an intra-machine restriction that rejects any swap or replica placement requiring a cross-machine 19

and assign each unit to the replica whose hosting rank currently has the minimum accumulated load, until the replicas’ ranks converge toward equal utilization. We use water-filling rather than the LP of Eq. (10) because intra-machine replicas do not affect cross-machine communication, leaving the LP’s joint computation-communication objective with little additional optimization space over the simpler heuristic.

Algorithm 3 Policy Update Stage Planner (Stages 2-4) Require: Base placement 𝐴base , per-micro-step loads {𝑊 (𝑖 ) }𝑖=1..𝑁 , machines 𝑀, ranks per machine 𝑅, redundant slots per rank 𝑁𝑟 Ensure: Per-micro-step placements {𝐴 (𝑖 ) }, assignments {𝑟 (𝑖 ) } 1: for each micro-step 𝑖 = 1..𝑁 in parallel do 2: 𝐴 (𝑖 ) ← 𝐴base 3: for each machine 𝑚 = 1..𝑀 in parallel do 4: // Stage 2: intra-machine relocation via LPT 5: local_experts ← {𝑒 : machine(𝐴 (𝑖 ) [𝑒]) = 𝑚} Í (𝑖 ) 6: Sort local_experts by descending 𝑠 𝑤𝑠,𝑒 7: RL[𝑟 ] ← 0 for 𝑟 ∈ ranks(𝑚) 8: for each 𝑒 in sorted local_experts do 9: 𝑟 ∗ ← arg min𝑟 ∈ranks(𝑚) RL[𝑟 ] Í (𝑖 ) 10: 𝐴 (𝑖 ) [𝑒] ← 𝑟 ∗ ; RL[𝑟 ∗ ] += 𝑠 𝑤𝑠,𝑒 11: end for 12: // Stage 3: intra-machine replication (𝑅 · 𝑁𝑟 slots) 13: for 𝑘 = 1 to 𝑅 · 𝑁𝑟 do 14: 𝑒 ∗ ← locally heaviest expert 15: 𝑟 ∗ ← arg min𝑟 ∈ranks(𝑚) RL[𝑟 ] 16: Place a replica of 𝑒 ∗ on 𝑟 ∗ and update RL 17: end for 18: // Stage 4: water-filling token assignment 19: for each replicated expert 𝑒 on 𝑚 do Assign each token unit to the replica whose rank 20: has the smallest accumulated load 21: end for 22: end for 23: end for 24: return {𝐴 (𝑖 ) }, {𝑟 (𝑖 ) }

transfer. However, we observe that once the search space is restricted to a single machine, the problem becomes considerably simpler, and a lighter-weight procedure with lower complexity attains the same balancing quality as the restricted form of Algorithm 2. We therefore adopt this simpler procedure, given in Algorithm 3. We still reuse the base placement 𝐴base from Stage 1 as the starting point of every micro-step. We then apply the simpler Stages 2-4 to refine the per-micro-step placement within each machine. Each of Stages 2-4 decomposes into 𝑀 independent per-machine subproblems, which are solved in parallel. We now describe these Stages 2-4 in detail. Stage 2: intra-machine relocation. For each machine, sort its hosted experts in descending order of micro-step-level Í (𝑖 ) load 𝑠 𝑤𝑠,𝑒 and redistribute them across the machine’s local ranks via LPT. This rebalances per-rank computation load within the machine. Stage 3: intra-machine replication. For each machine, fill its local redundant slots one at a time. At each step, choose the locally heaviest expert and place its replica on the local rank with the minimum current load. Stage 4: water-filling token assignment. Within each machine, distribute tokens for each replicated expert to its replicas by water-filling: iterate over source-rank tokens 20