MELINOE: 메모리 효율적인 MoE 추론을 위한 파인튜닝 기법

초록

MELINOE는 Mixture‑of‑Experts(모델)에서 시퀀스당 활성화되는 전문가 수를 제한하도록 파인튜닝하고, 사전 예측된 선호 전문가를 GPU 캐시에 미리 적재함으로써 CPU‑GPU 간 전송을 크게 감소시킨다. 실험 결과, 기존 오프로드 방식 대비 1.2‑3배, 전송‑중심 방식 대비 최대 14.7배의 처리량 향상을 보이며, 다운스트림 성능도 유지하거나 개선한다.

상세 분석

본 논문은 MoE 모델이 토큰당 일부 전문가만 활성화함으로써 연산 효율을 얻지만, 전체 파라미터 수가 방대해 GPU 메모리에 전체 모델을 적재해야 하는 문제를 지적한다. 기존 연구들은 전문가를 CPU DRAM에 오프로드하고 필요 시 GPU로 전송하는 방식으로 메모리 사용량을 줄였지만, 라우팅이 다양하게 발생하면 빈번한 캐시 미스와 PCIe 전송 지연이 전체 추론 속도를 저해한다. MELINOE는 라우팅 자체를 가변적으로 다루어, 시퀀스 수준에서 전문가 사용을 집중시키는 두 단계 접근법을 제안한다.

첫 번째 단계는 Cache Simulation Loss (L_cs) 라는 보조 손실을 도입해 라우터가 동일한 전문가를 반복적으로 선택하도록 유도한다. 라우터 확률 p(ℓ,t)와 Top‑K 선택 r(ℓ,t)를 이용해 가상의 캐시 상태 c(ℓ,t)를 재귀적으로 업데이트하고, r과 c의 차이를 최소화하는 형태의 손실을 정의한다. γ(감쇠 계수)와 캐시 용량 C를 조절해 LRU‑형태 혹은 LFU‑형태의 캐시 시뮬레이션을 구현한다. 이 손실은 전문가 전환을 최소화하면서도, 전체 시퀀스 간 전문가 다양성을 유지하도록 설계돼 라우터 붕괴(모든 시퀀스가 동일한 전문가만 사용) 현상을 방지한다.

두 번째 단계는 Activation Predictor 를 학습하는 것이다. 파인튜닝된 모델에서 각 레이어별로 선호 전문가 집합이 형성되면, 입력 프롬프트만으로 해당 시퀀스에 가장 많이 사용될 전문가를 예측하는 MLP를 별도로 훈련한다. 추론 시에는 이 예측기를 이용해 GPU‑resident 캐시를 사전 로드하고, 이후 표준 오프로드 추론을 수행한다. 따라서 캐시 미스가 크게 감소하고, 전송 지연이 최소화된다.

실험에서는 OLMoE와 Mixtral‑8x7B 같은 대형 MoE 모델을 대상으로, 다양한 GPU 메모리 제한(25 %~50 % 전문가 캐시) 하에서 throughput을 측정했다. MELINOE는 효율적인 베이스라인 대비 평균 1.5‑2.8배, 전송‑무거운 베이스라인 대비 최대 14.7배의 속도 향상을 기록했으며, GLUE·SQuAD·MT 등 여러 다운스트림 태스크에서 정확도 손실이 없거나 오히려 소폭 개선되었다. 또한, γ와 C에 대한 민감도 분석을 통해 γ=0 (LRU) 설정이 가장 높은 캐시 적중률을 보이며, 캐시 용량이 작을수록 보조 손실의 효과가 크게 나타난다.

이러한 접근법은 기존 오프로드, 프리패칭, 양자화 기법과도 호환 가능하다는 점에서 실용성이 높다. MELINOE는 라우팅을 고정된 것으로 보는 대신, 라우팅을 학습 가능한 변수로 전환함으로써 메모리 제약 환경에서도 MoE 모델을 효율적으로 활용할 수 있는 새로운 패러다임을 제시한다.