Compute-Optimal Is Not Cluster-Optimal: Systems-Aware Scaling for Sparse Mixture-of-Experts

기술 요약: 컴퓨트 최적화는 클러스터 최적화가 아니다

문제 정의
현재의 대규모 사전 학습 워크플로우는 일반적으로 알고리즘 스케일링 결정과 시스템 구현을 분리하여 처리합니다. "스케일링 법칙(scaling law)" 단계에서는 고정된 연산량(FLOPs) 제약 조건 하에서 손실(loss)을 최소화하기 위해 모델 아키텍처와 데이터 예산을 선택하며, 그 뒤를 이어 "시스템" 단계가 하드웨어 효율성을 위해 구현 방식을 최적화하는 과정이 단절된 상태로 진행됩니다. 저자들은 이러한 분리가 희소 혼합 전문가(sparse Mixture-of-Experts, MoE) 모델에는 부적절하다고 주장합니다. 전통적인 스케일링 법칙은 모델 FLOPs($C_{model}$)를 최적화하지만, 해당 아키텍처의 시스템 효율성은 무시합니다. 동일한 예측 손실과 $C_{model}$을 가진 두 아키텍처라도 모델 FLOPs 활용률(MFU), 통신 비용, 메모리 풋프린트의 차이로 인해 학습 처리량(throughput)은 판이하게 다를 수 있습니다. 결과적으로, 특정 FLOPs 예산으로부터 도출된 "컴퓨트 최적(compute-optimal)" 구성은 특정 하드웨어 클러스터에서 실제로 지속 가능한 "전달 가능한(deliverable)" 연산량을 고려하지 못하기 때문에, 해당 클러스터에서 실행 불가능하거나 최적이 아닐 수 있습니다.

방법론: MOSAIC
저자들은 모델 아키텍처와 시스템 공동 설계(co-design)를 하나의 최적화 문제로 정식화하는 프레임워크인 MOSAIC(Model Optimization via Systems-Aware Training Co-design)를 소개합니다. MOSAIC는 하드웨어 제약 조건 하의 예산 내에서 예측 손실을 최소화하도록 모델 아키텍처, 학습 토큰 예산, 그리고 분산 실행 레이아웃(병렬화 전략)을 공동으로 선택합니다.

이 프레임워크는 세 가지 핵심 구성 요소를 통합합니다:

1. 4차원 합동 스케일링 법칙: 저자들은 총 파라미터 수($N_{tot}$), 학습 토큰($D$), 희소도($S$), 그리고 **전문가 분할 계수($G$)**를 포함하는 새로운 희소 MoE용 스케일링 법칙을 피팅했습니다. $G$는 밀집 피드포워드 너비와 단일 전문가 너비의 비율($G = d_{ff}/d_{expert}$)로 정의되며, 전문가의 세밀도(granularity)를 포착합니다. 이 법칙은 다음과 같이 정식화됩니다:
   $$L(N_{tot}, S, D, G) = \frac{a}{N_{tot}^\alpha} + \frac{b}{D^\beta} + \frac{c}{(1-S)^\lambda} + \frac{j}{(1-S)^\delta N_{tot}^\gamma G^\eta} + e$$
   이 정식화를 통해 추상적인 파라미터 수가 아닌 이산적인 기하학적 선택을 바탕으로 손실을 예측할 수 있습니다.

2. 연산자 수준 성능 모델: 시스템 효율성을 추정하기 위해, 저자들은 주어진 아키텍처와 병렬화 레이아웃에 대한 MFU 및 메모리 풋프린트를 예측하는 분석적 성능 모델을 개발했습니다. 이 모델은 다음 사항들을 고려합니다:

   • 연산 비용 (dense 및 grouped GEMM, attention).
   • 통신 오버헤드 (expert parallelism을 위한 all-to-all, 파이프라인 버블, data-parallel collectives).
   • CPU 디스패치 지연 시간 (작은 마이크로 배치에 결정적임).
   • 노드 간 스트래글러(straggler) 및 노드 내 중첩(overlap)을 위한 하드웨어 특화 보정 계수.
     이 모델은 최대 180억 개의 활성 파라미터를 가진 학습 실행을 통해 검증되었으며, MFU 예측에서 15% 미만의 평균 절대 백분율 오차(MAPE)를 달성했습니다.

3. 결합 최적화 (MINLP): MOSAIC는 필요한 모델 FLOPs($C_{model}$)가 전달 가능한 모델 FLOPs($C_{deliverable}$)를 초과하지 않도록 하는 제약 조건 하에서 예측 손실을 최소화하는 혼합 정수 비선형 계획법(MINLP)을 해결합니다.
   $$C_{deliverable} = C_{peak} \times \text{MFU}(Z, P_\ell, B) \times \eta_{good}$$
   여기서 $C_{peak}$는 원시 하드웨어 천장이며, MFU는 아키텍처에 따라 달라집니다. 최적화 과정은 클러스터의 실제 역량 내에서 손실을 최소화하는 실현 가능한 기하학적 구조($Z$)와 실행 계획($P_\ell$)의 이산 그리드를 탐색합니다.

