[HOBBIT] A Mixed Precision Expert OffloadingSystem for Fast MoE Inference

 

 

 

MoE 아키텍처로 연산 속도는 높였지만, 여전히 모든 파라미터를 메모리에 올려둬야 해서 메모리가 작은 장치에서는 MoE 아키텍처로 만들어진 모델을 실행시키기 어렵다.

 

본투비 MoE 모델인 Mixtral 8x7B은 FFN 층을 8개로 나누고 매 토큰마다 전문가 2개만 선택해서 계산하는데, 활성 파라미터 수는 토큰당 12.9B 정도지만 전체 파라미터 수는 46.7B 정도여서.. 양자화 하지 않는다고 치면 파라미터를 저장하기 위해 메모리가 87GB 정도 필요하다.

 

 

 

 

 

대부분의 Expert-Offloading 기술은 중요한 파라미터를 VRAM에 저장하고, 나머지 파라미터를 RAM에 저장해 필요에 따라 VRAM에 있는 파라미터를 제거하고 RAM에서 가져오는 방식으로 작동한다. 

 

이런 방식도 좋긴 하지만.. 누락된 파라미터를 가져오는 작업의 오버헤드가 너무 크다.

EdgeMoE는 Expert에 적용하는 양자화 수준을 조작하고, AdapMoE는 특정 Expert를 건너뛰는 방식으로 최적화한다.

MoE-Offloading은 현재 레이어의 입력과 다음 레이어의 입력으로 필요한 전문가를 예측한다.

그런데 이런 최적화 방식도 GPU 계산 비용보다 Expert 로딩 시간이 훨씬 크기 때문에 한계가 있음;

 

모든 파라미터를 VRAM에 올리는건 불가능. 그러니 VRAM에는 진짜 자주 사용되는 Expert의 파라미터만 올려두고, 나머지는 RAM이나 SSD 등 보조 기억 장치에 저장하고 필요할 때 불러오자.

 

이 때 Cache Replacement Policy를 적절하게 설정해 VRAM에 올려둘 파라미터를 결정해야 함.

운영체제가 하는 역할과 같음.. LFU나 LRU를 사용하는 경우도 있고, Random Replacement Policy를 사용하는 경우도 있다.

 

하지만 모델의 고유한 특성을 고려하면 좀 더 최적화 할 수 있지 않을까? 

 

 

 

 

 

 

 

GPU 메모리에는 모든 토큰을 처리할 때 공통적으로 필요한 Non-Expert 파라미터와 자주 사용되는 Hot-Expert 파라미터의 일부를 올려놓고, CPU 메모리와 SSD 같은 보조 기억 장치에는 GPU 메모리에 넣지 못한 무거운 데이터들을 저장하고 필요에 따라 GPU 메모리로 불러와서 사용한다.

 

(b)를 보면 어차피 토큰 하나를 처리할 때는 전체 파라미터의 31%만 사용함을 확인할 수 있음. 

놀고 있는 Expert가 많으니 파라미터를 차등적으로 배치하는 것.. 

 

 

 

 

(a)를 보면 실제 계산 시간보다 Expert Load 시간이 훨씬 더 오래 걸린다.

목적은 명확해짐. 성능 저하를 최소화하면서, 양자화를 통해 모델의 파라미터 크기를 줄여보자.

중요도가 낮은 Expert를 양자화하면 Expert Load 비용을 줄이면서 성능 저하도 최소화 할 수 있다. (float16 -> int4 양자화 시 4배 빠름)

 

가중치는 보통 FP32(32비트 부동소수점) 형식을 사용하는데, INT8이나 INT4로 바꾸면 32비트를 표현해야 하던걸 8비트나 4비트로 표현할 수 있어 메모리를 크게 줄일 수 있다. - Quantization

 

단순히 소숫점을 버리는 방식이 아님. 최댓값과 최솟값을 기준으로 일정한 간격을 나눠서 매핑하는 방식을 사용한다.

대충 이 구간쯤 위치한다~ 라고 근사하는 방식.. 그러다 보니 원본 값과 차이가 발생하고, 이 차이가 쌓이면 계산 결과가 틀리거나 llm이 헛소리를 하게 됨. 

 

그러니 중요한 가중치는 원본을 최대한 유지하고, 덜 중요한 가중치만 깎아내는 Mixed Precision 방식을 사용해야 함.

 

 

 

 

 

 

