Scalable Training of Mixture-of-Experts Models with Megatron Core

이 논문은 메모리, 통신, 연산의 통합 최적화 및 다양한 병렬화 기법을 통해 MoE 모델의 확장성 문제를 해결하고, NVIDIA GB200/GB300 클러스터에서 높은 성능을 달성한 Megatron Core 기반의 실용적인 오픈소스 솔루션을 제시합니다.

핵심

거대한 AI 두뇌를 효율적으로 키우는 방법: NVIDIA 의 '메가트론 코어' 이야기

이 논문은 거대하고 복잡한 인공지능 (AI) 모델, 특히 '혼합 전문가 (Mixture-of-Experts, MoE)' 모델을 어떻게 하면 수천 개의 그래픽 카드 (GPU) 를 이용해 빠르고 효율적으로 훈련시킬 수 있는지에 대한 NVIDIA 의 비밀 레시피를 공개합니다.

마치 거대한 도시를 건설하는 것과 비슷합니다. 단순히 건물을 많이 짓는 것 (모델 크기 키우기) 만으로는 부족하고, 그 도시의 교통, 전력, 물 공급 시스템을 어떻게 최적화하느냐가 성패를 좌우합니다.

---

1. 문제: 왜 MoE 는 훈련하기가 너무 어려울까요?

기존의 AI 모델은 모든 데이터를 처리할 때 모든 부품 (파라미터) 을 동시에 작동시켰습니다. 하지만 MoE 는 다릅니다.

• 비유: 기존 모델은 전체 공장이 매일 모든 작업을 하는 것입니다.
• MoE 모델: 거대한 전문가 집단이 있는데, 들어온 작업 (데이터) 에 따라 가장 적합한 전문가 2~3 명만 골라서 일하게 합니다. 나머지 90% 는 쉬고 있죠.

이 방식은 매우 효율적이지만, 시스템 엔지니어에게는 **치명적인 세 가지 장벽 (Three Walls)**을 만들어냅니다.

🧱 장벽 1: 메모리 벽 (Memory Wall) - "창고가 너무 좁다"

• 문제: 비록 한 번에 3 명만 일하지만, 모든 전문가 (수백 명) 의 지식 (파라미터) 은 모두 창고 (메모리) 에 보관되어 있어야 합니다.
• 현실: 창고 크기는 한정되어 있는데, 전문가 수는 수조 개에 달합니다. 모든 전문가의 데이터를 한 번에 다 담을 수 없어 훈련 자체가 불가능해질 수 있습니다.

📡 장벽 2: 통신 벽 (Communication Wall) - "교통 체증이 심하다"

• 문제: 각 작업은 가장 적합한 전문가에게 실시간으로 배정되어야 합니다. 100 개의 GPU 가 있다면, 작업이 각 GPU 로 날아갈 때 모든 GPU 가 서로 데이터를 주고받는 (All-to-All) 혼란이 발생합니다.
• 현실: GPU 가 일을 기다리는 동안, 데이터가 이동하는 데만 시간이 너무 오래 걸립니다. 마치 고속도로에 차가 몰려서 제자리걸음을 하는 것과 같습니다.

⚙️ 장벽 3: 계산 효율 벽 (Compute Efficiency Wall) - "작업이 너무 작아서 비효율적이다"

• 문제: 각 전문가가 처리하는 작업량이 매우 작고 조각조각 나 있습니다.
• 현실: 거대한 공장이지만, 한 번에 처리하는 작업이 너무 작아서 기계 (GPU) 가 빈번히 멈추고 다시 시작해야 합니다. 또한, CPU 가 GPU 에게 "이제 이거 해!"라고 지시하는 횟수가 너무 많아 지시하는 사람 (CPU) 이 바빠서 기계 (GPU) 가 놀고 있는 상황이 발생합니다.

---

2. 해결책: NVIDIA 의 '메가트론 코어'가 쓴 마법

