기술 요약: DisagMoE

1. 문제 제기

대규모 전문가 혼합 (Mixture-of-Experts, MoE) 아키텍처는 전문가의 부분 집합을 희소하게 활성화함으로써 조 단위 파라미터를 가진 대규모 언어 모델 (LLM) 의 학습을 가능하게 합니다. 그러나 이러한 모델의 학습은 전문가 병렬화 (Expert Parallelism, EP) 에 필요한 all-to-all 통신으로 인해 심각한 병목 현상에 직면합니다. 모델이 여러 GPU 노드로 확장됨에 따라 제한된 노드 간 네트워크 대역폭 (예: InfiniBand) 이 통신 지연을 악화시켜, 학습 스텝 시간의 최대 50% 를 소모할 수 있습니다.

기존 최적화 전략은 주로 연산자 수준에서 통신과 계산을 겹치게 하거나 (예: FFN 계산의 청킹), 마이크로 배치 파이프라이닝을 사용하는 데 중점을 둡니다. 그러나 이러한 접근법은 두 가지 근본적인 한계를 겪습니다:

1. 잔류 정지 (Residual Stalls): 어텐션과 FFN 레이어 간의 계산 - 통신 비율의 고유한 불균형으로 인해 완벽한 겹침이 불가능합니다. 어텐션 계산은 시퀀스 길이에 따라 이차적으로 증가하는 반면, FFN 과 통신은 선형적으로 증가합니다. 이로 인해 특히 시퀀스 길이가 증가함에 따라 숨길 수 없는 "통신 꼬리 (communication tails)"가 발생합니다.
2. 자원 비효율성: 이전의 분리된 접근법 (어텐션과 FFN 을 서로 다른 하드웨어에 분리) 은 각 그룹이 전체 모델 파라미터를 보유해야 하므로 규모 확장 시 메모리 부족 (OOM) 오류를 초래하거나, 각 구성 요소의 특정 계산 - 통신 특성에 맞춰 자원 할당을 적응적으로 조정하지 못합니다.

2. 방법론: DisagMoE

DisagMoE 는 **분리된 아키텍처 (disaggregated architecture)**를 통해 모델 배치와 스케줄링을 공동으로 최적화함으로써 all-to-all 병목 현상을 완화하도록 설계된 학습 시스템입니다.

2.1 분리된 구성 요소 배치

전통적인 전문가 병렬화에서 어텐션과 FFN 레이어가 동일한 GPU 그룹에 공존하는 것과 달리, DisagMoE 는 구성 요소 유형별로 모델을 분할합니다:

• 어텐션 그룹 (A-Workers): 이 그룹들은 여러 트랜스포머 레이어의 자기 어텐션 (self-attention) 레이어를 포함합니다. 이들은 밀집 구성 요소를 복제하고 데이터 병렬화 (DP) 를 통해 학습 데이터를 분배합니다.
• FFN 그룹 (F-Workers): 이 그룹들은 여러 트랜스포머 레이어의 FFN 레이어 (전문가) 를 포함합니다. 전문가들은 그룹 내 GPU 들 간에 전문가 병렬화 (EP) 를 통해 샤딩됩니다.
• 레이어 할당: 레이어는 그룹에 인터리브된 방식으로 할당됩니다 (예: 그룹 0 은 레이어 0, $p$, $2p$를 처리하고, 그룹 1 은 레이어 1, $p+1$ 등을 처리함). 이는 메모리 제약 하에서 장치 활용도를 극대화합니다.

2.2 AF-Pipe: 겹침을 갖춘 다단계 파이프라인

DisagMoE 는 통신을 부수적 효과가 아닌 1 급 단계로 취급하는 파이프라인 스케줄인 AF-Pipe를 도입합니다.

• Many-to-Many (M2N) 통신: 단일 노드 그룹 내의 표준 all-to-all 교환 대신, AF-Pipe 는 어텐션 그룹과 FFN 그룹 간에 단방향, Many-to-Many 토큰 교환을 확립합니다.
• 단계 정렬: 파이프라인은 그룹 간의 경계를 정렬하여 전통적인 "디스패치 (dispatch)" 및 "컴바인 (combine)" 연산을 단일 통합 통신 단계로 융합합니다. 이는 통신 오버헤드를 약 $1/k$만큼 줄이고 ($k$는 top-$k$ 전문가 수), 중복 데이터 전송을 제거합니다.
• 비동기 실행: 시스템은 각 워커 내에서 세 가지 조정된 스트림 (전진, 역전파, 통신) 을 활용합니다. 이를 통해 FFN 그룹은 사용 가능한 데이터에 대한 전진 통과를 계산하는 동시에 그래디언트 (N2M) 를 전송할 수 있어, 블로킹되지 않는 실행과 거의 연속적인 GPU 활용을 보장합니다.

2.3 적응형 워커 할당

