Sparsely-Gated MoE 구현 관점에서 살펴보기: Router, Top-k, 부하 분산

MoE(Mixture of Experts): 여러 Expert 중 일부를 선택하여 실행하는 구조는 단순히 '거대하지만 가벼운 모델'을 의미하지 않습니다. 구현 관점에서는 다음과 같은 처리를 가진 희소한 FFN(Feed-Forward Network): 각 토큰을 개별적으로 변환하는 레이어입니다.

• Router가 토큰마다 Expert의 점수(score)를 계산합니다.
• Top-k를 통해 소수의 Expert를 선택합니다.
• 토큰을 선택된 곳으로 재배열하여 보냅니다.
• Expert의 출력을 합성하여 원래 토큰 순서로 되돌립니다.

본 글에서는 Router의 텐서 형태(tensor shape), Top-2 계산, 용량(capacity), 부하 분산, Expert Parallelism까지, 구현 시 혼동하기 쉬운 지점들을 중심으로 정리합니다.

MoE는 Transformer의 FFN을 희소하게 만듭니다

대표적인 Transformer MoE에서는 Attention을 Expert로 대체하는 것이 아니라, 블록 내의 Dense FFN을 여러 개의 FFN Expert로 대체합니다.

개념적인 처리 순서는 다음과 같습니다.

hidden states
-> Self-Attention
-> Router projection
...

2017년의 Sparsely-Gated MoE 논문은 stacked LSTM 사이에 MoE 레이어를 배치했습니다. Transformer의 position-wise FFN을 Expert화하는 구성은 GShard와 같은 후속 연구에서 사용되었습니다. 따라서, 원 논문의 구성과 현재 흔히 보는 Transformer MoE를 동일시하지 않는 것이 중요합니다.

또한, MoE 모델의 모든 파라미터가 희소해지는 것은 아닙니다. Embedding, Attention, normalization, 출력 헤드 등은 shared parameters(모든 토큰이 공통으로 사용하는 파라미터)로 남아 있습니다.

Router를 텐서 형태에서 추적하기

배치 크기와 시퀀스를 합친 토큰 수를
Router weight를
토큰

MoE 레이어의 출력은 선택된 Expert 출력들의 가중 평균입니다.
선택되지 않은 Expert에서는

여기서 라우팅 단위는 일반적으로 프롬프트 전체가 아니라 토큰 위치(token position)입니다. 같은 문장의 토큰이라도 다른 Expert로 전송될 수 있습니다. 또한, MoE 레이어가 바뀌면 hidden state도 바뀌기 때문에, 같은 토큰 위치가 다른 Expert를 선택할 수도 있습니다.

Top-2 라우팅의 수치 예시

8개의 Expert에 대한 특정 토큰의 Router logit이 다음 값이었다고 가정합니다.

Top-2는 E4의 2.0과 E6의 1.6입니다. 선택된 값을 softmax로 재정규화하면 가중치(weight)는 다음 값이 됩니다.

따라서, 이 토큰의 출력은 다음 가중 평균입니다.

구현에 따라 모든 Expert에 대한 softmax를 거친 후 Top-k를 선택할지, 아니면 Top-k 선택 후에 재정규화할지는 다릅니다. score function이나 training 시의 노이즈도 한 가지가 아닙니다. 공통적인 핵심은 Router가 희소한 할당과 합성 가중치(composite weight)를 만든다는 것입니다.

total parameters와 active compute 분리하기

MoE의 규모를 읽을 때는 total parameters(모델이 보유하는 모든 파라미터)와 active parameters(1 토큰의 계산 경로에서 사용하는 파라미터)를 구분해야 합니다.

각 Expert가
Top-k 라우팅으로 1 토큰이 사용하는 Expert 부분은 대략 다음 값입니다.

• 모든 Expert 가중치를 보유/배치하는 메모리
• Router 점수를 계산하는 연산 $T imes E$ - shared layer의 계산
• 토큰 재배열 및 장치 간 통신
• Expert별 토큰 수가 적을 경우의 연산 효율 저하

active parameters

어떤 파라미터를 공유할지는 자료마다 다릅니다. Switch Transformers 역시 알고리즘상의 계산량과 실제 학습 속도를 분리하여, 실측값이 저수준(low level) 최적화에 의존한다고 언급합니다. 활성 파라미터 비율만으로 지연 시간(latency) 비율을 추정하는 것은 위험합니다.

순전파 과정(forward pass)은 디스패치와 결합(combine)을 포함

구현에서는 Top-k 인덱스를 구하는 것만으로는 Expert를 효율적으로 실행할 수 없습니다. 개념적인 파이프라인은 다음과 같습니다.

• X @ W_r.T로 Router logits를 계산합니다 - 토큰별로 Top-k 인덱스와 가중치를 얻습니다
• Expert별 용량 슬롯을 할당합니다
• 토큰을 Expert 순서대로 재배열(permute)하여 디스패치합니다
• Expert별로 FFN을 배치 실행합니다
• 출력을 역재배열하고, Router 가중치로 결합합니다.

Expert마다 Python 루프를 돌리기만 해도 작은 행렬 곱셈이 대량으로 발생합니다. 실제 시스템에서는 Grouped GEMM(여러 행렬 곱셈을 묶어 실행하는 기법)이나 kernel fusion(여러 처리를 하나의 커널로 통합하는 최적화)이 사용됩니다. Megatron Core의 MoE 문서에서는 토큰 디스패처, Grouped GEMM, 통신 오버랩 등을 개별 설정으로 확인할 수 있습니다.