HOBBIT은 Dynamic Expert Loader, Adaptive Expert Predictor, Multidimensional Cache Manager로 구성된다.

각각 Expert가 VRAM에 없을 때 동적으로 Expert를 불러오고, 라우터 입력을 보고 다음 Expert를 예측하고, VRAM을 관리하는 알고리즘을 담당한다.

 

1. 라우터의 출력에 따라 MoE 계산에 필요한 상위 k개의 Expert를 선택하고 다음 레이어에 필요한 Expert를 예측함.

2. 캐시 관리자는 필요한 Expert가 VRAM에 있는지 확인하고 우선순위를 업데이트함. 

3. 모두 VRAM에 있으면 바로 계산하고 없으면 SSD나 RAM에서 VRAM으로 옮긴 후, 필요하다면 VRAM을 업데이트함.

 

어떤 Expert가 중요하지 않은지 판단할 때는.. 특정 데이터셋을 Profiling 해서 판단할 수도 있지만, 이러면 그 데이터셋에 종속적이니 다양한 배포 환경에서 적합하지 않다. 

그러니 런타임 입력을 기반으로 Expert의 중요도를 동적으로 평가하는 방법이 필요하다.

 

라우터의 출력값 G(x)와 Expert의 연산 결과인 E(x) 사이에 강한 상관관계가 있으니 이를 바탕으로 Expert를 세 가지로 분류한다.

 

High Precision - 오차가 T1 이하. FP16으로 표현한다.

Low Precision - 오차가 T1 ~ T2 사이. INT4 등으로 양자화 후 표현한다.

Skip - 오차가 T2 이상. 메모리에 올리지 않고, 계산도 생략한다.

 

$$se_i(x) = \begin{cases} \sum_{j=0}^{i-1} \|G(x)e_j\|, & i > 0 \\ 0, & i = 0 \end{cases}$$

 

여기서 T1과 T2는 사용자가 직접 설정해야 하는 Threshold.

위의 수식으로 계산한 Unimportance Degree Score(Expert가 결과값에 미치는 영향)을 기반으로 결정한다.

속도가 빨라야 하면 T1과 T2를 낮게 설정하고.. 정확도가 중요하면 T1과 T2를 높게 설정하고..

 

이 때 첫 번째 Expert는 라우터 점수가 가장 높은 Expert로 모델의 최종 출력 결과에 가장 결정적인 영향을 미치니 항상 High Precision으로 표현한다.

 

요점은, 실시간으로 연산해서 양자화 하는게 아니라 SSD / RAM 에서 VRAM으로 가져올 때 가져오는 수준을 결정하는 것.
그러니 SSD / RAM 에는 모든 Expert의 원본 파라미터와 양자화된 파라미터를 준비해 둬야 한다. 

 

 

 

 

 

빠르게 가져오는 방법은 양자화를 사용하는 것. 필요한 걸 미리 가져와 속도를 높여보자.

 

보통 MoE는 층 하나가 끝나야 다음 층의 라우터를 계산할 수 있지만, 여러 층의 라우터를 하나로 묶어서 한 번에 연산해보자.

연속된 층 끼리는 입력 데이터의 유사성이 높다. 그러니 현재 층의 입력으로 다음 층의 Expert를 예측할 수 있음.

 

예측해서 Expert를 미리 VRAM에 올렸는데, 예측이 틀리면... 패널티가 크게 발생함...

그래프의 c와 e를 보면, 예측한 Expert를 Low Precision으로 가져왔기에 패널티가 매우 적은걸 확인할 수 있음.

 

VRAM의 상황을 보고 여유가 좀 있다 싶으면 더 많은 층을 예측해서 가져오고, 부족하면 다음 층만 예측해서 가져온다.

 

 

 

 

 

캐시 교체 알고리즘인 Least Frequntly Used와 Least Recently Used 방식은 모든 데이터의 value와 가져오는 오버헤드가 같다고 가정하지만, HOBBIT 시스템에서는 그렇지 않음.

 

양자화를 적용하기에.. FP16으로 표현된 Expert를 SSD / RAM에서 가져오는 비용은 INT4로 표현된 Expert를 가져오는 비용보다 4배 더 비싸다.

 

그러니 새로운 알고리즘인 Least High Precision Frequently Used (LHU) 방식을 사용한다.

