MOE 시각화 – 토큰이 MoE 모델을 통해 실제로 어떻게 흐르는지 확인할 수 있는 도구 구축
AI 모델 내부의 토큰 흐름을 시각적으로 보여 주는 것은 AI 모델의 동작 방식을 쉽게 이해할 수 있게 합니다. AI 모델 설계 단계에서 지정한 Layer나 Expert 등을 실제 학습 종료한 모델에서 설계자의 의도대로 동작하는지 호기심에서 시작한 시험입니다. MoE (Mixture-of-Expert, 전문가 혼합 시스템) 모델은 토큰당 Expert 일부를 사용하도록 되어 있습니다. 실제로 어떤 Expert가 실행되는지 기록해 보니, 그 비율은 예상과 달랐습니다. 시험에 사용한 OLMoE 모델의 경우, allenai/OLMoE-1B-7B-0125는 이름에서 보듯이 전체 7B parameter를 가지고 활성 parameter 1B를 가지는 MoE 모델입니다. 이 모델은 16개 Layer, 64개 Expert, Top-8 Routing을 사용합니다. 산술적으로 대략 계산해 보면 전체 expert수는 16 x 64 = 1,024 Slot 을 가집니다. 토큰 하나의 흐름을 살펴 보았을 때 한 Layer에서 8개의 Expert 결과만 선택하니 활성 비율은 8/64 = 1/8 로서 7B의 1/8 = 0.875 B, 대략1B 활성 parameter를 가진다 할 수 있습니다.
저자는 각 MoE Layer 에 CUDA 이벤트를 추가해서 어떤 Expert가 활성화하는지 확인하여 이를 시각적인 heat map 형태로 표시했습니다. 또한, (일반 모델에서는 시험할 수 없지만) Sequence/Batch 를 조정해서 토큰 처리 과정과 비교를 해 보았습니다. 토큰 흐름은 실제 활성 expert 수가 Top-N 지정 수 만큼이었습니다. 반면에 Sequence/Batch 흐름에서는 동일한 decode 단계에서 거의 모든 expert가 활성화 되는 것을 확인할 수 있었습니다.
아래 내용은 일부 내용을 기계 번역한 것으로서 정확한 내용은 원문을 참고하시기 바랍니다.
※ Github: https://github.com/phenx-inc/moe-lightup
===============
“The sparse model that lights up 93% of itself · Phenx”
- 실험 대상: OLMoE-1B-7B (MoE: 16개 MoE 레이어 × 64개 전문가, top-8 라우팅) 를 NVIDIA RTX 4090에서 실제 디코드로 계측.
- 전체 전문가 슬롯: 16 × 64 = 1,024. 한 단계(한 스텝)에서 켜진(실행된) 슬롯 수는 배치 크기에 따라 크게 달라짐.
- Batch 크기 8: 628/1,024 ≈ 61% 활성화, 스텝당 30.281 ms, 처리량 264 토큰/s.
- Batch 크기 32: 952/1,024 ≈ 93% 활성화, 스텝당 38.128 ms, 처리량 839 토큰/s.
- 결과: Batch 크기 4배 증가에 스텝 시간은 약 26% 증가
— 토큰당 비용은 크게 감소(3.79 → 1.19 ms/token).
의미: MoE의 “매 토큰 당 모델의 일부만 쓴다”는 주장은 토큰 단위로는 맞지만, 실제 서비스에서는 Batch 때문에 대부분의 전문가가 한 번 이상 활성화되어 메모리 전송 비용이 커짐.
라우팅(토큰 → 전문가 선택) 관찰:
- 라우터는 각 레이어마다 독립적으로 실행되며, 토큰 하나는 16개 레이어에서 각각 top-8을 선택해 총 128번의 expert 방문을 함.
- 한 토큰이 도달한 서로 다른 expert 수(실험치): 평균 약 53/64. 가장 적은 경우도 51.
- 인접 레이어 간에 같은 expert가 유지되는 비율은 매우 낮아(평균 1.07/8), “영속적 코드 expert” 같은 단일 분류가 아님.
- 게이트(가중치) 분포가 평탄함: top-1 평균 가중치 약 8.8% (중간값 범위 7–11%), 93%는 15% 미만
→ 모델이 하나의 expert를 따르기보다 8개를 섞어 평균화함.
성능·메모리 세부:
- 각 Expert의 가중치: bfloat16으로 약 12.6 MB. 모든 expert(16×64 = 12.9 GB)는 GPU HBM에 상주.
- 문제는 상주 여부가 아니라 연산 전에 온칩 메모리로 스트리밍해야 하는 비용(데이터 이동)이 병목이라는 점.
- Batch 작업은 이 스트리밍 비용을 여러 토큰이 공유하게 해 효율을 크게 높임.
계측(도구) 관련 관찰
- 저자는 각 MoE 레이어 주위에 CUDA 이벤트를 넣어 실제 실행 타이밍을 기록한 녹화를 공개(플레이어 포함 인터랙티브 시각화 제공).
- 처음의 계측은 Python 루프(페일백 경로)를 계측했는데, 실제로는 퓨즈된 grouped GEMM 커널이 실행되어 루프 계측은 느리고 오해를 불러옴. (퓨즈 경로: ~30.3 ms/스텝, 루프 경로는 훨씬 느림)
- 실험은 실제 실행 경로(퓨즈된 경우)에 대한 수치로 보정됨.
결론적 시사점
- MoE의 “매 토큰 1/8만 계산”이라는 장점은 토큰 단위로는 사실이나, 실제 서비스/Batch 작업 상황에서는 많은 전문가가 활성화되어 메모리 시스템 비용이 여전히 중요하다.
- “어떤 부분이 사용되는가”와 “어떤 시간 범위에 대해”를 구분해 물어봐야 함(토큰-레벨 vs. Batch/워크로드-레벨).
- 라우팅이 얇은 선택(하나의 강한 승자)이 아니라 여러 전문가의 평균화라는 사실은, 전문가 교체·재구성에 대한 내구성/영향 해석에 중요함.
추가 메타정보
- 녹화 및 시각화 도구는 오픈소스로 제공되며 로컬(자기 머신, CUDA/MPS/CPU)에서 실행 가능. 페이지에는 실행/녹화 인터페이스와 실행 방법(레포 복제·로컬 서버 구동) 안내가 포함됨.
- 측정 중 특정 실행에서 3개 슬롯은 한 번도 활성화되지 않았음.
※ 지난 게시글:
- AI 뉴스 훑어보기 – 2026.7.16
- AI 뉴스 훑어보기 – 2026.7.15
- AI 뉴스 훑어보기 – 2026.7.13
- AI 뉴스 훑어보기 – 2026.7.10
- AI 뉴스 훑어보기 – 2026.7.9
※ 출처: r/LocalLLM, r/openclaw, r/unsloth, r/opencode, r/claude
