FEPLB: Exploiting Copy Engines for Nearly Free MoE Load Balancing in Distributed Training

Shuyao Qi; Haoyuan Liu; Shizhen Zhao

Home/Archive/arXiv

arXiv · 2604.19654

FEPLB: Exploiting Copy Engines for Nearly Free MoE Load Balancing in Distributed Training

Shuyao Qi·Haoyuan Liu·Shizhen Zhao

View original source ↗Read source document ↗Export metadata ↓

Abstract

Fine-grained, per-micro-batch load balancing is essential for efficient Mixture-of-Experts (MoE) training, yet every prior dynamic scheduling scheme pays for it with extra communication that is hard to hide. Especially on modern bulk-transfer backends such as DeepEP. We make a simple but consequential observation: on the NVIDIA Hopper architecture the NVLink Copy Engine can move data between intra-node GPUs without consuming any SM cycles, effectively providing a nearly free communication channel that runs in parallel with compute kernels. FEPLB turns this idle hardware into a new parallel dimension for MoE load rebalancing. Its Two-Phase Dispatch first routes tokens across nodes via the standard EP backend, then redistributes dynamic-expert tokens and weights within the NVLink domain through the Copy Engine at nearly zero cost, while a lightweight CPU scheduler runs concurrently with static expert computation. Because FEPLB uses only Copy Engine and CPU that are orthogonal to those consumed by EP and PP, it coexists with existing parallel strategies without reconfiguration. On GLM-5's MoE layers (128 experts, no auxiliary loss, up to 16 H100 GPUs), FEPLB reduces the token straggler by 51-70% and the GEMM straggler by 50-68% with no measurable EP communication overhead. Its advantage grows with the EP degree: at EP=8, it achieves 2x lower token straggler than FasterMoE.

Explore related subjects

cs.DC

Keep this discovery

Explore connections, maps & timelines

BibTeX RIS

Shuyao Qi, Haoyuan Liu, Shizhen Zhao. 2026-04-21. FEPLB: Exploiting Copy Engines for Nearly Free MoE Load Balancing in Distributed Training. https://arxiv.org/abs/2604.19654

Cite the original work for its findings. Save a collection to share your selection of sources.

KEEP EXPLORING

Related papers

arXiv2026

Profiling Concurrent Vision Inference Workloads on NVIDIA Jetson -- Extended

The proliferation of IoT devices and advancements in network technologies have intensified the demand for real-time data processing at the network edge. To address these demands, low-power AI accelerators, particularly GPUs, are increasingly deployed for inference tasks, enabling efficient computation while mitigating the latency and bandwidth limitations of cloud-based systems. Despite their growing deployment, GPUs remain underutilised even in computationally intensive workloads. This underutilisation stems from the limited understanding of GPU resource sharing, particularly in edge computing scenarios. In this work, we conduct a detailed analysis of both high- and low-level metrics, including GPU utilisation, memory usage, streaming multiprocessor (SM) utilisation, and tensor core usage, to identify bottlenecks and guide hardware-aware optimisations. By integrating traces from multiple profiling tools, we provide a comprehensive view of resource behaviour on NVIDIA Jetson edge devices under concurrent vision inference workloads. Our findings indicate that while GPU utilisation can reach $100\%$ with specific optimisations, critical low-level resources, such as SMs and tensor cores, often operate at only $15\%$ to $30\%$ utilisation. Moreover, we observe that certain CPU-side events, such as thread scheduling and context switching, frequently become bottlenecks, further constraining overall GPU performance. We provide several key observations for users of vision inference workloads on NVIDIA edge devices.

cs.DC↗

arXiv2026

Mask-Aware Execution for Efficient JEPA Training

Joint Embedding Predictive Architectures (JEPAs) are becoming a core representation-learning primitive and a building block for latent world models across vision, video, audio, brain dynamics, and time series. Despite (potential of) wide deployment, current JEPA training pipelines are inefficient: each input is executed through multiple mask-specific branches, with redundant target-side work, and memory-bound token routing. These costs grow with the number of masks and limit GPU efficiency. We present M-JEPA, a mask-aware execution architecture that restructures JEPA training without changing the learning objective. M-JEPA separates mask-independent computation from mask-dependent routing, enabling shared context encoder execution, fused token routing and slicing with backward support, sparse target encoder execution over the union of target tokens, and masked patch embedding for sparse inputs. The resulting pipeline preserves training semantics while reducing computation, memory traffic, and synchronization overhead. We implement M-JEPA for five JEPA variants and evaluate it on NVIDIA A100 GPUs. Compared against the state-of-the-art baselines, M-JEPA achieves up to 1.7x end-to-end training speedup for 2-10 masks. Separately, with masked patch embedding, 4.75x patch-embedding speedup at high sparsity. These results show that execution restructuring, rather than changes to the JEPA objective, is a key lever for efficient JEPA training.

cs.DC↗

arXiv2026

Multimmit: Extending Blocks for Faster Finality

To meet the throughput demands of modern blockchain systems, protocols for State Machine Replication (SMR) increasingly have many processors disseminate blocks of transactions in parallel, with consensus then establishing a total ordering on the blocks of all producers. Such designs face a choice as to when a block may enter the ordering. Certified approaches wait for a quorum to attest a block's availability, which is robust but adds message delays to every transaction. Uncertified approaches let proposals reference blocks immediately, which is fast but degrades rapidly when referenced data must be fetched on the critical path. Raptr, the state of the art, takes a middle course, finalising the longest prefix of the leader's proposal that a quorum holds, so that no processor ever blocks or fetches. The remaining weakness is sensitivity to order: if the data behind a single early batch is withheld, the proposal finalises little or nothing, so individual faulty producers can still deny the system its optimistic path. We present Multimmit, a protocol for $n \ge 5f+1$ processors combining a consensus layer requiring one round of voting per view with multi-chain data dissemination. Votes are cast relative to the leader's proposal, reporting per chain how far the voter can support it, and may themselves attest fresh blocks beyond it. A transaction block disseminated at time $t$ is ordered by $t+3δ$ in expectation and $t+2δ$ at best, measured from the block's dissemination rather than the leader's proposal. Degradation under faults is graceful: a faulty producer delays only its own chain's blocks, costing other chains at most a one-view wait for placement. No leader can both finalise its leader block and exclude a fresh, well-circulated block of an honest chain. Consensus traffic is tens of kilobytes per view, independent of transaction volume.

cs.DC↗