얼마나 자주 사용됐는지를 기준으로 잡지 않고, 얼마나 자주 High Precision으로 사용됐는지를 기준으로 잡는다.

 

Expert가 High Precision으로 사용된 횟수를 기록하고, 전체 사용 빈도는 낮더라도 한 번 사용될 때 High Precision으로 사용된 Expert는 VRAM에 좀 더 오래 남겨둔다.

 

(a)와 (b)를 보면, LFU는 4번을 남기지만 LHU는 6번을 남긴다. 
그러니 6번 Expert를 FP16으로 다시 읽어들이는 비용을 줄일 수 있음.

 

 

$$ p_t = w_{lru} p_t^{lru} + w_{lfu} p_t^{lfu} + w_{lhu} p_t^{lhu} + w_{fld} p_t^{fld} $$
$$ \text{subject to: } w_{lru} + w_{lfu} + w_{lhu} + w_{fld} = 1 $$
$$ p_t^{lru} = \frac{R_t}{T}, \quad p_t^{lfu} = \frac{F_t}{T}, \quad p_t^{lhu} = \frac{H_t}{T}, \quad p_t^{fld} = 1 - \frac{(l_i - l_t + l_n) \% l_n}{l_n} $$

 

 

캐시 관리 수식은 위와 같이 구성됨. (LHU + FLD + LFU + LRU)

 

최근에 사용한 토큰일수록 점수가 높고 (LRU)

문장에서 해당 Expert가 많이 호출될수록 점수가 높고 (LFU)

Expert가 High Precision으로 사용될수록 점수가 높고 (LHU)

현재 층과 물리적으로 가까울수록 점수가 높다. (FLD)

 

 

 

---

 

 

 

이론은 여기까지고.. 이제 테스트

 

 

하드웨어 

1. RTX 4090 - 24GB VRAM / 256GB RAM / 64 Core

2. Jetson AGX Orin - 32 GB VRAM, RAM (통합) / 12 Core / NVMe SSD

 

모델 및 데이터셋

Mixtral-8x7B (45B 파라미터와 8 Expert) 와 Phi-MoE (42B 파라미터와 16 Expert) 모델 사용 

Alpaca, GSM8K, TruthfulQA 데이터셋 사용 

 

기타

RTX 4090 에서는 float 16 정밀도 버전 사용

Jetson Orin 에서는 SSD를 사용하니 int 8 정밀도 버전 사용 

 

 

 

 

 

HOBBIT 아키텍처가 모든 면에서 최고의 성능을 보인다.. (LLama, MoE Infinity, HOBBIT 중)

 

Jetson Orin은 RAM과 VRAM을 공유하기에 GPU 연산을 위해 메모리를 할당하면 CPU 메모리가 부족해서 Page Fault가 발생함.

MoE Infinity는 GPU 서버용으로 설계된 아키텍처라서 그런지 SSD에서 읽어오는 속도가 매우 느렸음.

 

빠른 로딩을 위해 Low Precision Expert를 사용하니까 시스템 성능 저하를 방지하기 위해 모델 정확도를 확인해야 함.

모델의 공식 문서를 잘 읽어보고 설정할 것.. 

 

 

 

HB, MI, LL 은 Inference System

 

Llama.cpp : C++로 작성된 추론 소프투웨어로, 다양한 AI 모델을 CPU나 GPU에서 돌릴 수 있게 함

MoE-Infinity : 기존에 연구된 MoE 전용 추론 시스템. MoE 모델을 효율적으로 돌리는 소프트웨어

HOBBIT : 새로 제안한 추론 시스템으로, Llama.cpp을 기반으로 개조함.

 

즉, Inference System은 모델을 재생하는 프로그램이라고 보면 됨. 

 

MoE-Infinity와 HOBBIT은 MoE 아키텍처를 위해 설계된 시스템으로, 일반적인 Dense 모델에서는 최적화 기술을 사용할 수 없음.

Llama.cpp는 Dense 모델에는 최적화되어 있지만, MoE 모델을 리소스가 부족한 환경에서 돌릴 때는 효율이 떨어지니 HOBBIT을 도입함.

 

Llama 2, Llama 3, Mixtral-8x7B 같은건 모델 그 자체로, Inference System 위에서 모델이 돌아간다.

 

 

 

 

 

 

 

참고 

1. HOBBIT: A Mixed Precision Expert Offloading System for Fast MoE Inference

https://arxiv.org/abs/2411.01433