RoutePack: Expert Placement and Attention-Aware Data Packing for MoE Reinforcement Learning

2026-08-12

AI Summary

This work addresses the dual load-balancing bottlenecks in Mixture-of-Experts (MoE) reinforcement learning arising from attention computation and expert routing. To this end, we propose RoutePack, a hierarchical co-optimization framework that jointly optimizes expert placement and sequence packing for the first time. Within each optimizer-step window, RoutePack enables inter-layer expert rerouting and attention-expert-aware data packing. Its key innovations include a state-consistent rerouting mechanism, expert placement decoupled via aggregated routing demands, a token-capacity-aware row packing strategy, and a projection-based EDP sharding-aware objective function, with layout efficiently searched via parallel population annealing. Experiments demonstrate that RoutePack improves token throughput by 8.85% and 14.89% on Ling-3.0-Tiny and Ling-3.0-Flash, respectively, without modifying existing MoE kernels.

Abstract

Training Mixture-of-Experts (MoE) models for reinforcement learning (RL) couples two load-balancing problems: sequence composition determines dense attention work in each data-parallel microbatch, while token routing determines sparse expert work on expert-parallel ranks. Optimizing either alone can shift the bottleneck to the other. In MoE RL, rollout-time routing replay exposes every sample's sequence length and layer-wise expert demand before its training step. We present RoutePack, a hierarchical planner that coordinates state-consistent, layer-wise expert rerouting with joint attention- and expert-aware data packing over an optimizer-step window. RoutePack first places experts independently at each MoE layer using aggregate routing demand. It then packs samples into the smallest certified, or best-known feasible, number of token-capped execution rows and optimizes their DP layout with a projected EDP-shard-aware objective. The objective combines a window-normalized linear-quadratic attention proxy with per-layer physical EP-rank peaks and minimizes the accumulated cost of the slowest EDP shard. Parallel population annealing searches fixed-row feasible layouts while preserving sample coverage, capacity, nonempty cells, equal microbatch counts, and communicator topology. State-consistent materialization preserves logical top-k routing and existing MoE kernels without microbatch-level expert replication. Across Ling-3.0-Tiny and Ling-3.0-Flash, expert rerouting improves mean trainer-measured token throughput by 3.80% and 10.50%, while routing-aware packing adds another 4.86% and 3.98%, respectively. Overall, RoutePack improves throughput by 8.85% and 14.89% over the baseline.

Problem

• Mixture-of-Experts
• Reinforcement Learning
• Load Balancing
• Token Routing
• Attention Computation

Innovation

• Mixture-of-Experts
• reinforcement learning
• expert routing
• data packing
• load balancing