NVIDIA 는 이 세 가지 장벽을 동시에 뚫기 위해 **메가트론 코어 (Megatron Core)**라는 통합 시스템을 개발했습니다.

🔑 해결책 1: '병렬 접기 (Parallel Folding)' - "각자 최적의 길로"

• 기존 방식: 모든 전문가와 일반 부품이 같은 규칙으로 나뉘어 있었습니다. (예: 모든 GPU 가 같은 수의 전문가를 가져야 함)
• 메가트론의 혁신: 주의 (Attention) 부분과 전문가 (MoE) 부분을 완전히 분리했습니다.
  • 비유: 도시의 주거 지역과 산업 단지를 완전히 다르게 설계합니다. 주거지는 넓은 도로 (TP) 가 필요하고, 산업 단지는 많은 공장을 분산 (EP) 시키는 게 좋습니다. 메가트론은 이 두 가지를 각자 최적의 방식으로 배치할 수 있게 해줍니다.

🔑 해결책 2: '메모리 장벽 부수기' - "창고를 똑똑하게 쓰자"

• 기억력 재사용 (Recomputation): 데이터를 다 저장하지 말고, 필요할 때 다시 계산해서 쓰세요. (일시적인 계산 비용 증가로 영구적인 저장 공간 확보)
• 저장 공간 줄이기 (FP8/FP4): 데이터를 **더 작은 숫자 (FP8, FP4)**로 압축해서 저장합니다. (정밀도는 살짝 떨어지지만, 저장 공간은 절반 이상 줄어듭니다.)
• 오프로딩: GPU 메모리가 꽉 차면, CPU 메모리로 잠시 옮겼다 다시 가져옵니다. (PCIe 케이블을 통해 데이터를 빠르게 주고받으며 숨겨진 시간을 활용)

🔑 해결책 3: '통신 장벽 부수기' - "데이터를 숨겨라"

• DeepEP & HybridEP: 데이터를 보내는 전송 프로토콜을 완전히 새로 고쳐 속도를 극대화했습니다.
• 숨기기 (Overlap): 데이터를 보내는 동안 GPU 가 계산을 동시에 하도록 합니다.
  • 비유: 택배 트럭이 도로를 달리는 동안, 공장에서 다음 작업을 미리 준비하는 것입니다. 트럭이 도착할 때쯤 공장은 이미 준비가 끝난 상태입니다.

🔑 해결책 4: '계산 효율 장벽 부수기' - "작업 지시를 자동화하라"

• 그룹화 (Grouped GEMM): 작은 작업들을 하나의 큰 작업으로 묶어서 처리합니다.
• CUDA Graphs: CPU 가 매번 "이거 해, 저거 해"라고 지시하는 대신, **한 번에 모든 지시 사항을 녹음 (Graph)**해두고 재생만 시킵니다.
  • 비유: 지휘자가 매번 악기 하나하나에 지시를 내리는 대신, 악보 전체를 녹음해두고 한 번에 재생하는 것과 같습니다. CPU 의 부담이 사라져 GPU 가 계속 일할 수 있습니다.

---

3. 실제 성과: 얼마나 빨라졌나요?

이 기술들을 적용한 결과, NVIDIA 는 **DeepSeek-V3(6850 억 개 파라미터)**와 Qwen3(2350 억 개 파라미터) 같은 초대규모 모델을 훈련시켰습니다.

• GB300 (최신 칩) 기준: GPU 1 개당 **초당 1,233 테라플롭스 (TFLOPS)**의 속도를 기록했습니다.
• H100 기준: 1,024 개의 GPU 를 연결해 훈련했을 때, 기존 방식보다 훨씬 높은 효율을 보여줍니다.

이는 마치 수천 명의 전문가가 한 팀이 되어, 교통 체증 없이, 창고 부족 없이, 지시 지연 없이 거대한 AI 두뇌를 키우는 것과 같습니다.

---

4. 결론: 왜 이것이 중요한가요?

