GitHub - deepseek-ai/EPLB: Expert Parallelism Load Balancer

Expert Parallelism Load Balancer (EPLB)는 Expert Parallelism (EP) 환경에서 GPU 간의 부하 불균형 문제를 해결하기 위한 알고리즘을 제공합니다. EP에서는 서로 다른 전문가(experts)가 다른 GPU에 할당되며, 전문가들의 현재 워크로드에 따라 부하가 달라질 수 있어 GPU 간의 부하 균형을 맞추는 것이 중요합니다. 이 저장소는 DeepSeek-V3 논문에서 설명된 배포된 EP 부하 분산 알고리즘을 오픈 소스로 공개합니다.

핵심 방법론 (Core Methodology)

EPLB는 추정된 전문가 부하를 기반으로 균형 잡힌 전문가 복제 및 배치 계획을 계산합니다. 전문가 부하를 예측하는 구체적인 방법(예: 과거 통계의 이동 평균 사용)은 이 저장소의 범위 밖에 있습니다.

중복 전문가 전략 (Redundant Experts Strategy): 부하가 높은 전문가를 복제(duplicate)하여 총 전문가 복제본 수를 늘립니다. 이는 전체 GPU에 걸쳐 부하를 분산시키기 위한 기본 전략입니다.

휴리스틱 패킹 (Heuristic Packing): 복제된 전문가들을 휴리스틱하게 GPU에 할당하여 GPU 간의 부하 균형을 보장합니다.

그룹 한정 전문가 라우팅 (Group-limited Expert Routing) 고려: DeepSeek-V3에서 사용되는 group-limited expert routing의 특성을 활용하여, 가능한 경우 동일한 그룹에 속하는 전문가들을 동일한 노드(node)에 배치하여 노드 간(inter-node) 데이터 트래픽을 줄이려 시도합니다.

알고리즘 (The Algorithm)

부하 분산 알고리즘은 두 가지 정책을 사용하며, 각기 다른 상황에 적용됩니다.

계층적 부하 분산 (Hierarchical Load Balancing):

* 적용 조건: 서버 노드(server nodes) 수가 전문가 그룹(expert groups) 수를 나눌 수 있는 경우 (즉, $num_expert_groups$인 경우). 이는 group-limited expert routing의 이점을 최대한 활용하기 위함입니다.
* 동작 방식:
* 1단계: 그룹을 노드에 패킹: 먼저 전문가 그룹들을 노드에 균등하게 패킹하여 다른 노드의 부하가 균형을 이루도록 합니다. 이는 노드 간 트래픽을 최소화하는 데 기여합니다.
* 2단계: 노드 내 전문가 복제: 각 노드 내에서 전문가들을 복제합니다. 이 단계에서 복제 예산(replication budget)이 각 노드에 할당됩니다.
* 3단계: 복제된 전문가를 GPU에 패킹: 마지막으로 복제된 전문가들을 개별 GPU에 패킹하여 다른 GPU의 부하가 균형을 이루도록 합니다. 이 과정은 각 노드 내에서 이루어집니다.
* 사용 사례: 일반적으로 더 작은 expert-parallel size를 사용하는 prefilling 단계에서 활용됩니다.

전역 부하 분산 (Global Load Balancing):

* 적용 조건: 계층적 부하 분산 정책을 사용할 수 없는 경우, 즉 서버 노드 수가 전문가 그룹 수를 나누지 못하는 경우.
* 동작 방식: 전문가 그룹의 제약 없이 전역적으로 전문가를 복제하고, 복제된 전문가들을 개별 GPU에 패킹합니다. 이 정책은 모든 GPU를 대상으로 전체 복제본을 직접 할당합니다.
* 사용 사례: 일반적으로 더 큰 expert-parallel size를 사용하는 decoding 단계에서 활용됩니다.

인터페이스 및 예시 (Interface and Example)

부하 분산기의 주 함수는 eplb.rebalance_experts 입니다. 이 함수는 다음의 입력 파라미터를 받습니다:
* weight: 추정된 전문가 부하를 나타내는 텐서.
* num_replicas: 총 전문가 복제본 수.
* num_groups: 전문가 그룹 수.
* num_nodes: 서버 노드 수.
* num_gpus: 총 GPU 수.

함수는 전문가 복제 및 배치 계획을 나타내는 phy2log, log2phy, logcnt를 반환합니다. phy2log는 물리적 GPU 인덱스에 매핑되는 논리적 전문가 ID를 포함하는 텐서입니다.

이 코드는 MIT License 하에 배포됩니다.