PagedWeight: Efficient MoE LLM Serving with Dynamic Quality-Aware Weight Quantization

기술 요약: PagedWeight

문제 정의

Mixture-of-Experts (MoE) 대규모 언어 모델(LLM)은 토큰당 일부 전문가만을 활성화함으로써 높은 효율성과 정확성을 제공합니다. 그러나 서빙 시나리오에서 MoE 모델은 심각한 메모리 긴장 상태에 직면합니다. 즉, GPU는 방대한 모델 가중치와 계속 증가하는 Key-Value (KV) 캐시를 모두 저장해야 합니다. PagedAttention과 같은 기술이 KV 캐시의 파편화를 관리할 수는 있지만, GPU 메모리의 60% 이상을 차지할 수 있는 MoE 가중치의 상당한 메모리 점유 문제를 해결하지는 못합니다.

기존의 양자화 방법은 대부분 추론 시작 전에 정밀도를 고정하는 정적 방식입니다. 일부 런타임 정밀도 적응 방법이 존재하지만, 이들은 정밀도 변화를 활용하여 GPU 메모리를 가중치와 KV 캐시 사이에서 동적으로 재할당하지 못합니다. 결과적으로, 작업 정확도를 희생하거나 추론을 중단하지 않으면서, KV 캐시가 확장됨에 따라 덜 중요한 전문가 가중치로부터 메모리를 동적으로 확보할 수 있는 시스템이 부족한 실정입니다.

방법론: PagedWeight

저자들은 전문가 가중치의 정밀도와 KV 캐시 크기 사이의 균형을 맞추기 위해 런타임에 모델 가중치를 동적으로 양자화하는 새로운 MoE LLM 서빙 관리 시스템인 PagedWeight를 제안합니다. 이 시스템은 전문가 가중치를 페이지화된 KV 캐시 블록과 유사하게 취급하여, 동적인 오프로딩(offloading) 및 리로딩(reloading)을 가능하게 합니다.

핵심 구성 요소

1. 페이지 기반 가중치 표현 (Paged Weight Representation):

   • PagedWeight는 레이어, 전문가, 모듈 유형(gate_up 또는 down)에 의해 식별되는 전문가별 선형 블록(linear-block) 단위로 동작합니다.
   • 이 시스템은 가중치를 연관된 룩업 테이블(LUT)과 함께 중첩된 비트 평면(bit-plane) 형식으로 표현하는 Any-Precision LLM (APL) 기술을 활용합니다.
   • **가중치 페이지 테이블(Weight Page Table)**은 각 선형 블록에 대한 희망 비트 너비($d_i$)와 확정된 비트 너비($q_i$)를 추적합니다. 페이지는 GPU 상주, CPU 상주 또는 전송 상태로 존재할 수 있습니다.

2. 품질 인식 런타임 플래너 (Quality-Aware Runtime Planner):
   플래너는 정확도 손실을 최소화하기 위해 다음 세 가지 요소를 기반으로 어떤 가중치를 양자화(비트 너비 축소)할지 결정합니다.

   • 오프라인 글로벌 민감도 (Offline Global Sensitivity): 헤시안 가중치 점수(Hessian-weighted scores)를 사용하여 양자화에 대한 각 선형 블록의 내재적 민감도를 추정합니다.
   • 온라인 라우팅 통계 (Online Routing Statistics): 전문가의 "라우팅 질량"(선택 빈도 및 가중치)을 추적합니다. 자주 라우팅되는("hot") 전문가는 더 높은 비트 너비 하한선과 손상 승수(damage multiplier)를 통해 보호됩니다.
   • 프롬프트 잔차 (Prompt Residuals): 선형 회귀 헤드를 통해 예측된 프롬프트별 보정 항입니다. 이는 입력 시퀀스에 따라 양자화 영향이 달라진다는 점을 인지하여, 현재 프롬프트에 기반해 글로벌 민감도 추정치를 조정합니다.

   플래너는 잠재적인 비트 너비 축소에 대한 "손상(damage)" 점수를 계산합니다. KV 캐시 압박으로 인해 메모리 목표치가 트리거되면, 플래너는 목표를 달ู하기 위해 가장 낮은 손상을 주는 축소 방식(절약된 바이트당 가장 높은 품질)을 탐욕적(greedy)으로 선택합니다.