이 논문은 단순히 "더 빠른 AI"를 만드는 법을 알려주는 것이 아닙니다. 거대 AI 시대의 시스템 공학이 어떻게 진화해야 하는지를 보여줍니다.

• 단순한 확장 (Scaling) 은 불가능합니다: 단순히 GPU 를 더 많이 사서 연결하는 것만으로는 한계가 있습니다.
• 통합적 설계 (Co-design) 가 필수입니다: 메모리, 통신, 계산 능력을 서로 연결된 하나의 시스템으로 보고, 한 부분을 최적화하면 다른 부분에도 영향을 미친다는 점을 고려해야 합니다.

NVIDIA 의 메가트론 코어는 이러한 복잡한 균형을 맞춰, 수조 개의 파라미터를 가진 AI 모델을 실제로 훈련시킬 수 있는 **생산-ready(실제 사용 가능)**한 길을 열었습니다. 이제 연구자와 기업들은 이 기술을 바탕으로 더 똑똑하고, 더 길고, 더 복잡한 AI 를 만들 수 있게 되었습니다.

기술 요약

Megatron Core 를 활용한 확장 가능한 MoE 모델 훈련에 대한 기술 요약

이 보고서는 NVIDIA 가 개발한 오픈소스 프레임워크인 Megatron Core 내에서 혼합 전문가 (Mixture-of-Experts, MoE) 모델을 대규모로 확장하여 훈련하기 위한 체계적인 접근법과 최적화 기법을 다룹니다. 특히 DeepSeek-V3(685B 파라미터) 및 Qwen3(235B) 와 같은 최신 MoE 아키텍처를 GB300, GB200, H100 등 최신 GPU 클러스터에서 훈련할 때 발생하는 시스템적 병목 현상을 해결하는 방법을 제시합니다.

---

1. 문제 정의 (Problem)

MoE 모델은 희소성 (Sparsity) 을 통해 토큰당 계산량을 줄이면서 모델 용량을 극대화할 수 있지만, 이는 밀집 (Dense) 모델 훈련과는 다른 시스템적 도전을 야기합니다. Megatron Core 는 이를 **세 가지 주요 장벽 (Three Walls)**으로 정의합니다.

1. 메모리 장벽 (Memory Wall):
   • MoE 는 전체 파라미터 수는 많지만, 각 토큰이 활성화하는 전문가 (Expert) 는 소수입니다. 이로 인해 전체 파라미터와 활성화 (Activation) 메모리가 급증합니다.
   • 동적 라우팅으로 인한 부하 불균형 (Load Imbalance) 은 예측 불가능한 메모리 피크를 유발하며, 기존 Dense 모델용 최적화만으로는 GPU 메모리 용량을 초과하기 쉽습니다.
2. 통신 장벽 (Communication Wall):
   • 전문가 병렬화 (Expert Parallelism, EP) 는 모든 토큰을 해당 전문가가 있는 GPU 로 라우팅하기 위해 All-to-All 통신을 필요로 합니다.
   • EP 규모가 커질수록 노드 간 (Inter-node) 통신이 증가하며, 이는 대역폭이 낮은 링크를 통해 수행되어 훈련 시간을 60% 이상 차지할 수 있습니다.
3. 계산 효율성 장벽 (Compute Efficiency Wall):
   • 미세한 전문가 (Fine-grained experts) 는 작은 GEMM 연산을 생성하여 GPU 연산 유닛 (Tensor Cores) 을 충분히 활용하지 못합니다.
   • 라우팅, 페러미테이션 (Permutation), 동적 텐서 모양 (Dynamic Shapes) 으로 인한 많은 커널 런칭 (Kernel Launch) 과 호스트 (CPU) - 디바이스 동기화는 GPU 가 유휴 상태가 되는 원인이 됩니다.

또한, Attention 레이어와 MoE 레이어는 서로 다른 최적의 병렬화 전략 (Attention 은 TP/CP 선호, MoE 는 EP 선호) 을 요구하는데, 기존 프레임워크는 이를 하나의 구성으로 강제하여 성능 저하를 초래했습니다.

