Beyond Task-Agnostic: Task-Aware Grouping for Communication-Efficient Multi-Task MoE Inference

기술 요약: 태스크 불가지론을 넘어: 통신 효율적인 멀티태스크 MoE 추론을 위한 태스크 인지형 그룹화

1. 문제 정의

희소 활성화(Sparsely activated) Mixture-of-Experts (MoE) 모델은 조건부 연산을 통해 효율적으로 용량을 확장하지만, 분산 추론 시 심각한 시스템적 과제에 직면합니다. 구체적으로는 과도한 GPU 간 통신 오버헤드와 심각한 부하 불균형(load imbalance) 문제입니다. 기존의 전문가 배치 전략들은 자주 함께 활성화되는 전문가들을 같은 위치에 배치함으로써 이러한 비용을 완화하려고 시도하지만, 일반적으로 태스크 불가지론적(task-agnostic, 전역적으로 집계된) 라우팅 트레이스로부터 단일한 글로벌 배포 계획을 도출한다는 한계가 있습니다.

저자들은 이러한 접근 방식의 결정적인 결함을 지적합니다. 즉, 전문가의 공동 활성화(co-activation)는 강력하게 **태스크 조건부(task-conditioned)**라는 점입니다. 논문의 분석(그림 1)에서 입증되었듯이, 특정 태스크 군(예: 코드 생성)에서 밀접하게 결합된 전문가 쌍은 다른 태스크(예: 수학적 추론)에서는 서로 상관관계가 없거나 심지어 경쟁 관계에 있을 수 있습니다. 이러한 이질적인 패턴을 하나의 글로벌 그래프로 평균화하면 태스크별 전문가 그룹이 붕괴되어, 통신 지역성(communication locality)을 활용하지 못하고 온라인 워크로드의 분포 변화에 대응하지 못하게 됩니다.

2. 방법론

제안된 프레임워크인 **태스크 인지형 공동 활성화 그룹화(Task-Aware Coactivation Grouping, TACG)**와 **범용 전문가 공유 복제(Generic Expert Shared Replication, GESR)**는 다음과 같은 2단계 배포 시점 프로세스를 통해 문제를 해결합니다.

A. 태스크 인지형 공동 활성화 그룹화 (TACG)
TACG는 태스크 조건부 공동 활성화 통계를 용량 제약이 가능한 배치 계획으로 변환합니다:

