기술 요약: Piper – 리소스 모델링과 파이프라인 하이브리드 병렬화를 통한 대규모 MoE 학습의 효율성 제고

1. 문제 제기

고성능 컴퓨팅 (HPC) 플랫폼에서 혼합 전문가 (Mixture-of-Experts, MoE) 모델을 학습시키는 것은 대규모에서 밀집 모델의 효율성을 달성하는 것을 방해하는 상당한 장애물에 직면해 있습니다. MoE 아키텍처 (예: Mixtral, DeepSeek, Qwen) 는 희소하게 활성화되는 매개변수를 통해 계산 비용을 줄이면서 우수한 성능을 제공하지만, 그 학습은 다음과 같은 요인들로 인해 저해됩니다:

• 큰 메모리 풋프린트: 모델 매개변수와 일시적 활성화 텐서를 모두 저장하는 것은 상당한 GPU 메모리를 소모합니다.
• 통신 병목 현상: 전문가 병렬화 (Expert Parallelism, EP) 는 이종 네트워크 전반에 걸친 빈번하고 대규모의 all-to-all 통신을 필요로 합니다. Frontier 슈퍼컴퓨터의 Dragonfly 토폴로지와 같은 비균일 인터커넥트에서 이 지연 시간은 지배적인 성능 제한 요인이 됩니다.
• 작업 부하 불균형: 편향된 전문가 할당은 심각한 부하 불균형을 초래하여, 특히 학습 초기 단계나 외부 균형 조정 메커니즘이 부재할 때 GPU 활용도를 낮춥니다.
• 비효율적인 계산 패턴: 세분화된 MoE 아키텍처는 하드웨어 활용도가 낮은 "키가 크고 마른 (tall-and-skinny)" GEMM(일반 행렬 곱셈) 을 생성합니다.
• 종합적 전략의 부재: 기존 프레임워크 (예: DeepSpeed-MoE, X-MoE) 는 종종 메모리, 계산, 통신 제약을 공동으로 최적화하는 플랫폼 인식형 하이브리드 병렬화 전략이 부족합니다.

2. 방법론

저자들은 수학적 리소스 모델링, 경험적 마이크로 벤치마킹, 그리고 플랫폼 인식형 성능 추정을 통해 분산 학습 전략을 공동 설계하는 Piper라는 프레임워크를 제안합니다. 방법론은 네 가지 핵심 구성 요소로 이루어집니다:

A. 분석적 및 경험적 리소스 모델링

Piper 는 다양한 병렬화 방식 (데이터 병렬화, 텐서 병렬화, 파이프라인 병렬화, 전문가 병렬화) 하에서 다양한 MoE 아키텍처에 대한 메모리, 계산, 통신 요구 사항을 정량화하기 위한 수학적 모델을 개발합니다.

• 메모리 모델링: 이 프레임워크는 GPipe 및 1F1B 파이프라인 스케줄에 대한 정적 메모리 (매개변수, 옵티마이저 상태) 와 활성화 메모리를 계산합니다. 여기에는 비전문가 모듈 (어텐션) 의 복제와 $EP \times PP$ 디바이스 메쉬 전반에 걸친 전문가의 분포가 고려됩니다.
• 통신 모델링: 메시지 크기, 네트워크 대역폭, 토폴로지 제약을 기반으로 all-to-all 지연 시간을 추정합니다.
• 검증: 이러한 모델들은 대상 HPC 플랫폼 (Frontier) 의 마이크로 벤치마크를 사용하여 매개변수화되며, 코드 계측 및 하드웨어 프로파일링을 통해 검증됩니다.

B. 파이프라인 하이브리드 병렬화

Piper 는 **파이프라인 병렬화 (PP)**와 **전문가 병렬화 (EP)**의 새로운 조합을 도입합니다.

• 모든 전문가를 대규모 글로벌 그룹에 분산시켜 대규모 all-to-all 을 강제하는 대신, Piper 는 GPU 를 $PP \times EP$ 메쉬로 구성합니다.
• 전문가들은 작고 토폴로지가 로컬인 그룹 (예: 단일 노드 내 또는 단일 홉 스위치 그룹 내) 으로 제한됩니다.
• 이러한 로컬라이제이션은 비용이 많이 드는 집단 통신의 범위를 줄여 Piper 가 빠른 노드 내부 인터커넥트 (NVLink/Infinity Fabric) 를 활용하고 대규모 all-to-all 집단의 높은 지연 시간을 피할 수 있게 합니다.

C. 토폴로지 인식형 All-to-All 알고리즘 (HALO)

Dragonfly 토폴로지에서의 통신 비효율성을 해결하기 위해 Piper 는 HALO(계층적 친화도 인식 로컬리티 최적화 All-to-All) 를 구현합니다.

• 설계: NIC 친화도와 Dragonfly 그룹화 (노드 내부, 그룹 내부, 그룹 간) 를 기반으로 통신 그룹을 구성합니다.
• 3 단계 실행:
  1. 1 단계 (노드 내부): 로컬 데이터의 완전한 독립적 추출.
  2. 2 단계 (노드 간): 원격 노드 간의 배치된 포인트 투 포인트 전송.
  3. 3 단계 (노드 내부): 노드 내 원격 데이터의 재분배.
• 최적화: 1 단계는 지연 시간을 숨기기 위해 2 단계/3 단계와 동시에 시작됩니다. 이 알고리즘은 노드당 네 개의 NIC 를 모두 포화시키고, 가능한 경우 노드 할당을 단일 랙으로 제한하여 느린 랙 간 링크를 피합니다.