주요 기여

• 컴퓨트 예산의 재정의: 본 논문은 고정된 $C_{model}$ 예산 하에서 최적화할 경우, 최적의 희소도가 데이터 지원의 상한선(즉, 가능한 가장 희소한 모델)에 위치하는 "경계 추구(boundary-seeking)" 동작이 발생함을 보여줍니다. 이는 시스템 비용을 무시한 채, 더 희소한 모델이 고정된 FLOPs에서 더 낮은 손실을 예측하기 때문입니다.
• 시스템 인지 공동 설계: 스케일링 법칙을 성능 모델과 결합함으로써, MOSAIC는 예산의 기준을 $C_{model}$에서 $C_{deliverable}$로 전환합니다. 이를 통해 최적의 희소도는 경계가 아니라 클러스터의 특정 제약 조건(메모리, 대역폭, 병렬화 오버헤드)에 의해 결정되는 **내부점(interior point)**임을 밝혀냈습니다.
• 성능 모델 검증: 저자들은 희소 MoE를 위한 최적의 병렬 레이아웃(Tensor, Expert, Pipeline, Data Parallelism)을 예측하고 MFU를 정확히 예측하는 검증된 연산자 수준 성능 모델을 제공하며, "최적"의 레이아웃이 특정 기하학 및 클러스터 크기에 매우 의존적임을 보여줍니다.
• 실증적 검증: 이 프레임워크는 총 250B 파라미터 규모의 실험을 포함하여 AWS p6 노드(NVIDIA B200 GPU)에서 테스트되었습니다. 결과는 모델 FLOPs를 기준으로 한 구성 순위가 하드웨어-컴퓨트 효율성(벽시계 시간)을 기준으로 한 순위와 크게 다르다는 것을 확인시켜 주었습니다.

결과

• 경계 vs 내부 최적점: 고정된 $C_{model}$ 하에서는 피팅된 손실이 희소도가 증가함에 따라 단조 감소하여 최적점을 경계($S_{max}$)로 밀어냅니다. 그러나 고정된 하드웨어 예산($C_{deliverable}$) 하에서는 손실 곡선이 내부 최솟값을 나타냅니다. 극도로 희소한 모델의 시스템 비용(통신, 메모리)이 결국 이론적인 손실 이득을 상쇄하기 때문입니다.
• 손실 vs 처리량 트레이드오프: 32개 노드, 20일 예산 실험에서 가장 많은 모델 FLOPs를 실현한 구성(더 밀집되고 거친 전문가 기하 구조)이 가장 낮은 손실을 달기 위해서는 아니었습니다. 반대로, 손실 최적 구성(더 희소하고 세밀한 전문가 기하 구조)은 더 낮은 MFU를 유지했지만, 전달 가능한 FLOPs 제약 내에서 유리한 손실 궤적을 유지할 수 있었기에 최종적으로 더 낮은 손실을 달성했습니다.
• 실제 환경 검증: 8개 노드 클러스터에서의 단계별 사전 학습 실행을 통해 프레임워크를 검증했습니다. 측정된 MFU 순서는 예측된 순서와 일치했습니다. 또한, "피크 등가 하드웨어 컴퓨트(MFU를 고려한)"에 대해 손실을 플롯했을 때, 구성의 순위가 원시 모델 FLOPs에 대한 순위와 뒤바뀌는 것을 확인하였으며, 이는 "컴퓨트 최적" 모델이 "클러스터 최적" 모델은 아니라는 점을 입증합니다.

의의
본 논문은 프런티어 언어 모델 학습 방식의 근본적인 변화를 주장합니다. 즉, 단절된 단계인 "스케일링 법칙 선택"과 "시스템 튜닝"에서 통합된 아키텍처-시스템 공동 설계로 이동해야 한다는 것입니다. 저자들은 희소 MoE 모델의 경우, 모델의 시스템 효율성을 가격에 반영하지 못하기 때문에 모델 FLOPs 예산에만 의존하는 것은 불충분하다고 주장합니다. MOSAIC를 도입함으로써, 저자들은 최적의 모델 기하 구조가 보편적인 상수가 아니라 특정 하드웨어 클러스터, 학습 윈도우, 그리고 병렬화 제약 조건에 따라 결정되는 가변적인 것임을 입증하였습니다. 이 접근 방식은 선택된 아키텍처가 이론적으로 손실 최적일 뿐만 아니라, 대상 하드웨어에서 실질적으로 실행 가능하고 효율적임을 보장합니다.