최소 부하 전문가 병렬화: 불균형 전문가 혼합 모델의 부하 분산

초록

전문가 혼합(MoE) 모델은 일반적으로 통계적으로 균형 잡힌 전문가 라우팅을 보장하기 위해 명시적 부하 분산 제약 조건과 함께 사전 학습됩니다. 그럼에도 불구하고, 잘 학습된 MoE 모델조차도 상당히 불균형한 라우팅을 보인다는 것을 관찰했습니다. 이러한 동작은 자연스러운—심지어 바람직한—것으로 여겨질 수 있습니다. 불균형 라우팅은 모델이 도메인 특화 지식을 전문가 하위 집합 내에 집중할 수 있게 하기 때문입니다. 전문가 병렬화(EP)는 전문가를 여러 장치에 분산하여 MoE 모델의 규모 확장성을 위해 설계되었지만, 균형 잡힌 라우팅이라는 덜 논의된 가정을 수반합니다. 극단적인 불균형 상황에서 EP는 과도한 수의 토큰을 소수의 전문가로 집중시켜, 명시적 부하 분산이 종종 적용되지 않는 사후 학습 또는 추론 단계에서 과부하 장치의 계산 및 메모리 한계 초과 오류를 야기할 수 있습니다. 우리는 과부하 장치에서 활용도가 낮은 장치로 초과 토큰과 관련 전문가 매개변수를 동적으로 재라우팅하는 새로운 EP 알고리즘인 최소 부하 전문가 병렬화(LLEP)를 제안합니다. 이를 통해 메모리 제약을 준수하면서 모든 장치가 최소 집합 지연 시간 내에 작업 부하를 완료할 수 있습니다. 다양한 모델 규모에서 LLEP는 표준 EP 대비 최대 5배의 속도 향상과 최대 4배의 피크 메모리 사용량 감소를 달성했습니다. 이는 gpt-oss-120b의 경우 약 1.9배 더 빠른, 더 빠르고 높은 처리량의 사후 학습 및 추론을 가능하게 합니다. 우리는 이 방법을 포괄적인 이론적 분석과 절제 연구를 포함한 체계적인 실증 평가로 뒷받침합니다. 이러한 결과는 주요 절충점을 밝히고 하드웨어별 하이퍼파라미터 조정을 위한 원칙적인 프레임워크를 제공하여 최적의 성능을 달성할 수 있게 합니다.

English

Mixture-of-Experts (MoE) models are typically pre-trained with explicit load-balancing constraints to ensure statistically balanced expert routing. Despite this, we observe that even well-trained MoE models exhibit significantly imbalanced routing. This behavior is arguably natural-and even desirable - as imbalanced routing allows models to concentrate domain-specific knowledge within a subset of experts. Expert parallelism (EP) is designed to scale MoE models by distributing experts across multiple devices, but with a less-discussed assumption of balanced routing. Under extreme imbalance, EP can funnel a disproportionate number of tokens to a small number of experts, leading to compute- and memory-bound failures on overloaded devices during post-training or inference, where explicit load balancing is often inapplicable. We propose Least-Loaded Expert Parallelism (LLEP), a novel EP algorithm that dynamically reroutes excess tokens and associated expert parameters from overloaded devices to underutilized ones. This ensures that all devices complete their workloads within the minimum collective latency while respecting memory constraints. Across different model scales, LLEP achieves up to 5x speedup and 4x reduction in peak memory usage compared to standard EP. This enables faster and higher-throughput post-training and inference, with ~1.9x faster for gpt-oss-120b. We support our method with extensive theoretical analysis and comprehensive empirical evaluations, including ablation studies. These results illuminate key trade-offs and enable a principled framework for hardware-specific hyper-parameter tuning to achieve optimal performance.

최소 부하 전문가 병렬화: 불균형 전문가 혼합 모델의 부하 분산

Least-Loaded Expert Parallelism: Load Balancing An Imbalanced Mixture-of-Experts

초록

Support