[논문 리뷰] MoESys: A Distributed and Efficient Mixture-of-Experts Training and Inference System for Internet Services

Dianhai Yu, Liang Shen외 6명|arXiv (Cornell University)|2022. 05. 20.

Advanced Neural Network Applications인용 수 12

한 줄 요약

SE-MoE는 2D 프리패칭, 계층적 스토리지 융합, CPU-GPU 메모리 링 기반 추론을 통해 로드 밸런싱, 통신 효율성, 메모리 활용도를 향상시키는 분산형이며 효율적인 믹스처 오브 응용(모델) 학습 및 추론 시스템이다. DeepSpeed 대비 33% 높은 학습 처리량과 13% 높은 추론 처리량을 달성하였으며, 불균형 MoE 워크로드 상황에서는 처리량이 64% 높고 메모리가 18% 적게 사용되었다. 48개의 A100 GPU에서 12B 파라미터를 가진 UFO 모델을 8일 만에 성공적으로 학습시켰다.

ABSTRACT

While modern internet services, such as chatbots, search engines, and online advertising, demand the use of large-scale deep neural networks (DNNs), distributed training and inference over heterogeneous computing systems are desired to facilitate these DNN models. Mixture-of-Experts (MoE) is one the most common strategies to lower the cost of training subject to the overall size of models/data through gating and parallelism in a divide-and-conquer fashion. While DeepSpeed has made efforts in carrying out large-scale MoE training over heterogeneous infrastructures, the efficiency of training and inference could be further improved from several system aspects, including load balancing, communication/computation efficiency, and memory footprint limits. In this work, we present a novel MoESys that boosts efficiency in both large-scale training and inference. Specifically, in the training procedure, the proposed MoESys adopts an Elastic MoE training strategy with 2D prefetch and Fusion communication over Hierarchical storage, so as to enjoy efficient parallelisms. For scalable inference in a single node, especially when the model size is larger than GPU memory, MoESys builds the CPU-GPU memory jointly into a ring of sections to load the model, and executes the computation tasks across the memory sections in a round-robin manner for efficient inference. We carried out extensive experiments to evaluate MoESys, where MoESys successfully trains a Unified Feature Optimization (UFO) model with a Sparsely-Gated Mixture-of-Experts model of 12B parameters in 8 days on 48 A100 GPU cards. The comparison against the state-of-the-art shows that MoESys outperformed DeepSpeed with 33% higher throughput (tokens per second) in training and 13% higher throughput in inference in general. Particularly, under unbalanced MoE Tasks, e.g., UFO, MoESys achieved 64% higher throughput with 18% lower memory footprints.

연구 동기 및 목표

• 로드 불균형, 통신 병목 현상, 높은 메모리 사용량으로 인한 분산 MoE 학습 및 추론의 비효율성을 해결하기 위해.
• 이질적인 GPU-CPU 클러스터에서 확장성 있고 고처리량의 MoE 학습 및 추론을 가능하게 하기 위해.
• GPU 메모리 용량을 초과하는 MoE 모델을 배포할 때 메모리 압박을 줄이고 시스템 활용도를 향상시키기 위해.
• 새로운 시스템 수준 기법을 통해 MoE 시스템의 통신 및 계산 워크로드 최적화를 위해.
• 인터넷 서비스에서 대규모 MoE 모델을 최소한의 자원 오버헤드로 생산 규모로 구현할 수 있도록 지원하기 위해.

제안 방법

• 장치 간 데이터 및 모델 병렬성 효율성을 향상시키기 위해 2D 프리패칭을 활용한 유연한 MoE 학습을 도입한다.
• 노드 간 통신 오버헤드를 줄이고 대역폭 활용도를 향상시키기 위해 계층적 스토리지 기반의 융합 통신을 구현한다.
• GPU 메모리 용량을 초과하는 MoE 모델을 효율적으로 실행하기 위해 CPU-GPU 메모리 링 아키텍처를 활용한다.
• GPU 메모리 소비를 줄이고 학습 처리량을 향상시키기 위해 데이터 병렬성에서 임bedding 분할을 적용한다.
• 계층적 스토리지 및 메모리 관리를 활용해 데이터 이동과 D2H/H2D 오버헤드를 최소화한다.
• 전문가 불균형 문제를 완화하기 위해 최적화된 라우팅 및 로드 밸런싱 전략을 통합한다.

실험 결과

연구 질문

• RQ1 어떻게 분산형 이질 컴퓨팅 환경에서 MoE 학습 및 추론의 효율성을 높일 수 있는가?
• RQ2 MoE 학습에서 통신 및 메모리 병목 현상을 줄이기 위해 어떤 시스템 수준 최적화가 가능한가?
• RQ3 GPU 메모리 용량을 초과하는 MoE 모델은 어떻게 효율적으로 추론할 수 있는가?
• RQ4 시스템 수준 설계를 통해 MoE 라우팅의 로드 불균형 문제를 어느 정도 완화할 수 있는가?
• RQ5 통합된 시스템이 MoE 학습 및 추론에서 고처리량과 낮은 메모리 프로파일을 동시에 달성할 수 있는가?

주요 결과

• SE-MoE는 표준 MoE 워크로드에서 DeepSpeed 대비 33% 높은 학습 처리량(초당 토큰 수)을 달성했다.
• SE-MoE는 일반 MoE 작업 전반에서 DeepSpeed 대비 13% 높은 추론 처리량을 기록했다.
• UFO 모델과 같은 불균형 MoE 워크로드 상황에서 SE-MoE는 DeepSpeed 대비 처리량이 64% 높고 메모리 프로파일이 18% 낮았다.
• 시스템은 48개의 A100 GPU를 사용해 12B 파라미터를 가진 통합 기능 최적화(UFO) MoE 모델을 8일 만에 성공적으로 학습시켰다.
• 데이터 병렬성에서의 임bedding 분할은 V100 GPU에서 GPU 메모리 사용량을 최대 26.3% 감소시키고 처리량을 최대 15.6% 향상시켰다.
• CPU-GPU 메모리 링 추론 메커니즘은 라운드로빈 작업 스케줄링을 통해 GPU 메모리 용량을 초과하는 MoE 모델의 효율적 실행을 가능하게 했다.