[논문 리뷰] Fast Inference of Mixture-of-Experts Language Models with Offloading

Artyom Eliseev, D. Peter Mazur|arXiv (Cornell University)|2023. 12. 28.

Topic Modeling인용 수 7

한 줄 요약

이 논문은 전문가 접근 패턴과 예측적 사전 로딩을 활용하여 소비자용 하드웨어에서 대규모 믹스처 오브 전문가(MoE) 언어 모델의 추론을 가속화하는 MoE 전용 오프로딩 전략을 제안한다. 혼합 양자화(전문가에 대해 2–3비트, 주의 히든 레이어에 대해 4–16비트)와 LRU 캐싱, 그리고 사전 로딩 전략을 결합함으로써, 데스크톱급 GPU와 무료 등급 Google Colab 인스턴스에서 Mixtral-8x7B-Instruct의 인터랙티브 추론을 2–3토큰/초 속도로 구현한다.

ABSTRACT

With the widespread adoption of Large Language Models (LLMs), many deep learning practitioners are looking for strategies of running these models more efficiently. One such strategy is to use sparse Mixture-of-Experts (MoE) - a type of model architectures where only a fraction of model layers are active for any given input. This property allows MoE-based language models to generate tokens faster than their dense counterparts, but it also increases model size due to having multiple experts. Unfortunately, this makes state-of-the-art MoE language models difficult to run without high-end GPUs. In this work, we study the problem of running large MoE language models on consumer hardware with limited accelerator memory. We build upon parameter offloading algorithms and propose a novel strategy that accelerates offloading by taking advantage of innate properties of MoE LLMs. Using this strategy, we build can run Mixtral-8x7B with mixed quantization on desktop hardware and free-tier Google Colab instances.

연구 동기 및 목표

• 제한된 GPU 메모리로 소비자용 하드웨어에서 Mixtral-8x7B-Instruct와 같은 대규모 MoE 언어 모델의 인터랙티브 추론을 가능하게 한다.
• 비정규적인 전문가 접근 패턴과 높은 메모리 대역폭 요구로 인해 기계적 파라미터 오프로딩이 비효율적인 MoE 모델의 문제를 해결한다.
• 전문가 재사용과 조기 게이팅 예측과 같은 내재된 접근 규칙성을 활용해 MoE 아키텍처에 최적화된 오프로딩을 최적화한다.
• 양자화와 지능적인 오프로딩의 조합을 통해 저사양 하드웨어, 특히 T4 GPU를 탑재한 무료 등급 Google Colab에서도 실용적인 성능을 달성한다.
• 고성능 데이터센터 GPU가 필요 없이 현대적인 MoE 모델을 로컬에서 배포하고 실행할 수 있는 오픈소스 솔루션을 제공한다.

제안 방법

• 자주 사용되는 전문가의 반복적인 GPU-메모리 이동을 줄이기 위해 LRU 캐싱을 활용하는 MoE 인식 오프로딩 파이프라인을 설계한다.
• 게이팅 함수의 출력을 기반으로 전문가를 사전에 로딩하는 사전 예측 전략을 구현하여 계산 시간과 겹쳐 전용 로딩 시간을 줄인다.
• 혼합 양자화를 통합: MoE 전문가에 대해 2–3비트의 HQQ 양자화를, 주의 히든 레이어, 임bedding, 정규화에 대해 4–16비트 정밀도를 사용한다.
• 동적 전문가 가중치 오프로딩을 통해 GPU에서 CPU 메모리로 전환하는 하이브리드 추론 엔진을 사용하며, 핀치드 메모리와 비동기 전송을 통한 효율적인 메모리 관리 기법을 적용한다.
• 초기 레이어의 히든 상태가 후속 레이어에서 필요한 전문가를 예측할 수 있음을 활용하여 예측적 로딩 전략을 설계한다.
• PCIe 밸런스 문제로 인해 제약을 받는 시스템을 고려해 오프로딩 빈도를 최소화하고 캐싱 및 사전 로딩을 통해 데이터 재사용을 극대화한다.

Figure 1 : An example of expert loading pattern in Mixtral-8x7B-Instruct for select layers. Blue cells indicate that a certain expert was active when encoding a certain token; deeper blue indicates higher gating weight. Small gray squares show which experts are cached with an LRU cache for $k{=}2$ .

실험 결과

연구 질문

• RQ1 메모리 제약이 있는 하드웨어에서 MoE 모델의 전문가 접근 패턴을 어떻게 활용해 오프로딩 오버헤드를 줄이고 추론 속도를 향상시킬 수 있는가?
• RQ2 소비자용 GPU에서 실행할 경우 MoE 레이어의 경우 최적의 양자화 정밀도와 모델 성능 간의 균형은 무엇인가?
• RQ3 LRU 캐싱과 사전 예측 로딩 전략이 기계적 오프로딩 대비 MoE 추론에서 GPU-메모리 통신을 얼마나 줄일 수 있는가?
• RQ4 무료 등급 Colab나 데스크톱급 GPU에서 오프로딩과 양자화만으로 Mixtral-8x7B-Instruct를 2–3토큰/초 속도로 인터랙티브로 실행할 수 있는가?
• RQ5 낮은 메모리 환경에서 표준 오프로딩 프레임워크인 Hugging Face Accelerate와 비교해 MoE 전용 오프로딩 전략의 성능은 어떠한가?

주요 결과

• 제안된 MoE 전용 오프로딩 전략은 RTX 3060(12GB), RTX 3080 모바일(16GB), T4(16GB) GPU에서 Mixtral-8x7B-Instruct의 추론 속도를 2–3토큰/초로 달성한다.
• T4 GPU를 탑재한 무료 등급 Google Colab에서 2비트 양자화된 전문가와 완전한 오프로딩을 적용한 결과, 모델은 2.092토큰/초의 속도로 생성한다.
• LRU 캐싱과 사전 예측 로딩의 조합은 기계적 오프로딩 대비 RTX 3060에서 추론 시간을 최대 15% 감소시킨다.
• 전문가에 대해 3비트 HQQ 양자화, 주의 레이어에 대해 4비트 정밀도를 사용할 경우, 퍼플렉서티 저하가 최소화되면서도 성능-정확도 균형이 우수한 결과를 얻는다.
• 특히 PCIe 제약이 있는 시스템에서 기계적 오프로딩(예: Hugging Face Accelerate를 통한 방식) 대비 최대 2.5배 빠른 성능을 기록한다.
• 낮은 VRAM을 가진 RTX 3060 데스크톱이 더 고성능의 RTX 3080 모바일과 거의 유사한 성능을 보이며, 호스트-디바이스 대역폭이 주요 밸런스 요소임을 시사한다.

Figure 2 : (left) LRU cache hit ratio for different cache size $k$ ; (right) speculative loading recall when pre-loading a different number of experts. Regular lines represent loading 1 layer ahead; dashed line stands for 2 layers ahead; dotted line is 10 layers ahead.