Mozart: Modularized and Efficient MoE Training on 3.5D Wafer-Scale Chiplet Architectures

이 논문은 인간 뇌의 모듈화 구조에서 영감을 받아 3.5D 웨이퍼 스케일 칩릿 아키텍처에서 MoE 기반 대규모 언어 모델의 효율적인 훈련을 가능하게 하는 알고리즘 - 하드웨어 공동 설계 프레임워크인 '모차르트 (Mozart)'를 제안합니다.

핵심

모차르트 (Mozart): 거대한 AI 두뇌를 위한 '초고속 칩' 설계도

이 논문은 거대한 인공지능 (LLM) 을 더 빠르고 효율적으로 훈련시키기 위해, 인간의 뇌 구조에서 영감을 받아 새로운 컴퓨터 칩과 소프트웨어를 함께 설계한 이야기를 담고 있습니다.

이 내용을 일반인이 이해하기 쉽게 비유를 들어 설명해 드리겠습니다.

---

1. 문제: 거대한 AI 의 '교통 체증'

지금까지의 AI 는 거대한 도서관처럼 모든 지식을 한곳에 저장하고 처리하는 방식이었습니다. 하지만 최신 AI 는 '모듈형 (MoE)' 방식을 사용합니다. 마치 거대한 병원을 상상해 보세요.

• 일반적인 AI: 모든 환자가 한 명의 의사에게만 가서 진료를 받습니다. (비효율적, 병목 현상 발생)
• 모듈형 AI (MoE): 환자가 들어오면, '심장 전문의', '신경과 전문의', '피부과 전문의' 등 **상황에 맞는 특정 의사 (전문가)**만 호출합니다.

하지만 여기서 문제가 생깁니다.
병원 전체가 너무 넓고, 의사들이 서로 다른 건물 (칩) 에 흩어져 있어서, 환자가 적절한 의사를 찾아가는 동안 시간이 너무 많이 걸리고, 의사가 환자를 만나기 위해 이동하는 통행료 (데이터 전송 비용) 가 너무 비쌉니다. 또한, 어떤 전문의는 바쁘고 어떤 전문의는 한가해서 자원 낭비가 심합니다.

2. 해결책: '모차르트 (Mozart)' 프로젝트

저자들은 이 문제를 해결하기 위해 인간의 뇌를 벤치마킹했습니다. 인간의 뇌는 특정 기능을 담당하는 부위들이 서로 가까이 배치되어 있어, 정보 전달이 매우 빠르고 효율적입니다.

저자들은 이 원리를 적용하여 **'모차르트'**라는 새로운 시스템을 만들었습니다. 이는 **소프트웨어 (알고리즘)**와 **하드웨어 (칩)**를 함께 설계한 '알고리즘 - 하드웨어 공동 설계'입니다.

🎻 비유: 모차르트의 오케스트라

이 시스템을 거대한 오케스트라에 비유해 볼까요?

1. 스마트한 악기 배치 (전문가 배치 전략)

   • 기존 방식: 바이올린, 트럼펫, 드럼이 모두 무대 끝에서 무작위로 흩어져 있어, 악보가 오가느라 시간이 걸립니다.
   • 모차르트 방식: 함께 연주해야 하는 악기들 (함께 활성화되는 전문가들) 을 무대 중앙에 가까이 배치합니다. 예를 들어, 바이올린과 비올라가 자주 함께 연주된다면, 이들을 같은 무대 구역에 둡니다. 이렇게 하면 악보 (데이터) 가 이동할 필요가 줄어들어 속도가 빨라집니다.

2. 흐르는 물처럼 흐르는 데이터 (세밀한 스케줄링)

   • 기존 방식: 악보가 다 도착할 때까지 기다렸다가 연주를 시작합니다. (기다리는 시간이 길어짐)
   • 모차르트 방식: 연주 (계산) 와 악보 전달 (데이터 전송) 을 동시에 진행합니다. 한 악기가 연주하는 동안, 다음 악기를 위한 악보는 이미 다음 무대로 흘러가고 있습니다. 이렇게 하면 '기다리는 시간'을 완전히 없앨 수 있습니다.

3. 초고속 3.5D 칩 (새로운 무대 구조)

   • 기존 방식: 악기들이 2 차원 평면 위에 넓게 퍼져 있어, 소리가 전달되는 데 시간이 걸립니다.
   • 모차르트 방식: 칩을 수직으로 여러 층 쌓아 올린 (3.5D) 구조로 만들었습니다. 마치 고층 빌딩처럼, 계산하는 층과 기억하는 층이 바로 위에 바로 붙어 있습니다. 데이터가 이동할 거리가 매우 짧아져 초고속으로 처리됩니다.

