通过智能调度专家和数据打包，显著提升MoE强化学习的训练吞吐量。
RoutePack: Expert Placement and Attention-Aware Data Packing for MoE Reinforcement Learning
- 分层规划专家位置与数据打包，协同优化注意力与专家负载。
- 在两个模型上分别提升吞吐量8.85%和14.89%。
- 适合大规模MoE强化学习系统优化，尤其关注训练效率的团队。
在强化学习中训练混合专家（MoE）模型面临双重负载均衡挑战：序列组成决定数据并行微批次中的密集注意力工作量，而令牌路由决定专家并行秩上的稀疏专家工作量。单独优化任一环节可能将瓶颈转移至另一环节。在MoE RL中，回放阶段的路由可提前知晓每个样本的序列长度与逐层专家需求。我们提出RoutePack，一种分层规划器，协调状态一致的逐层专家重路由，并在优化器步窗口内联合考虑注意力与专家负载进行数据打包。RoutePack首先基于聚合路由需求在各MoE层独立放置专家；随后将样本打包进最少的、满足令牌上限的执行行，并以投影的EDP分片感知目标优化其数据并行布局。该目标结合窗口归一化的线性-二次注意力代理与逐层物理专家分片峰值，最小化最慢EDP分片的累积成本。并行种群退火在保持样本覆盖、容量、非空单元、均等微批次数与通信拓扑的前提下，搜索固定行可行布局。状态一致的实例化保留逻辑top-k路由与现有MoE内核，无需微批次级专家复制。在Ling-3.0-Tiny和Ling-3.0-Flash上，专家重路由分别提升平均训练吞吐量3.80%和10.50%，路由感知打包进一步带来4.86%和3.98%提升，整体吞吐量较基线提高8.85%和14.89%。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。