계산에 구애받는 어텐션과 통신에 구애받는 FFN 간의 불균형을 해결하기 위해, DisagMoE 는 **계산 - 통신 루프라인 모델 (Compute-Communication Roofline Model)**에 의해 안내되는 적응형 할당기를 사용합니다.

• 모델링: 시스템은 구성 요소가 네트워크 바운드에서 계산 바운드로 전환되는 전환점 ($\hat{I}$) 을 모델링합니다. 시퀀스 길이가 증가함에 따라 어텐션 레이어는 FFN 레이어보다 빠르게 계산 루프에 도달합니다.
• 최적화: 시스템은 GPU 와 네트워크 인터페이스 카드 (NIC) 분할을 혼합 정수 선형 계획법 (MILP) 문제로 공식화합니다. 이는 FFN 그룹의 통신 루프를 높이기 위해 더 많은 네트워크 대역폭을 할당하고, 어텐션 그룹에는 더 많은 계산 자원을 할당하여 파이프라인 단계를 균형 있게 조정하고 버블을 최소화합니다.
• 검색 전략: 알고리즘은 루프라인 모델을 사용하여 시드 구성을 생성한 후, 모델링되지 않은 오버헤드를 고려하기 위해 측정된 스텝당 시간을 사용하여 로컬 검색을 통해 이를 정제합니다.

3. 주요 기여

1. 분리된 구성 요소 배치: 어텐션과 FFN 레이어를 분리된 워커 그룹으로 분리하여 독립적인 확장과 자원 할당을 가능하게 하는 새로운 아키텍처.
2. AF-Pipe 스케줄: 기존 all-to-all 을 그룹 간 M2N 교환으로 대체하고, 어텐션 및 FFN 계산과 통신을 체계적으로 겹치게 하는 다단계 파이프라인.
3. 적응형 자원 할당: 루프라인 모델에 기반하여 어텐션 및 FFN 그룹에 대한 GPU 및 NIC 자원을 공동으로 크기 조절하는 메커니즘으로, 다양한 시퀀스 길이와 모델 구성에 걸쳐 계산 - 통신 비율을 효과적으로 균형 맞춤.

4. 실험 결과

저자들은 16 노드 (128× H800 GPU) 클러스터에서 DeepSeek-MoE, GPT-OSS-120B, Qwen3 세 가지 MoE 모델을 사용하여 4K 에서 32K 까지의 시퀀스 길이로 DisagMoE 를 평가했습니다.

• 처리량 향상: DisagMoE 는 Megatron-LM(1F1B 인터리브), Tutel, Comet, DualPipe 를 포함한 강력한 베이스라인보다 일관되게 우수한 성능을 발휘했습니다.
  • Megatron-1F1B 대비 최대 1.81 배 속도 향상 달성.
  • 최첨단 MoE 겹침 시스템 대비 최대 1.34 배 속도 향상 달성.
• 통신 감소: 베이스라인과 비교하여 DisagMoE 는 겹치지 않은 통신 시간을 Tutel 대비 최대 88%, Comet 대비 75%, DualPipe 대비 45% 감소시켰습니다.
• 확장성: 시스템은 다양한 top-$k$ 값과 전문가 병렬화 (EP) 크기에서 효과성을 입증했습니다.
• 자원 민감성: 실험은 어텐션과 FFN 그룹 간의 최적 GPU 비율이 시퀀스 길이에 따라 달라짐 (예: 긴 시퀀스는 더 많은 어텐션 자원을 선호) 을 확인하여 적응형 할당의 필요성을 입증했습니다.

5. 의의 및 주장

이 논문은 DisagMoE 가 현재 MoE 학습의 근본적인 한계인 어텐션과 FFN 레이어 간의 구조적 불균형으로 인해 통신을 완벽하게 겹칠 수 없다는 문제를 해결한다고 주장합니다. 이러한 구성 요소를 분리하고 루프라인 모델에 기반하여 자원을 적응적으로 할당함으로써, 시스템은 대칭적인 대역폭 제공이 달성할 수 있는 것보다 더 넓은 유효 겹침 창을 확장합니다.

저자들은 그들의 접근 방식이 고정된 시퀀스 길이를 가진 대규모 사전 학습 작업 부하에 특히 중요하다고 강조합니다. 그들은 현재 설계가 안정적인 형태 (고정된 시퀀스 길이 및 마이크로 배치) 를 가정하고 있지만, 미래의 하드웨어와 소프트웨어가 단순히 원시 대역폭을 확장하는 것이 아니라 구성 요소 간의 유리한 계산 - 통신 비율을 목표로 해야 한다는 핵심 통찰은 MoE 학습을 효율적으로 확장하기 위한 새로운 방향을 제시한다고 지적합니다. 이 시스템은 Megatron-LM 에 구현되었으며, 이전의 분리된 학습 시도들을 괴롭혔던 메모리 제약 없이 수백 조 파라미터 모델을 학습할 수 있는 실현 가능한 솔루션임이 입증되었습니다.