3. 결과: 얼마나 빨라졌을까?

이 새로운 '모차르트' 시스템을 테스트한 결과, 기존 방식보다 약 2 배 (1.9 배~2.4 배) 더 빠른 속도를 보여주었습니다.

• 핵심: AI 가 더 큰 지능을 가지면서도, 전기 요금과 시간을 아낄 수 있게 되었습니다.

4. 요약: 왜 이것이 중요한가?

지금까지 AI 는 더 커질수록 컴퓨터가 감당하기 힘들어졌습니다. 하지만 모차르트는 마치 인간의 뇌처럼 모듈화되고 효율적으로 설계되어, 거대한 AI 모델을 훈련시키는 데 드는 비용과 시간을 획기적으로 줄여줍니다.

한 줄 요약:

"거대한 AI 병원을, 함께 일하는 의사들을 가까이 배치하고, 데이터가 흐르는 대로 즉시 처리되도록 설계한 초고속 3D 빌딩을 지어, AI 훈련을 2 배나 빠르게 만든 혁신적인 프로젝트입니다."

기술 요약

1. 연구 배경 및 문제 정의 (Problem)

• MoE (Mixture-of-Experts) 아키텍처의 한계: 대규모 언어 모델 (LLM) 에서 MoE 아키텍처는 모듈화된 계산을 통해 효율성을 높였으나, 내재된 희소성 (sparsity) 으로 인해 하드웨어 배포 시 심각한 문제를 야기합니다.
  • 메모리 지역성 (Memory Locality) 문제: 동적이고 불균형한 워크로드로 인해 메모리 접근 패턴이 비효율적입니다.
  • 통신 오버헤드: 전문가 (Expert) 간 데이터 이동 (All-to-All 통신) 이 빈번하여 병목 현상이 발생합니다.
  • 자원 활용도 저하: 고정된 하드웨어 구조가 MoE 의 동적 활성화 패턴을 지원하지 못해 연산 자원이 낭비됩니다.
• 기존 하드웨어의 부적합성: 기존 2.5D/3D 칩렛 아키텍처나 GPU 기반 시스템은 주로 조밀하고 균일한 계산을 가정하며, 와퍼 스케일 (Wafer-Scale) 통합을 고려하지 않거나, MoE 의 미세한 모듈성 (Fine-grained modularity) 을 반영한 시스템 수준의 최적화가 부족합니다.

2. 제안 방법론: Mozart (Methodology)

Mozart 는 MoE 기반 LLM 의 효율적인 학습을 위해 알고리즘 - 하드웨어 공동 설계 (Algorithm-Hardware Co-design) 프레임워크를 제안합니다.

A. 알고리즘 측면 (Algorithm Side)

1. 전문가 배치 전략 (Expert Allocation Strategy):
   • 사전 학습된 모델의 라우팅 패턴 (활성화 빈도 및 공동 활성화 패턴) 을 분석하여 전문가를 클러스터링합니다.
   • 자주 함께 활성화되는 전문가들을 동일한 칩렛 또는 인접 칩렛에 배치하여, 칩렛 간 All-to-All 통신 데이터 양을 최소화하고 워크로드 균형을 맞춥니다.
   • 이를 위해 2 단계 접근법 (전문가 클러스터링 및 클러스터 할당) 을 사용합니다.
2. 세분화된 스케줄링 (Fine-grained Scheduling):
   • 스트리밍 토큰 및 전문가: DRAM 과 칩렛 간의 통신 오버헤드를 온칩 연산과 겹쳐서 (Overlap) 처리합니다.
   • 스트리밍 전문가: 워크로드가 많은 전문가 클러스터를 먼저 로드하여 통신 - 연산 중첩을 최적화합니다.
   • 스트리밍 토큰: 글로벌 토큰 배치 (Batch) 를 마이크로 배치로 분할하여, Attention 연산과 MoE 연산을 파이프라인 방식으로 동시에 실행합니다.

B. 하드웨어 측면 (Architecture Side)

