[Paper] HyperParallel-MoE: 멀티코어 인터리브 스케줄링을 통한 Ascend NPU에서의 빠른 MoE 훈련

Source: arXiv - 2605.23764v1

Overview

Modern Mixture‑of‑Experts (MoE) 모델은 데이터를 전문화된 “전문가” 서브‑네트워크의 일부에 라우팅함으로써 대규모 확장을 달성합니다. 이러한 모델을 Ascend NPU—Huawei의 AI 가속기—에서 효율적으로 학습하는 데는 기존 프레임워크가 MoE 연산자를 하나씩 순차적으로 실행하여 칩의 이기종 연산 유닛(행렬‑지향 AIC와 벡터‑지향 AIV)의 활용도가 낮아지는 문제가 있었습니다. HyperParallel‑MoE는 MoE 실행을 미세한 타일‑단위 작업 흐름으로 재작성하고, AIC와 AIV 자원을 동시에 활용하도록 스케줄링하는 새로운 컴파일 및 스케줄링 시스템으로, 핵심인 Dispatch‑to‑Combine 단계에서 최대 1.58배 낮은 지연 시간을 제공합니다.

Key Contributions

• Tile‑level heterogeneous taskflow: 전체 연산자 MoE 파이프라인을 작은 계산·통신 타일들의 정적 그래프로 변환하여 AIC(매트릭스)와 AIV(벡터) 유닛 모두에 매핑할 수 있게 함.
• AIV‑driven one‑sided communication: 집합 라우팅 및 데이터 셔플링을 호스트 CPU에서 분리하여 비용이 많이 드는 호스트 측 동기화를 없앰.
• Unified task abstraction: 통신과 계산을 동일하게 다루어 의존성을 유지하면서도 적극적인 겹침을 허용함.
• Event‑driven static scheduler: 거의 제로에 가까운 런타임 오버헤드로 교차 큐 실행을 조정하여 단일 커널 실행으로 AIC와 AIV 워커를 모두 구동함.
• Integration with MindSpore/MindFormers: 실제 DeepSeek 스타일 MoE 모델에 적용해 기존 최적화 연산자를 재작성하지 않고도 실용성을 입증함.

방법론

1. Static Compilation: MoE 전방 패스(Dispatch → Expert FFN → Combine)를 타일 격자로 분해합니다. 각 타일은 작은 행렬 곱셈 또는 벡터 연산과 필요한 데이터 이동을 나타냅니다.
2. Heterogeneous Mapping: 타일에 선호하는 연산 유닛을 표시합니다—밀집 행렬 작업은 AIC, 벡터 중심 라우팅 및 축소는 AIV.
3. One‑Sided Communication Primitives: 맞춤형 AIV 커널이 Ascend의 내장 교차 큐 동기화 프리미티브를 사용해 칩 내에서 직접 scatter/gather를 수행하므로 호스트가 집합 연산을 기다리며 차단되지 않습니다.
4. Taskflow Generation: 종속성을 보존하는 DAG가 생성되며, 엣지는 데이터 준비 상태를 인코딩합니다. 이 DAG는 이벤트 기반 스케줄러에 전달되어 AIC와 AIV 워커 스레드를 모두 포함하는 single 커널을 실행합니다.
5. Runtime Execution: 통합 런타임은 타일 단위에서 통신, 행렬 곱셈, 벡터 축소를 교차 실행하여 세밀한 겹침을 구현하면서도 전문가 FFN을 위한 기존 고성능 AIC 커널을 계속 활용합니다.

결과 및 발견

• 지연 감소: Ascend A3 클러스터의 여러 전문가‑병렬 구성에서 HyperParallel‑MoE는 기준 MindSpore 구현에 비해 Dispatch‑to‑Combine MoE‑FFN 지연을 최대 1.58배 줄였습니다.
• 자원 활용도: 프로파일링 결과, 작업 흐름이 활성화될 때 AIC 점유율이 ~45 % (기준)에서 >80 %로 상승했으며, AIV 유닛은 유휴 상태가 아니라 통신을 처리하며 바쁘게 동작합니다.
• 확장성: 이 접근 방식은 전문가와 NPU 수가 증가함에 따라 확장되며, 모델이 클러스터 내 여러 노드에 걸쳐 분산될 때에도 성능 향상을 유지합니다.
• 호환성: 기존 전문가 커널(예: fused GEMM+Bias+Activation)은 수정이 필요 없었으며, HyperParallel‑MoE는 단지 타일 스케줄에 감싸서 사용했습니다.

실용적 시사점

• 더 빠른 MoE 훈련: 개발자는 Ascend 하드웨어에서 더 큰 MoE 모델을 실제 시간보다 짧은 시간에 훈련시킬 수 있어 클라우드 비용을 절감하고 실험 속도를 높일 수 있습니다.
• 간소화된 배포: 시스템이 표준 MindSpore/MindFormers 스택 내에서 동작하므로 팀은 모델 코드를 재작성하거나 커스텀 커널을 만들 필요 없이 이를 도입할 수 있습니다.
• 향상된 칩 활용도: 이 방법론은 이기종 가속기에서 숨겨진 병렬성을 추출하며, 이는 다른 워크로드(예: 트랜스포머 스타일 어텐션, 희소 커널)에도 적용할 수 있는 패턴입니다.
• 낮은 에너지 발자국: 높은 연산 활용도는 유휴 사이클을 줄여 훈련 단계당 에너지를 감소시킬 수 있으며, 이는 지속 가능성을 중시하는 AI 연구소에 이점이 됩니다.

Limitations & Future Work

• Static Scheduling Assumptions: 현재 스케줄러는 컴파일‑타임에 정적으로 결정됩니다; 동적 워크로드 변동(예: 전문가 간 부하 불균형)은 여전히 정체 현상을 초래할 수 있습니다.
• Hardware Specificity: 이 설계는 Ascend‑전용 프리미티브(AIC/AIV, 크로스‑큐 이벤트)를 활용합니다. 다른 가속기 계열로 포팅하려면 유사한 이기종 자원이 필요합니다.
• Memory Overheads: 타일‑단위 분해는 중간 데이터에 대한 추가 버퍼링을 도입하는데, 이는 메모리 제한이 있는 디바이스에서 병목이 될 수 있습니다.
• Future Directions: 저자들은 적응형 타일 크기 조정 탐색, 런타임 부하‑균형 통합, 그리고 추론‑시 MoE 라우팅 및 기타 이기종 AI 칩을 지원하도록 프레임워크를 확장할 계획입니다.

저자

• Zewen Jin
• Congkun Ai
• Guangpeng Zhang
• Hanbo Zhang
• Haoran Wang
• Shihan Xiao
• Da Lei
• Xuefeng Jin
• Teng Su
• Cheng Li

논문 정보

• arXiv ID: 2605.23764v1
• 카테고리: cs.DC
• 출판일: 2026년 5월 22일
• PDF: PDF 다운로드