---

2. 방법론 및 핵심 기여 (Methodology & Key Contributions)

Megatron Core 는 위 세 가지 장벽을 동시에 해결하기 위해 통합된 최적화 스택을 제공합니다.

2.1. 다차원 병렬화 및 Parallel Folding

• 문제 해결: Attention 과 MoE 레이어의 병렬화 요구사항 충돌을 해결합니다.
• 해법 (Parallel Folding): Attention 과 MoE 레이어의 병렬화 매핑을 분리하여 Decoupled Parallelism을 가능하게 합니다.
  • Attention 레이어: 높은 Tensor Parallelism (TP) 및 Context Parallelism (CP) 사용.
  • MoE 레이어: 높은 Expert Parallelism (EP) 사용 (Expert TP 는 1 로 고정하여 GEMM 효율성 극대화).
  • 기존 EP ≤ DP 제약을 깨고, EP 를 TP/CP 그룹을 가로질러 '접을 (Fold)' 수 있게 하여 최소 GPU 수를 줄이고 유연성을 높였습니다.

2.2. 메모리 장벽 극복

• 메모리 효율적 페러미테이션 (Memory-Efficient Permutation): 라우팅 가중치를 활성화 계산 전에 적용하여 역전파 시 저장해야 하는 중간 텐서를 제거합니다 (0 오버헤드).
• 정밀도 인식 최적화 (Precision-Aware Optimizer): Adam 옵티마이저의 모멘텀/분산 상태를 BF16/FP8 로 저장하여 메모리를 50% 절감합니다.
• 세분화된 재계산 (Fine-grained Recomputation) 및 오프로딩: 메모리 소모가 큰 모듈 (LayerNorm, MLA 등) 에만 재계산을 적용하고, 나머지는 CPU 메모리로 오프로딩하여 GPU 메모리 부담을 줄입니다.
• FSDP for MoE: 전문가 데이터 병렬 (EDP) 그룹 내에서 파라미터, 그래디언트, 옵티마이저 상태를 분할 (Sharding) 하는 맞춤형 FSDP 를 구현했습니다.

2.3. 통신 장벽 극복

• 최적화된 디스패처 (DeepEP & HybridEP): NCCL 기반의 표준 All-to-All 대신, NVLink 및 RDMA 하드웨어 원시 (TMA, IBGDA) 를 활용한 커널을 사용하여 대역폭을 극대화하고 지연 시간을 줄입니다.
• 지연 시간 숨기기 (Communication Overlapping): 1F1B 스케줄링과 결합하여 All-to-All 통신을 인접한 마이크로배치 (Micro-batch) 의 계산 (Forward/Backward) 과 겹쳐서 실행합니다. 이를 통해 통신 오버헤드를 5% 미만으로 줄였습니다.

2.4. 계산 효율성 장벽 극복

• Grouped GEMM: 여러 전문가의 연산을 하나의 커널로 배치하여 GPU 활용도를 높입니다.
• CUDA Graphs 및 Sync-Free Execution:
  • CUDA Graphs: 호스트 오버헤드를 제거하기 위해 커널 시퀀스를 그래프로 캡처합니다.
  • Sync-Free Kernels: Dropless MoE(토큰 드롭 없는 MoE) 의 동적 모양 문제를 해결하기 위해, 호스트가 모양을 알 필요 없이 GPU 가 직접 모양 정보를 읽고 커널을 실행하도록 설계했습니다.
  • ECHO (Elastic Cloning for Hot Experts): 인기 있는 전문가를 클론하여 부하 불균형을 해소하고 메모리 단편화를 줄입니다.
  • Paged Stashing: CUDA Graph 내에서의 메모리 단편화를 해결하기 위해 페이지 기반 메모리 관리 기법을 도입했습니다.

2.5. 저정밀도 훈련 (FP8/FP4)