1. 3.5D 웨이퍼 스케일 칩렛 아키텍처:
   • 2.5D NoP-Tree 토폴로지: Attention 칩렛을 중앙 허브 (Dispatcher) 로, 전문가 칩렛을 리프 (Leaf) 노드로 구성합니다. 스위치를 통해 네트워크 내 MoE 집계 (In-network aggregation) 를 수행하여 통신 지연과 대역폭 비용을 줄입니다.
   • 계층적 메모리 구조:
     • 3D Logic-on-Memory: 각 연산 칩렛에 SRAM 다이를 수직 적층 (Hybrid Bonding) 하여 중간 결과 (활성화 데이터) 를 로컬에 캐싱하고 접근 지연을 최소화합니다.
     • 분산 DRAM: 모델 가중치는 웨이퍼 주변에 분산된 DRAM 에 저장합니다.
2. 알고리즘 - 하드웨어 매핑:
   • Attention 모듈과 MoE 전문가 모듈을 각각 최적화된 칩렛에 매핑하고, 데이터 흐름을 조정하여 메모리 병목과 연산 병목을 해결합니다.

3. 주요 기여 (Key Contributions)

1. Mozart 프레임워크 개발: MoE 기반 LLM 의 효율적인 학습을 위한 포괄적인 알고리즘 - 하드웨어 공동 설계 프레임워크를 제시했습니다.
2. 최적화 기법 도입:
   • 전문가 클러스터링 및 할당을 통한 워크로드 균형 및 통신 데이터 감소.
   • 스트리밍 기반의 세분화된 스케줄링을 통한 통신 - 연산 중첩 (Overlap) 극대화.
3. 새로운 하드웨어 아키텍처 설계: 3.5D 통합 기술과 2.5D NoP-Tree 토폴로지를 결합한, MoE 의 희소성과 모듈성에 최적화된 웨이퍼 스케일 칩렛 아키텍처를 설계했습니다.
4. 성능 검증: 다양한 규모의 MoE 모델 (Qwen3, OLMoE, DeepSeek-MoE) 에서 기존 방법 대비 1.9 배 이상의 가속화를 달성함을 입증했습니다.

4. 실험 결과 (Results)

• 실험 설정: Qwen3-30B-A3B, OLMoE-1B-7B-0924, DeepSeek-MoE-16B-Base 세 가지 모델을 대상으로 NVIDIA A100 기반 시뮬레이션 및 Verilog 기반 하드웨어 평가를 수행했습니다.
• 성능 향상:
  • Mozart-C (최적화 기법 모두 적용) 는 베이스라인 대비 Qwen3-30B-A3B 에서 1.92 배, OLMoE-1B-7B-0924 에서 2.37 배, DeepSeek-MoE-16B-Base 에서 2.17 배의 학습 속도 향상을 보였습니다.
• 세부 분석:
  • 통신 - 연산 중첩의 효과: DRAM 대역폭 (HBM2 vs SSD) 이 낮을 때 성능 향상이 제한적이었으나, HBM2 환경에서는 Mozart 의 최적화 기법이 전체적인 지연 시간을 크게 단축시켰습니다.
  • 시퀀스 길이 영향: 시퀀스 길이가 길어질수록 통신 병목이 심화되지만, Mozart 는 이를 완화하여 긴 시퀀스에서도 낮은 지연 시간을 유지했습니다.
  • 병목 분석: Mozart 시스템은 현재 연산 병목이 아닌 메모리 (가중치 로드) 병목 상태임을 확인했습니다.

5. 의의 및 결론 (Significance)

• 모듈형 LLM 의 확장성: MoE 아키텍처의 모듈성이 하드웨어의 물리적 모듈성 (칩렛) 과 자연스럽게 매칭될 수 있음을 보여주었습니다.
• 차세대 하드웨어 패러다임: 기존 조밀한 연산 중심의 하드웨어 설계에서 벗어나, 희소하고 동적인 워크로드를 처리하기 위한 3.5D 웨이퍼 스케일 아키텍처의 가능성을 제시했습니다.
• 실용적 가치: 대규모 MoE-LLM 의 사후 학습 (Post-training) 및 배포 단계에서 자원 활용도와 병렬화 효율을 극대화하여, 더 크고 효율적인 모델 개발을 가능하게 합니다.

이 논문은 MoE 의 알고리즘적 특성과 차세대 칩렛 하드웨어 아키텍처를 깊이 있게 결합하여, 대규모 언어 모델 학습의 효율성 문제를 해결하는 새로운 방향성을 제시했다는 점에서 의의가 큽니다.