로드 밸런싱 손실(load-balancing loss)과 용량(capacity)은 별개

Router가 특정 Expert만 계속 선택하면, 바쁜(busy) Expert는 더 많은 데이터로 업데이트되고, 유휴(idle) Expert는 학습 기회를 잃게 됩니다. 이는 계산 자원의 편중일 뿐만 아니라, 학습상의 긍정적인 피드백 루프가 됩니다.

이 문제에 대해 Sparsely-Gated MoE는 Expert의 중요도(importance)와 부하(load)를 균등하게 하는 보조 손실(auxiliary loss)을 사용했습니다. GShard 역시 디스패치 비율과 평균 게이트 값에 기반한 보조 손실을 도입하고 있습니다.

다만, 보조 손실은 '반드시 균등하게 분배하는 규칙'이 아닙니다. 태스크 손실(task loss)과 동시에 최적화되기 때문에, 계수가 너무 약하면 편향이 남아있고, 너무 강하면 내용에 따른 라우팅보다 균등화를 우선할 가능성이 있습니다.

한편, 용량은 실행 시 하나의 Expert가 받을 수 있는 토큰 수의 상한선입니다. 개념적인 식은 다음과 같은 형태로 나타낼 수 있습니다.

| 메커니즘 | 주요 단계 | 목적 | 너무 강하거나 크면 | :--- |
| load-balancing loss | 학습(training) | 라우팅 분포의 극단적인 편향을 억제 | 태스크에 유용한 편향까지 억제할 수 있음 |
| expert capacity | 순전파 과정 (forward pass) | 버퍼와 계산량에 상한선을 설정 | 빈 슬롯, 메모리, 통신 증가 |

용량을 초과하는 할당은 오버플로우(overflow)입니다. 구현에 따라 다른 후보로 돌리거나, Expert 계산을 건너뛰고 잔차 경로(residual path)만 통과시키는 등의 처리가 있습니다. 용량을 작게 하면 패딩은 줄어들기 쉬우나, 오버플로우가 늘어납니다.

ST-MoE의 router z-loss는 Router logit이 너무 커지는 것을 억제하는 안정화 기법입니다. 이는 로드 밸런싱이나 용량과는 역할이 다릅니다.

Expert 병렬화(Expert Parallelism)에서는 All-to-All이 포함됨

모든 Expert가 하나의 가속기(accelerator)에 들어가지 않는 경우, Expert를 여러 장치(device)로 나누는 Expert 병렬화를 사용합니다.

예를 들어 GPU 0이 가진 토큰의 선택처 Expert가 GPU 3에 있다면, 그 hidden state를 GPU 3으로 보내야 합니다. 각 장치가 서로 다른 전송지와 데이터(data)를 가지기 때문에, 개념적으로 다음 두 번의 All-to-All(모든 장치 간에 서로 다른 데이터를 교환하는 통신)이 발생합니다.

• 토큰을 선택한 Expert의 장치로 보낸다
• Expert 출력을 원래 토큰 소유 장치로 돌려준다

Expert FFN의 FLOPs가 작더라도, dispatch 대기 시간이 길면 wall-clock latency는 짧아지지 않습니다. Expert Parallelism에서는 통신량(communication volume), network topology, Expert별 batch size, 계산과 통신의 overlap을 함께 평가해야 합니다.

pure Python으로 Top-k Router 확인하기

다음 코드는 하나의 토큰에 대한 logit에서 Top-k를 선택하고, 선택된 값만으로 stable softmax를 계산합니다. capacity, tensor batch, 분산 dispatch, backpropagation은 포함하지 않은 교육용 최소 예시입니다.

from __future__ import annotations
import logging
import math
...

max_logit을 인수로 받는 이유는 지수 계산의 overflow를 피하기 위함입니다. Expert ID는 0부터 시작하므로, 실행 결과의 3과 5가 표의 E4와 E6에 대응합니다.

실제(production) 구현에서는 Router projection, Top-k, permute, Expert GEMM, 역permute를 tensor operation으로 통합합니다. 이 함수가 재현하는 것은 선택(selection)과 weight 정규화만입니다.

production에서 무엇을 기록할 것인가

loss와 throughput만으로는 MoE 고유의 문제를 분리하여 진단하기 어렵습니다. 적어도 다음 지표들을 layer 단위로 기록해야 라우팅(routing)과 실행 효율성을 분리하여 추적할 수 있습니다.

Dense FFN과 MoE FFN의 트레이드오프

요약

MoE 구현은 Router의 Top-k만으로는 완성되지 않습니다.

설계 및 운영 측면에서는 다음 3가지 쌍을 분리하여 평가하면 정리하기 쉽습니다.

• total parameters와 active compute
• load-balancing loss와 expert capacity
• algorithmic FLOPs와 통신을 포함한 wall-clock time

다음 시간에는 라우팅을 Top-1로 단순화한 Switch Transformer를 다루며, Top-1이 정말 단순히 그런 것인지 깊이 파고들겠습니다.

MoE의 전제부터 원 논문 결과, 부하 분산과 capacity의 구체적인 예시까지 한국어 도해로 확인하고 싶다면 개인 블로그 완전판도 참조해 주세요.

참고 자료

• Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer
• GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding
• Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity
• ST-MoE: Designing Stable and Transferable Sparse Expert Models
• NVIDIA Megatron Core: Mixture of Experts

Discussion