EPLB: MoE 모델에서 GPU들 간의 부하를 분배(Load Balancing)하는 라이브러리 (feat. DeepSeek)

들어가며 :pytorch::kr:

DeepSeek에서 이번 주, 하루에 하나씩, 인프라와 관련한 GitHub 저장소를 공개하기로 하였습니다. 오늘은 DeepSeek-V3 및 DeepSeek-R1 모델 개발 시 사용했던 최적화 전략 및 코드들에 대한 저장소 3개를 한꺼번에 공개했습니다. 그 중 첫번째로 EPLB 저장소를 먼저 살펴보겠습니다. :smiley:

EPLB 소개

DeepSeek-AI에서 오늘 공개한 EPLB(Expert Parallelism Load Balancer)는 MoE (Mixture of Experts) 모델에서 GPU 간 부하를 균형 있게 분배하는 라이브러리입니다. 특히, DeepSeek-V3에서 사용한 redundant experts 전략과 group-limited expert routing 기법을 활용하여 최적의 부하 균형을 유지합니다.

즉, EPLB는 EP(Expert Parallelism) 환경에서 GPU별로 분산 배치한 전문가(Expert)들의 부하를 균형있게 조정하는 알고리즘입니다. 기본적으로 각 Expert는 특정 GPU에 할당되지만, 작업량이 전문가마다 다를 수 있어 GPU 간 부하 불균형이 발생할 수 있습니다.

이를 해결하기 위해 EPLB는 다음과 같은 기법을 사용합니다:

  1. Redundant Experts 전략: 부하가 많은 Expert를 복제하여 여러 GPU에 분산 배치

  2. Group-Limited Expert Routing : 같은 그룹의 Expert들을 최대한 같은 노드에 배치하여 노드 간 데이터 전송을 최소화

이러한 전략을 통해 GPU 간 부하를 균형 있게 유지하고, 효율적인 분산 학습이 가능합니다.

EPLB의 주요 기능

EPLB는 MoE 모델의 부하(Load) 균형을 유지하기 위해 두 가지 로드 밸런싱 전략을 제공합니다:

  • 계층적 부하 균형 (Hierarchical Load Balancing): 계층적 부하 균형 전략은 전문가 그룹의 수가 노드 수의 배수(N배)인 경우 사용할 수 있는 전략입니다. 먼저, Expert 그룹을 각 노드에 균형 있게 배치한 후, 각 노드 내에서 전문가 복제 및 배치하는 방식으로 동작합니다. 이 전략은 사전 워밍업과 같은 작은 크기의 EP(Expert Parallel)에서 활용 가능합니다.

  • 전역 부하 균형 (Global Load Balancing): 전역 부하 균형 전략은 전문가 그룹 수와 노드 수가 배수 관계가 아닐 때 사용할 수 있는 전략입니다. 이는 전문가 그룹과 상관없이 전역적으로 전문가를 복제하여 개별 GPU에 배치하는 방식으로 동작합니다. 이 전략은 디코딩과 같은 큰 EP(Expert Parallel)에서 활용 가능합니다.

라이선스

EPLB 프로젝트는 MIT License로 공개되어 있으며, 자유롭게 수정 및 배포가 가능합니다.

:github: EPLB 라이브러리 GitHub 저장소

더 읽어보기




이 글은 GPT 모델로 정리한 글을 바탕으로 한 것으로, 원문의 내용 또는 의도와 다르게 정리된 내용이 있을 수 있습니다. 관심있는 내용이시라면 원문도 함께 참고해주세요! 읽으시면서 어색하거나 잘못된 내용을 발견하시면 덧글로 알려주시기를 부탁드립니다. :hugs:

:pytorch:파이토치 한국 사용자 모임:kr:이 정리한 이 글이 유용하셨나요? 회원으로 가입하시면 주요 글들을 이메일:love_letter:로 보내드립니다! (기본은 Weekly지만 Daily로 변경도 가능합니다.)

:gift: 아래:arrow_lower_right:쪽에 좋아요:+1:를 눌러주시면 새로운 소식들을 정리하고 공유하는데 힘이 됩니다~ :star_struck: