기술 요약: OASIS

문제 정의

대규모 언어 모델(LLM)은 추론 과정에서 상당한 메모리와 계산 자원을 요구합니다. 양자화는 압축을 위한 주요 방법이지만, 기존 방식들은 효율성과 정확도 사이의 트레이드오프 문제에 직면해 있습니다:

• 가중치 전용 양자화 (Weight-Only Quantization, WOQ): 메모리는 줄여주지만, 실행 전 가중치를 FP16으로 역양자화하는 데 많은 비용이 들며, 이는 전체 런타임의 20~90%를 소비하기도 합니다.
• 정수 가중치 및 활성화 양자화 (Integer Weight-and-Activation Quantization, INT-WAQ): 완전한 저정밀도 연산을 가능하게 하지만, 제한된 표현 능력으로 인해 심각한 정확도 저하를 겪습니다.
• 비균등 가중치 및 활성화 양자화 (Non-Uniform Weight-and-Activation Quantization, NU-WAQ): 데이터 분포에 더 잘 부합하도록 학습된 코드북을 사용하며 정확도를 개선합니다. 그러나 이 방식의 인덱스 코딩 형식은 표준 저정밀도 하드웨어와 호환되지 않아, 결국 FP16으로 역양자화를 수행해야 하며 이로 인해 계산상의 이점이 상쇄됩니다.

또한, LLM의 활성화 이상치(activation outliers)는 가중치 이상치보다 더 많고 크기가 커서 높은 양자화 노이즈를 유발합니다. 추론 중에 이러한 이상치를 동적으로 식별하는 것은 일반적으로 일반 행렬 곱셈(GEMM) 연산의 크리티컬 패스(critical path)에 탐지 과정을 배치함으로써 레이턴시 오버헤드를 발생시킵니다.

방법론

본 논문은 비균등 양자화된 가중치와 활성화 사이의 효율적인 GEMM을 역양자화 없이 수행할 수 있는 알고리즘-아키텍처 공동 설계 프레임워크인 OASIS(Outlier-Aware LUT-Based GEMM with Dual-Side Quantization)를 제안합니다.

1. WAQ LUT 기반 GEMM

OASIS는 비균등 WAQ(NU-WAQ)의 세 가지 특정 기회를 활용하여 룩업 테이블(LUT) 기반 계산을 최적화합니다:

• 오프라인 사전 계산 (Offline Pre-computation): 스트리밍 활성화에 의존하는 WOQ 방식과 달리, NU-WAQ는 가중치와 활성화 모두에 대해 오프라인으로 학습된 중심점(centroids)을 사용합니다. 이를 통해 전체 LUT를 추론 전에 온칩(on-chip)에 사전 계산하여 저장할 수 있습니다.
• 카테시안 곱 LUT (Cartesian Product LUTs): 전체 내적(inner product)을 저장하는 대신(이는 리덕션 길이에 따라 거대한 LUT를 필요로 함), OASIS는 가중치와 활성화 중심점의 카테시안 곱을 저장합니다. 이는 LUT 크기를 획기적으로 줄여줍니다.
• 높은 병렬성 (High Parallelism): 카테시안 곱 LUT는 리덕션 길이(입력 채널 수)와 독립적이므로, 대규모 입도(granularity)의 GEMM을 가능하게 하며 기존의 그룹 단위 LUT 방식보다 훨씬 높은 계산 병렬성을 제공합니다.

2. 룩어헤드(Look-Ahead)를 통한 이상치 인식 양자화

활성화 이상치를 런타임 페널티 없이 처리하기 위해, OASIS는 이중 브랜치 계산 스킴을 도입합니다:

• 메인 브랜치 (Look-Ahead): 전체 활성화 벡터에 대해 WAQ LUT-GEMM을 수행하며, 이상치를 일반적인 방식으로 양자화된 것처럼 취급합니다. 이 브랜치는 초기에는 이상치의 양자화 오차를 무시합니다.
• 이상치 브랜치 (Error Compensation): 동시에 이상치를 식별하고 그 양자화 오차(잔차, residuals)를 계산합니다. 이 잔차들은 역양자화된 가중치와 곱해져 메인 브랜치에서 발생한 오차를 보정하기 위해 누적됩니다.
• 결과: 최종 출력은 룩어헤드 결과와 오차 보정값의 합이며, 이는 수학적으로 기존의 동적 이상치 탐지와 동일하지만 탐지 과정이 크리티컬 패스에서 제거되었습니다.

3. 하드웨어 아키텍처

OASIS 가속기는 다음을 포함합니다:

• Orizuru: 실시간으로 가장 크거나 작은 활성화 값을 식별하는 경량 트리 기반 top-k 탐지 엔진입니다. 비교 및 레이턴시를 최소화하기 위해 공유 리프 노드를 가진 2단계 이진 트리를 사용합니다.
• 연산 요소 (Processing Elements, PE): 가중치와 활성화 인덱스를 병합하여 인덱스 카운터(Index Counters)에 전달함으로써 각 고유 인덱스 쌍의 빈도를 결정하는 "Concat Unit"을 활용합니다.
• 리덕션 (Reduction): 인덱스 카운트와 LUT에 사전 저장된 카테시안 곱 값을 사용하여 가중 합(weighted sums)을 수행하는 MAC 트리를 수행합니다.

주요 기여

1. WAQ LUT-GEMM: 역양자화 없이 비균등 양자화된 가중치와 활성화 사이의 효율적인 GEMM을 가능하게 하는 새로운 방법으로, 사전 계산된 카테시안 곱 LUT를 사용하여 기존 WOQ LUT-GEMM 방식 대비 LUT 크기를 64배 줄이고 병렬성을 1024배 높였습니다.
2. 이상치 인식 스킴: 런타임 오버헤드 없이 활성화 이상치를 처리하는 동시 룩어헤드 계산 및 오차 보정 메커니즘을 제공합니다.
3. OASIS 가속기: 제안된 WAQ LUT-GEMM 스킴을 효율적으로 지원하는 아키텍처 설계입니다.
4. Orizuru 엔진: 데이터 스트림에서 실시간으로 활성화 이상치를 식별하기 위한 경량 top-k 탐지 엔진입니다.

실험 결과

저자들은 다양한 LLM(OPT, LLaMA, Mistral)과 작업(perplexity, zero-shot accuracy)에 대해 알고리즘 벤치마크를 위해 NVIDIA A100 GPU를, 하드웨어 지표를 위해 사이클 단위 정밀 시뮬레이터를 사용하여 OASIS를 평가했습니다.

• 알고리즘 성능:

  • FP16 베이스라인과 비교했을 때, OASIS는 평균 **1.94%**의 정확도 저하(perplexity 기준)를 기록했으며, 이는 Atom 베이스라인보다 6.34% 낮은 수치입니다.
  • W4A4 정밀도에서, OASIS는 Atom 대비 6.44%, QuaRot 대비 6.92% 향상된 정확도를 보여줍니다.
  • 동적 이상치 탐지(OASIS)는 정적 임계값 방식(OASIS-S)보다 우수한 성능을 보였으며, perplexity를 W4A4에서 0.05, W4A3에서 0.27 감소시켰습니다.

• 하드웨어 성능:

  • FIGLUT ASIC 가속기와 비교하여, OASIS는 평균 3.00배의 속도 향상과 1.44배의 에너지 효율 개선을 달달성했습니다.
  • A100 GPU 및 QuaRot과 비교했을 때, OASIS는 처리량(throughput) 측면에서 각각 5.41배 및 3.12배의 속도 향상을 달성했습니다.
  • OASIS는 기존 WOQ LUT-GEMM 설계(FIGLUT, LUT Tensor Core, LUT-GEMM)와 비교하여 LUT 크기를 62.1배에서 994.2배, 리덕션 FLOPs를 248.6배에서 497.1배 줄였습니다.

의의

본 논문은 OASIS가 비균등 가중치 및 활성화 양자화를 위한 알고리즘 효율성과 하드웨어 구현 사이의 간극을 메운다고 주장합니다. 역양자화의 필요성을 제거하고 카테시안 곱 LUT를 통해 대규모 병렬성을 가능하게 함으로써, OASIS는 높은 모델 정확도를 유지하면서도 상당한 처리량 및 에너지 효율 이득을 얻었습니다. 이 연구는 기존에 하드웨어 호환성 문제로 어려움을 겪었던 NU-WAQ가 FP16 베이스라인에 필적하는 모델 품질을 유지하면서도 커스텀 가속기에서 효율적으로 실행될 수 있음을 입증했습니다.