1. 태스크 군 선호도 스코어링: 캘리브레이션 트레이스를 사용하여 각 태스크 군에 대한 사용 벡터와 공동 활성화 행렬을 계산합니다. 시스템은 각 전문가가 특정 태스크 군에 대해 가지는 "소프트 선호도 점수"($s^{(f)}_{l,e}$)를 계산하고, 이를 정규화하여 전문가의 군 선호도를 나타내는 확률 분포($p^{(f)}_{l,e}$)를 생성합니다.
2. 태스크 변조 공동 활성화 그래프: 통합된 공동 활성화 그래프를 사용하는 대신, TACG는 태스크 변조 그래프($\tilde{A}_l$)를 구축합니다. 이 그래프는 전문가들의 군 선호도 곱($p^{(f)}_{l,e} p^{(f)}_{l,e'}$)에 의해 표준 공동 활성화 엣지의 가중치를 재설정합니다. 이를 통해 공동으로 활성화되면서 동시에 강한 군 선호도를 공유하는 전문가 사이의 엣지는 증폭시키고, 태스크 간 엣지는 약화시킵니다.
3. 용량 제약 기반 그룹화: 프레임워크는 엄격한 용량 제약 하에서 태스크 변조 그래프 상의 그룹 내 친화도를 최대화함으로써 전문가들을 $M$개의 그룹(GPU당 하나)으로 분할합니다. 이는 2단계 스펙트럴 클러스터링(spectral clustering)과 정확한 GPU 부하 제한을 보장하기 위한 로컬 "용량 복구(capacity repair)" 단계를 통해 해결됩니다.

B. 범용 전문가 공유 복제 (GESR)
모든 태스크 군에 걸쳐 활성화되어 핫스팟이 되기 쉬운 "범용(generic)" 전문가들로 인한 부하 왜곡을 처리하기 위해, GESR은 경량화된 동반 메커니즘을 도입합니다:

1. 식별: 글로벌 중심성과 태스크 간 일관성에 보상을 주고, 강한 태스크 특이적 편차에는 페널티를 주는 점수를 사용하여 범용 전문가를 식별합니다.
2. 복제: 이 전문가들을 구조적 호환성(공동 활성화 친화도)에 따라 선택된 소수의 보조 GPU로 복제합니다.
3. 온라인 가드: 서빙 시점에, 지수적으로 감쇠하는 부하 신호를 사용하여 핫스팟 발생을 방지하면서도 MoE 라우터를 변경하지 않고 최적의 후보 GPU를 선택하는 지역성 및 부하 인지형 선택 규칙을 적용합니다.

3. 주요 기여

• 통찰: 본 논문은 멀티태스크 MoE 서빙에서의 전문가 공동 활성화가 강력하게 태스크 조건부임을 입증하며, 효과적인 배치가 글로벌 평균 친화도가 아닌 태스크 인지형 공동 활성화에 의해 주도되어야 한다고 주장합니다.
• TACG: 전문가와 태스크 군 간의 선호도를 사용하여 공동 활성화 그래프의 가중치를 재설정함으로써, 태스크 불가지론적 그룹화보다 엄격하게 일반화된 배치를 생성하는 태스크 인지형 전략입니다.
• GESR: 태스크 간 범용 전문가를 선택적으로 복제하고 부하 인지형 온라인 선택을 수행하여 워크로드 변화 하에서도 균형을 유지하는 경량 복제 모듈입니다.
• 실증적 검증: 세 가지 대표적인 오픈 소스 MoE 모델(DeepSeek-MoE-16B, Qwen1.5-MoE-A2.7B, Moonlight-16B-A3B)에 대한 광범-한 평가를 통해 강력한 베이스라인 대비 일관된 개선을 입증했습니다.

4. 실험 결과

프레임워크는 네 가지 태스크 군(Code, Query, Math, Complex Reasoning)을 가진 세 가지 모델에 대해 16-GPU 전문가 병렬 설정에서 평가되었습니다.

• 통신 감소: 전체 프레임워크(TACG + GESR)는 베이스라인 대비 평균 토큰당 통신 비용을 31.39% 감소시켰으며, 이는 가장 강력한 기존 배치 전략인 ExFlow보다 12.58% 높은 성능입니다.
• 부하 균형: 이 방법은 0.9975의 높은 평균 Jain 공정성 지수와 0.0736의 낮은 평균 최대 부하 위반(MaxVio)을 유지하여, 통신 효율을 크게 높이는 동시에 베이스라인과 대등하거나 더 나은 부하 균형을 달。
• 강건성: 심각한 태스크 분포 변화(OOD 단일 태스크 워크로드) 상황에서도 태스크 인지형 방식은 안정적인 통신 감소를 유지한 반면, 태스크 불가지론적 변형 모델들은 큰 성능 변동과 핫스팟 형성(Code 태스크에서 MaxVio가 약 7배 급증)을 겪었습니다.
• Ablation 연구: 태스크 변조 공동 활성화 신호를 제거했을 때 통신 이득이 31.39%에서 16.51%로 감소하여, 태스크 조건화가 지배적인 요인임을 확인했습니다. 또한 GESR이 통신 효율을 희생하지 않으면서 부하 위반(MaxVio)을 제어하는 데 필수적임을 입증했습니다.

5. 의의 및 주장

본 논문은 MoE 전문가 배치의 패러다임을 정적인 태스크 불가지론적 관점에서 동적인 태스크 인지형 관점으로 전환했다는 점에서 그 의의가 있다고 주장합니다. 다양한 태스크 군에 걸친 전문가 공동 활성화의 이질성을 명시적으로 모델링함으로써, 글로벌 집계 방식에서는 보이지 않는 통신 지역성을 확보할 수 있습니다.

저자들은 이 접근 방식이 통신-부하 균형의 경계선(frontier)에서 엄격한 개선을 제공한다고 단언합니다. 통신 이득을 위해 부하 균형을 희생하거나 그 반대의 경우가 빈번한 기존 방법들과 달리, TACG + GESR은 거의 완벽한 부하 균형을 유지하면서도 상당한 통신 감소를 달성합니다. 이 연구는 효과적인 멀티태스크 서빙을 위해서는 단순히 평균을 내는 것이 아니라, 태스크별 전문가 협업의 고유한 구조를 존중하는 배포 전략이 필요함을 강조합니다.

본 논문은 연구 범위를 겸허히 정의하며, 현재는 정적 오프라인 배치에 집중하고 있으며, 요청 윈도우 내의 동적 재그룹화, 메모리 공동 설계(프리페칭/캐싱), 그리고 도메인 분할 모델의 특화된 전문가에 대한 연구를 향후 과제로 남겨두고 있습니다.