D. 동적 전문가 이동

보조 손실 함수에만 의존하지 않고 부하 불균형을 해결하기 위해 Piper 는 전문가 이동 (Expert Migration) 기법을 사용합니다.

• 외부 스케줄러가 토큰 분포를 모니터링합니다. 불균형이 임계값을 초과하면 힐 클라이밍 스왑 알고리즘이 전문가 병렬 그룹 내에서 부하를 재균형화하기 위해 필요한 최소한의 전문가 이동 수를 식별합니다.
• 전문가가 빠른 인터커넥트에 로컬화되어 있으므로, 이동의 오버헤드 (그룹 내 all-to-all) 는 총 학습 시간의 5% 미만으로 분산됩니다.

3. 주요 기여

1. 리소스 모델링 프레임워크: 다양한 병렬화 구성에 걸쳐 메모리, 계산, 통신을 추정하는 통합 분석 모델로, Frontier 슈퍼컴퓨터에서 경험적으로 검증되었습니다.
2. Piper 프레임워크: 레이어 내부 전문가 병렬화에 파이프라인 병렬화를 적용하고, 통신을 로컬라이제이션하며 대상 HPC 플랫폼에 대한 효율적인 구성을 자동으로 식별하는 학습 시스템입니다.
3. HALO 알고리즘: 의존성 구조와 NIC 친화도를 활용하여 벤더 구현체보다 1.5 배~4 배의 대역폭을 달성하는 Dragonfly 토폴로지 인식형 all-to-all 알고리즘입니다.
4. 전문가 이동: 관찰된 부하를 기반으로 전문가를 물리적으로 재분배하여 최소한의 오버헤드로 동적 부하 균형을 맞추는 기법입니다.
5. 조단위 규모 학습: Frontier 슈퍼컴퓨터에서 최첨단 MoE 모델 (최대 1.7 조 개 매개변수) 에 대한 학습 전략을 시연했습니다.

4. 결과

이 논문은 Frontier 슈퍼컴퓨터에서 수행된 실험에 기반한 다음과 같은 성능 지표를 보고합니다:

• 모델 FLOP 활용도 (MFU): Piper 는 X-MoE 와 같은 최첨단 프레임워크에 비해 2~3.5 배 더 높은 MFU를 달성합니다. 참고로, X-MoE 는 5,450 억 개 매개변수 모델에 대해 5.23% MFU 를 보고한 반면, Piper 는 유사하거나 더 큰 모델에 대해 20~50% MFU를 달성합니다.
• 통신 효율성: HALO all-to-all 알고리즘은 벤더 구현체 (RCCL 기반 torch.dist.all_to_all) 보다 1.2 배~9 배의 대역폭을 제공하며, 가장 큰 향상 (최대 9 배) 은 랙 간 통신이 지배적인 16 노드 이상의 규모에서 관찰됩니다.
• 확장성:
  • 1,024 개 GPU 에서 33 TFLOPS 처리량을 기록하며 1.7 조 개 매개변수 MoE 모델을 성공적으로 학습했습니다.
  • 약한 확장성 (weak-scaling) 실험에서 64 개에서 1,024 개 GPU 로 확장할 때 73% 의 확장 효율성을 시연했습니다.
  • 조 단위 매개변수 모델에 대해 20% MFU를 달성하여 이전 벤치마크를 크게 능가했습니다.
• 단일 노드 성능: 단일 Frontier 노드에서 다양한 최첨단 모델의 단일 레이어를 학습했을 때 78~129 TFLOPS 의 처리량이 관찰되어 모델 성능의 상한선을 확립했습니다.

5. 중요성과 주장

이 논문은 Piper가 정적이고 일률적인 병렬화 전략에서 시스템 인식형, 리소스 모델링 기반 접근법으로 전환함으로써 대규모 MoE 모델의 효율적인 학습에서 중요한 진전을 이루었다고 주장합니다.

• 종합적 최적화: 저자들은 메모리, 계산, 통신을 다양한 병렬화 차원에 걸쳐 통합적으로 다루는 통일된 모델을 제공한 이전 연구는 없으며, 이는 실제 HPC 플랫폼에서 검증되었다고 주장합니다. Piper 는 이 격차를 메꿉니다.
• 플랫폼 인식: 이 프레임워크는 대상 HPC 플랫폼 (특히 비균일 인터커넥트와 토폴로지) 에 대한 깊은 이해가 직접적으로 학습 처리량 향상으로 이어진다는 것을 보여줍니다.
• 병목 현상 완화: 파이프라인 병렬화를 통해 전문가 병렬화를 로컬라이제이션함으로써 Piper 는 이전 MoE 학습 노력에서 지배적인 병목 현상이었던 대규모 all-to-all 지연 시간을 효과적으로 완화합니다.
• 실용성: 모듈식 설계는 커널 최적화와 동적 부하 균형 조정 전략의 점진적 통합을 가능하게 하여, emerging MoE 아키텍처에 적응할 수 있게 합니다.

저자들은 그들의 접근 방식이 이전에 보고된 것 (예: 5.23% MFU) 보다 훨씬 높은 효율성 (20~50% MFU) 으로 조 단위 매개변수 모델의 학습을 가능하게 한다고 결론지으며, 엄격한 리소스 모델링과 토폴로지 인식형 알고리즘 설계를 통해 효율적인 대규모 MoE 학습이 달성 가능함을 증명합니다.