RoutePack: 전문가 배치 및 주의 메커니즘 기반 데이터 패킹을 통한 MoE 강화 학습

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

강화 학습(RL)을 위한 Mixture-of-Experts (MoE) 모델 학습은 두 가지 로드 밸런싱 문제를 동시에 해결해야 합니다. 시퀀스 구성은 각 데이터 병렬 마이크로 배치 내에서 집중적인 어텐션 연산을 결정하며, 토큰 라우팅은 전문가 병렬 랭크에서의 희소한 전문가 활용을 결정합니다. 이 중 어느 하나만 최적화하면 다른 부분에 병목 현상이 발생할 수 있습니다. MoE RL에서 롤아웃 시간의 라우팅 리플레이는 각 샘플의 시퀀스 길이와 레이어별 전문가 수요를 학습 단계 전에 파악할 수 있도록 합니다. 본 논문에서는 RoutePack이라는 계층적 플래너를 제안합니다. RoutePack은 최적화 단계 내에서 상태 일관성을 유지하며, 레이어별 전문가 재라우팅을 조정하고 어텐션 및 전문가 정보를 모두 고려한 데이터 패킹을 수행합니다. 먼저 RoutePack은 각 MoE 레이어에서 집계된 라우팅 수요를 사용하여 전문가를 독립적으로 배치합니다. 그런 다음, 샘플들을 토큰 수 제한된 실행 행의 최소 인증 가능 개수 또는 가장 잘 알려진 실현 가능한 개수로 묶고, 프로젝션된 EDP (Effective Data Parallelism) 샤드 인지 목표 함수를 사용하여 DP (Data Parallelism) 레이아웃을 최적화합니다. 이 목표 함수는 윈도우 정규화된 선형-이차 어텐션 근사치와 각 레이어의 물리적인 EP (Effective Parallelism) 랭크 피크를 결합하고, 가장 느린 EDP 샤드의 누적 비용을 최소화합니다. 병렬 파퓰레이션 어닐링은 샘플 커버리지, 용량, 비어 있지 않은 셀, 동일한 마이크로 배치 수 및 통신 토폴로지를 유지하면서 고정된 행의 실현 가능한 레이아웃을 검색합니다. 상태 일관성을 유지하는 구현 방식은 논리적인 top-k 라우팅과 기존 MoE 커널을 마이크로 배치 수준의 전문가 복제 없이 보존합니다. Ling-3.0-Tiny 및 Ling-3.0-Flash 모델에서 실험 결과, 전문가 재라우팅은 평균적으로 트레이너가 측정한 토큰 처리량을 각각 3.80% 및 10.50% 향상시켰으며, 라우팅 인지 데이터 패킹은 추가로 각각 4.86% 및 3.98%의 성능 향상을 가져왔습니다. 전체적으로 RoutePack은 기준 모델 대비 각각 8.85% 및 14.89%의 처리량 향상을 달성했습니다.

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.