Mixture-of-Experts (MoE) models have recently emerged as a powerful technique for enhancing the scalability and performance of neural networks, primarily by leveraging learnable gating networks to allocate input tokens to different expert models. However, training MoE models on GPUs presents unique challenges, including insufficient GPU memory capacity for a large number of experts and computational inefficiency due to token load imbalance. To address these issues, we introduce Expert Server MoE (ES-MoE), a novel method that offloads all expert parameters and their optimizer states to CPUs. This approach not only mitigates the memory constraints of GPU-based training but also enhances training throughput by creating a unified pool of experts that allows for more efficient scheduling. Furthermore, ES-MoE employs pipelined expert optimization to minimize the iteration latency, effectively circumventing the issue of extended CPU optimization time. We validate our approach using GPT-based MoE architectures, demonstrating that ES-MoE scales up to 16 times better, and improves throughput up to 4.55x over the existing frameworks.
전문가 혼합 모델은 주로 학습 가능한 게이팅 네트워크를 활용하여 입력 토큰을 여러 전문가 모델에 할당함으로써 신경망의 확장성과 성능을 향상시키는 강력한 기법으로 최근 부상하고 있다. 그러나 그래픽 처리 장치에서 전문가 혼합 모델을 훈련하는 데는 많은 수의 전문가를 위한 그래픽 처리 장치 메모리 용량 부족, 토큰 부하 불균형으로 인한 계산 비효율성 등 고유한 과제가 있습니다. 이러한 문제를 해결하기 위해 모든 전문가 매개변수와 최적화 도구 상태를 중앙 처리 장치의 메모리로 오프로드하는 새로운 방법인 전문가 서버 기법을 도입했다. 이 접근 방식은 그래픽 처리 장치 기반 훈련의 메모리 제약을 완화할 뿐만 아니라 보다 효율적인 스케줄링이 가능한 통합된 전문가 풀을 생성하여 훈련 처리량을 향상시킨다. 또한 전문가 서버 기법은 모델 최적화를 파이프라인을 이용하여 늦어지는 학습 시간을 최소화하고 모델 성능을 향상시켜 중앙 처리 장치 최적화 문제를 효과적으로 개선한다. 전문가 서버 기법이 전문가 혼합 모델을 학습시키는 대에 있어서 기존 프레임워크에 비해 최대 16배 더 잘 확장되고 처리량이 최대 4.55배 향상됨을 입증했다.