3. 비동기 실행 및 커널 최적화 (Asynchronous Execution & Kernel Optimization):

   • 비동기 페이지 이동 (Asynchronous Page Movement): 지연 시간을 숨기기 위해 시스템은 가중치 페이지를 CPU RAM으로 오프로드하고 비동기적으로 리로드합니다. 비트 너비 확정은 안전한 경계(예: 플랜이 완전히 실행된 후)에서만 업데이트되어 추론이 중단되지 않도록 보장합니다.
   • 융합된 혼합 정밀도 커널 (Fused Mixed-Precision Kernel): 사용자 정의 CUDA 커널은 Any-Precision 비트 평면과 LUT를 직접 읽어 들여, 단일 융합 연산 내에서 각 선형 블록에 대해 서로 다른 비트 너비를 지원함으로써 오버헤드를 줄입니다.

주요 기여

1. 시스템 설계: 온라인 Any-Precision MoE 가중치를 위해 확정/희망 비트 너비, 페이지 상태 및 메모리 소비를 동적으로 관리하는 페이지 기반 가중치 시스템인 PagedWeight를 제안했습니다.
2. 플래닝 알고리즘: 메모리 압박 상황에서 저손상 가중치 축소 전략을 선택하기 위해 오프라인 민감도, 온라인 라우팅 통계 및 프롬프트 잔차를 합성하는 품질 인식 런타임 플래너를 제안했습니다.
3. 실행 전략: 처리량 손실을 최소 соответственно 최소화하면서 오프로드/리로드 지연을 숨기는 비동기 페이지 이동과 융합된 혼합 정밀도 MoE 커널을 구현했습니다.
4. 경험적 검증: 세 가지 MoE 모델(Qwen1.5-MoE-A2.7B, Mixtral-8×7B, Gemma-4-26B-A4B)에 대한 포괄적인 평가를 통해 기존 베이스라인 대비 우수한 성능을 입증했습니다.

결과

저자들은 언어 모델링, 추론 및 롱 컨텍스트(long-context) 작업에 대해 PagedWeight를 FP16, APL, DP-LLM, MxMoE와 비교 평가했습니다.

• 품질-메모리 트레이드오프: PagedWeight는 베이스라인 대비 최대 72.0%의 GPU 메모리 절감과 1.94배의 처리량 향상을 달성하면서도 FP16에 준하는 정확도를 유지합니다.
• 고정 예산에서의 품질: 유사한 메모리 예산에서, PagedWeight는 양자화 방법들보다 최대 39.3% 향상된 작업 품질을 보여주며, 처리량 손실은 최대 4.1%에 불id합니다.
• 롱 컨텍스트 성능: 엄격한 메모리 예산(예: 10 GB) 하에서도 PagedWeight는 FP16에 필적하는 롱 컨텍스트 품질(Passage Retrieval, NarrativeQA)을 유지하는 반면, 정적 양자화 베이스라인(APL 등)은 심각한 성능 저하를 겪습니다.
• 처리량: PagedWeight는 균일 양자화(uniform quantization) 베이스라인과 대등한 처리량을 보여주며, 관찰된 최대 처리량 감소폭은 배치 크기 4에서 4.1%입니다.
• 어블레이션(Ablation) 연구: 라우팅 통계, 프롬프트 잔차 또는 페이지 이동과 같은 구성 요소를 제거하면 성능이 일관되게 저하되어, 전체 시스템 설계의 필요성을 확인했습니다.

의의

본 논문은 PagedWeight가 모델 작업 정확도, 메모리 소비, 그리고 처리량/지연 시간 사이의 복잡한 트레이드오프를 효과적으로 탐색한다고 주장합니다. 전문가 가중치를 정적인 자산이 아닌 페이지화되고 관리 가능한 단위로 취급함으로써, PagedWeight는 작업 정확도와 GPU 메모리(가중치 vs KV 캐시) 할당 사이의 미개척된 트레이드오프 공간을 열어줍니다. 이 시스템은 동적인 품질 인식 양자화가 정적 양자화나 런타임 메모리 압박에 적응하지 못하는 방법들보다 MoE 서빙을 위한 유효하고 우월한 전략임을 입증합니다.