글로벌 배치 부하 분산: MoE LLM 훈련을 개선하는 거의 공짜 점심

Mixture-of-Experts(MoE) 아키텍처는 모델 파라미터를 확장하는 인기 있는 기술이지만, Megatron-core와 같은 기존 훈련 프레임워크는 마이크로 배치 수준의 부하 균형 손실(micro-batch level load balancing loss)을 사용하여 마이크로 배치에 동질적인 데이터(예: 코드만 포함)가 있을 때 성능이 저하될 수 있습니다. 알리바바의 큐원(Qwen)은 글로벌 배치 균형 손실(global-batch balance loss)을 도입하여 모든 병렬 그룹에서 전문가 선택 빈도를 동기화하고 손실을 전역적으로 계산합니다. 세 가지 MoE 구성(3.4B/0.6B 활성화, 15B/2.54B 활성화, 43B/6.6B 활성화)과 최대 400B 토큰에 대한 실험 결과, 글로벌 배치 균형이 다양한 작업에서 마이크로 배치 균형을 일관되게 능가하는 것으로 나타났습니다. 또한, 글로벌 배치 균형은 전문가의 상당한 도메인 특화를 유도하는데, 이는 마이크로 배치 균형 모델에서는 나타나지 않습니다. 글로벌 손실 위에 작은 마이크로 배치 균형 손실(가중치 0.01)을 추가하면 효과성에 영향을 주지 않으면서 훈련 속도가 향상됩니다. 이 방법은 전문가 수와 동일한 크기의 벡터만 동기화하면 되므로 거의 무료에 가깝습니다. 이 연구는 arXiv 프리프린트 2501.11873에 게재되었습니다.