• 선택적 정밀도 (Selective Precision): 라우터, 임베딩, 옵티마이저 상태는 고정밀도 (FP32/BF16) 를 유지하고, 전문가 GEMM 과 활성화는 FP8/FP4 로 양자화합니다.
• NVFP4 및 MXFP8: Blackwell 및 Hopper 아키텍처에 최적화된 양자화 레시피를 제공하여 메모리, 통신, 계산 효율성을 동시에 개선합니다.

2.6. 긴 컨텍스트 및 RL 지원

• 긴 컨텍스트: Attention 이 계산의 주를 이루는 긴 시퀀스 (16K~64K+) 를 위해 Context Parallelism (CP) 과 Tensor Parallelism (TP) 을 결합하고, Dynamic Context Parallelism을 통해 가변 길이 시퀀스에 맞춰 CP 크기를 동적으로 조정합니다.
• 강화학습 (RL): 가변 길이 시퀀스 처리 (Packed Sequences), 온라인 가중치 내보내기, 라우터 재생 (Router Replay) 등을 지원하여 RL 파인튜닝 워크로드를 최적화합니다.

---

3. 성능 결과 (Results)

Megatron Core MoE 스택은 다양한 하드웨어 플랫폼에서 상태-of-the-art(SOTA) 성능을 입증했습니다.

모델	하드웨어	GPU 수	정밀도	성능 (TFLOPS/GPU)	토큰/초/GPU
DeepSeek-V3 (685B)	GB300 (Blackwell)	256	MXFP8	1,233	4,730
DeepSeek-V3	GB200 (Blackwell)	256	MXFP8	1,048	4,020
DeepSeek-V3	H100 (Hopper)	1,024	FP8-Blockwise	368	1,412
Qwen3 (235B)	GB300	256	MXFP8	974	6,583
Qwen3	GB200	256	MXFP8	919	6,212

• GB300/GB200 vs H100: GB300/GB200 플랫폼은 H100 대비 약 3 배 높은 토큰 처리량을 보여주며, 이는 메모리 대역폭 향상과 MXFP8 네이티브 지원 덕분입니다.
• 긴 컨텍스트: 131K 시퀀스 길이에서 Qwen3 를 훈련할 때에도 1,150 TFLOPS/GPU 의 높은 효율성을 유지했습니다.

---

4. 의의 및 결론 (Significance)

이 보고서는 MoE 모델의 대규모 훈련이 단순히 모델 크기를 늘리는 문제가 아니라, 메모리, 통신, 계산 효율성이라는 세 가지 시스템 장벽을 통합적으로 해결해야 하는 과제임을 강조합니다.

1. 실용적인 확장성: Megatron Core 는 수천 개의 GPU 에서 조밀한 모델부터 조밀한 MoE 모델까지, 수조 개의 파라미터를 가진 모델을 훈련할 수 있는 생산 준비 (Production-ready) 된 오픈소스 솔루션을 제공합니다.
2. 하드웨어 최적화: Blackwell (GB300/GB200) 과 Hopper (H100) 아키텍처의 특성을 최대한 활용하는 소프트웨어 스택을 제공하여 하드웨어 투자 대비 효율을 극대화합니다.
3. 유연한 아키텍처 지원: Parallel Folding 을 통해 Attention 과 MoE 레이어를 독립적으로 최적화할 수 있게 하여, DeepSeek-V3, Qwen3 등 다양한 최신 MoE 아키텍처에 적용 가능합니다.
4. RL 및 긴 컨텍스트 대응: 강화학습 (RL) 과 긴 컨텍스트 (Long-context) 훈련과 같은 최신 트렌드에 맞춰 동적 병렬화 및 메모리 관리 기법을 도입했습니다.

결론적으로, Megatron Core MoE 는 MoE 모델의 확장성을 제한하던 시스템적 한계를 극복하고, 차세대 대규모 언어 모델 훈련의 새로운 표준을 제시합니다.