Optimum Intel으로 Intel Xeon에서 StarCoder 가속

Hugging Face와 Intel은 4세대 Intel Xeon Scalable 프로세서에서 StarCoder-15B 모델에 대해 7배 이상의 추론 가속을 달성했습니다. 이 성능 향상은 optimum-intel 라이브러리를 통해 8비트 및 4비트 양자화 기법과 보조 생성(추측 디코딩)을 통합함으로써 실현됩니다.

추론 성능 벤치마크

양자화와 보조 생성의 통합은 HumanEval 데이터셋에서 모델 정확도를 유지하면서 Time Per Output Token(TPOT)을 크게 감소시킵니다.

StarCoder 양자화 정밀도 HumanEval (pass@1) TTFT (ms) TTFT 가속 TPOT (ms) TPOT 가속
Baseline None A16W16 33.54 357.9 1.00x 181.0 1.00x
INT8 SmoothQuant A8W8 33.96 163.4 2.19x 82.4 2.20x
INT4 RTN (g128) A16W4 32.80 425.1 0.84x 54.0 3.35x
INT8 + AG SmoothQuant A8W8 33.96 183.6 1.95x 24.8 7.30x

Intel Xeon을 위한 양자화 전략

StarCoder를 최적화하기 위해 팀은 자동 회귀 토큰 생성에서 주요 제약인 메모리 대역폭 병목 현상을 두 가지 별도 양자화 방법으로 해결했습니다:

8비트 정적 양자화 (INT8)

SmoothQuant 알고리즘을 사용하여 모델을 최소한의 정확도 손실로 INT8로 양자화합니다. SmoothQuant는 활성화와 가중치 모두에 스무딩 스케일링 팩터를 적용해 활성화에서 큰 크기의 이상치가 미치는 영향을 완화합니다. 이 접근법은 TPOT에서 약 2.20배, Time To First Token(TTFT)에서 약 2.19배의 속도 향상을 제공했습니다.

4비트 가중치 전용 양자화 (INT4)

메모리 사용량과 로딩 시간을 더욱 줄이기 위해 모델 가중치를 그룹별 Round To Nearest(RTN) 양자화(128개 그룹)로 4비트로 양자화했습니다. 이 방법은 TPOT에서 3.35배의 속도 향상을 달성했지만, 4비트 가중치를 계산 전에 16비트로 복원하는 데 필요한 연산 오버헤드 때문에 TTFT에서 0.84배의 속도 저하가 발생했습니다.

보조 생성 (추측 디코딩)

보조 생성(AG)은 작은 고속 초안 모델을 사용해 후보 토큰을 예측하고, 이를 큰 대상 모델이 병렬로 검증함으로써 추론을 가속화합니다.

  • Draft Model: 팀은 bigcode/tiny_starcoder_py(164M 파라미터)를 사용했으며, 이는 대상 StarCoder-15B 모델보다 약 95배 작습니다.
  • Optimal Configuration: 8비트 SmoothQuant를 초안 모델과 대상 모델 모두에 적용하여 최상의 결과(~7.30배 TPOT 가속)를 얻었습니다.

보조 생성의 기술적 트레이드오프

보조 생성에서 8비트 양자화를 선택하고 4비트를 사용하지 않는 이유는 병목 현상의 변화에 있습니다:

  1. 초안 모델: 8비트 양자화된 164M 파라미터 모델은 대부분 CPU 캐시 내에 들어가 메모리 대역폭 병목을 없애고 4비트 WOQ와 관련된 복원 오버헤드를 피합니다.
  2. 대상 모델: 대상 모델은 여러 토큰을 병렬로 검증하므로 병목이 메모리 대역폭에서 연산으로 이동합니다. 8비트 양자화 모델은 복원 연산이 무거운 4비트 모델보다 연산을 회피함으로써 더 나은 성능을 보입니다.

Optimum Intel을 통한 구현

사용자는 표준 AutoModelForCausalLM 클래스를 optimum-intel 라이브러리의 IPEXModelForCausalLM으로 교체하여 이러한 최적화를 구현할 수 있습니다. 이 라이브러리는 4세대 Xeon 프로세서에서 Intel Extension for PyTorch(IPEX)와 Intel Advanced Matrix Extensions(AMX)를 활용합니다.

pip install --upgrade-strategy eager optimum[ipex]

from optimum.intel import IPEXModelForCausalLM
from transformers import AutoTokenizer, pipeline

model = IPEXModelForCausalLM.from_pretrained(model_id)
tokenizer = AutoTokenizer.from_pretrained(model_id)
pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)

Sources