MoE 학습용 토폴로지 인식 Expert Replica 배치

이 요약은 AI가 원문을 분석해 생성했습니다. 정확한 내용은 원문 기준으로 확인하세요.

TL;DR

MoE의 동적 라우팅은 일부 Expert에 토큰을 몰리게 해 분산 학습의 Rank별 부하 불균형을 만들며, 기존 Replica 방식은 이를 완화하는 대신 다중 노드에서 발생하는 Expert weight 이동 비용을 놓쳤습니다. TAOT는 overload와 여유 용량을 통신비용 행렬이 포함된 엔트로피 정규화 Optimal Transport 문제로 구성하고 Sinkhorn-Knopp 반복으로 Rank 간 흐름을 계산한 뒤, 정수 Replica 매칭과 토큰 할당을 실행 스케줄로 연결합니다. guest-weight transfer를 home-expert computation과 중첩한 결과 end-to-end MoE training 속도가 1.43배 향상됐고, weighted expert-communication cost는 최대 74% 감소했습니다. 부하 균형뿐 아니라 네트워크 토폴로지와 가중치 이동을 함께 최적화해야 MoE 학습의 실제 처리량을 높일 수 있다는 결과입니다.

섹션별 상세

MoE는 LLMs의 규모를 확장하는 데 쓰이지만 동적 라우팅이 일부 Expert와 Rank에 토큰을 집중시켜 Expert-parallel training의 부하 불균형을 키웁니다. 기존 동적 Replica 방식은 hot Expert를 유휴 Rank에 복사해 계산을 분산했지만, 부하 균형만 최적화하고 여러 노드 사이에서 Expert weight를 옮기는 비용은 충분히 고려하지 않았습니다. 그 결과 통신량이 균형 개선 효과보다 커져 전체 학습 비용이 증가할 수 있다는 문제가 TAOT의 출발점입니다.

TAOT는 hot Rank의 overload와 가벼운 Rank의 spare capacity를 통신비용 행렬을 포함한 균형 엔트로피 정규화 Optimal Transport 문제로 바꿉니다. Sinkhorn-Knopp 반복이 Rank 수준의 flow hint를 계산하면, TAOT는 이를 정수 Replica 매칭과 토큰 할당으로 결합해 실제 실행 가능한 스케줄을 만듭니다. 이 과정은 단순히 Expert를 복제할 위치가 아니라 가중치 이동 경로와 계산 분산을 함께 결정해 통신 비용이 낮은 배치를 찾도록 구성됩니다.

시스템 수준에서는 guest-weight transfer를 home-expert computation과 겹쳐 실행해 가중치 전송에 걸리는 시간을 계산 작업 뒤에 숨깁니다. 실험에서 TAOT는 end-to-end MoE training을 1.43배 빠르게 했고, 부하 균형 품질은 기존 최첨단 방법과 비슷하거나 더 나은 수준에 도달했습니다. 모든 구성에서 weighted expert-communication cost가 가장 낮았으며 최대 74% 감소해, Replica 배치에서 통신 토폴로지를 함께 최적화해야 하는 이유를 뒷받침합니다.

용어 해설

Mixture-of-Experts: — 여러 Expert 중 일부만 선택해 토큰을 처리하는 모델 구조입니다. 라우터가 토큰을 Expert에 분배하므로 전체 파라미터를 유지하면서 계산량을 줄일 수 있지만, 특정 Expert에 요청이 몰리면 Expert-parallel training에서 Rank별 부하 불균형과 통신 병목이 발생합니다.
최적 수송(Optimal Transport): — 공급 지점의 초과 부하를 수요 지점의 여유 용량으로 옮기는 비용 최소화 기법입니다. TAOT는 과부하 Rank와 유휴 용량을 각각 수송 문제의 공급·수요로 표현하고, Rank 사이의 Expert weight 이동 비용을 함께 최적화합니다.
Sinkhorn-Knopp 알고리즘(Sinkhorn-Knopp): — 엔트로피 정규화가 적용된 Optimal Transport 문제의 수송 행렬을 반복적으로 계산하는 알고리즘입니다. TAOT는 Sinkhorn-Knopp 반복으로 Rank 단위의 데이터 이동 힌트를 얻은 뒤, 정수 Replica 매칭과 토큰 할당을 실행 가능한 스케줄로 변환합니다.
Expert 복제본(Expert Replica): — MoE training에서 처리 요청이 몰린 hot Expert의 가중치를 다른 Rank에 복사한 실행 인스턴스입니다. 유휴 Rank가 복제본의 계산을 나누어 맡으면 부하를 완화할 수 있지만, 다중 노드 환경에서는 가중치 이동 통신 비용이 성능 이득을 상쇄할 수 있습니다.
엔트로피 정규화(Entropy Regularization): — Optimal Transport 해가 특정 이동 경로에만 집중되지 않도록 목적 함수에 엔트로피 항을 추가하는 방식입니다. TAOT는 이를 포함한 균형 수송 문제를 사용해 부하 초과량, 여유 용량, Rank 간 통신 비용을 함께 반영한 흐름을 계산합니다.

기술

Mixture-of-Experts (MoE)
LLMs
Optimal Transport
Entropy Regularization
Sinkhorn-Knopp
Expert Replica
Expert-parallel training

활용 사례

다중 노드 환경의 MoE 학습
동적 Expert Replica 배치
hot Expert의 계산 부하 분산
Expert weight 전송과 계산의 중첩 실행

언급된 리소스

논문TAOT: Topology-Aware Optimal Transport for Dynamic Expert Replica Placement in MoE Training