LLM.int8(): 8-bit Matrix Multiplication for Transformers at Scale

2024. 5. 28. 16:17LLM/Paper Reading

LLM 훈련 및 추론에 엄청난 GPU 메모리가 요구된다. 추론 단계만 해도, 모델 파라미터 수 x 자료형 x 1.2 정도의 용량이 필요하다. 예를 들어, 40B 크기의 모델을 4바이트 자료형으로 구동하면 대략 190GB의 메모리가 필요하다. (여기서 1.2는 어텐션 캐시나 토큰 생성과 같은 추가 메모리 요구를 고려한 배수이다.)

이러한 메모리 요구 사항을 줄이려면 자료형을 조정하는 양자화 기법을 사용할 수 있다. 

양자역학에 따르면, 거시적인 세계는 연속적이지만 미시적인 세계는 이산화된 형태에 가까워서, 분절된 행렬 값으로 표현하기 적절하다고 한다. 딥러닝에서의 양자화란, 연속적인 형태의 정보를 이산화된 형태로 변환하여 의미를 유지하면서 정보 크기를 압축하는 기법을 말한다. 그러므로 그 어원은 양자역학에서 온 것이다.

 

자료형을 조정할 때 고려해야 할 몇 가지 옵션은 다음과 같다. 먼저, FP32(32비트 = 4바이트, 완전 정밀도) 및 FP16(16비트 = 2바이트, 반 정밀도)과 같은 일반적인 자료형이 있다. 더 나아가, BF16(같은 16비트이되, 범위 확장 및 정밀도 축소)와 같이 딥러닝을 위해 특별히 설계된 자료형도 있다. 모델을 FP16 또는 BF16로 양자화하고 서빙하는 것은 메모리 사용량을 대폭 줄일 수 있는 흔한 방법 중 하나이다.

- 대부분의 HW는 FP32를 지원하며 (1bit 부호, 8 bits 지수 23 bits 가수(value의 range) ) 정보를 담습니다. 

- FP16에서는  (1bit 부호, 5 bits 지수  , 10 bits 가수 ) 정보를 담습니다.

- BFLOAT16은 FP32과 동일한 지수 부를 가지면서 가수부를 7bits로 만든 것입니다. 이는 곳 가수부의 3bit 손실을 의미한다. 즉, FP32 처럼 큰수를 표현할 수 있으면서 정확도는 조금 포기하는 것입니다. 

- 이상적으로 학습과 추론은 FP32로 이뤄져야 합니다. 이는 FP16/BF16 보다 2배 느립니다. 따라서, 처리 속도를 높이기 위해서 FP16과 정확도 유지를 위해서 FP32를 섞어 학습하는 Mixed Precision 접근법이 사용되어 가중치는 32bits로 저장하되, 순전파/역전파에서는 FP16/BF16으로 수행하여 학습 속도를 높입니다. 역전파 시에 계산된 FP16/BF16 기울기는 FP32 포맷 가중치를 업데이트 합니다. 

==> 경험상 Mixed Precision을 사용하면 배치는 1.5배 이상 늘었고, 학습 속도도 빨랐다


INT8(8비트 = 1바이트, 0 ~ 255 또는 -128 ~ 127의 정수)로 메모리를 더 줄여보는건 어떨까? 이렇게 단순하게 자료형을 줄이면 모델의 성능이 급격하게 감소할 수 있다. 이것은 양자화로 인한 오차가 층 별로 전파되면서 누적되기 때문이다. 특히 저자인 Tim Dettmers는 6B 이상의 트랜스포머 기반 모델에서 단순한 양자화 방식이 실패함을 발견했으며, 이는 이상값으로 인해 발생하는 문제와 관련이 있다. 그는 이러한 이상값을 특정 임계값을 갖는 패턴으로 정의하고 이를 '이머전트 피처'라고 명명하였다.



먼저 벡터를 INT8로 단순히 양자화하는 방법을 설명해야겠다. 최대 절대값 양자화 방법으로써, 다음과 같이 동작한다.

- 절대 최대(absmax) 양자화에서는 FP16 / BF16 숫자와 해당 int8 숫자 사이의 매핑을 계산( 부동 소수점 수를 정수로 변환하는 것을 의미 )하려면 먼저 텐서의 절대 최대값으로 나눈 다음 데이터 유형의 전체 범위를 곱해야 합니다.

- 예를 들어, [1.2, -0.5, -4.3, 1.2, -3.1, 0.8, 2.4, 5.4]를 포함하는 벡터에 절대최대 양자화를 적용한다고 가정해 보겠습니다. 이 경우 절대 최대값인 5.4를 추출합니다. Int8의 범위는 [-127, 127]이므로 127을 5.4로 나누고 배율 인수로 23.5를 얻습니다.  (127을 텐서의 가장 큰 값으로 나눈 값 알파를 얻는다)

- 따라서 원래 벡터에 이를 곱하면 양자화된 벡터 [28, -12, -101, 28, -73, 19, 56, 127]가 생성됩니다.

(텐서 값을 알파 값으로 스케일링한다, 당연 텐서의 가장 큰 값은 127이 된다.)

 

- 최신 값을 검색하려면 int8 전체 정밀도 숫자를 양자화 인자(알파)로 나눌 수 있지만 위의 결과는 "반올림"되므로 일부 정밀도가 손실되긴 합니다.

 

==> 즉, 정밀도를 유지하면서 데이터표현범위를 줄여줌으로서 오차가 층 별로 전파되면서 누적되는것을 방지할 수 있음.


이 방법의 장점 중 하나는 0 값이 오차 없이 그대로 정수로 변환된다는 것이다. 즉, 0 값은 양자화 과정에서 정보 손실이 발생하지 않는다. 이러한 특성은 특히 희소한 벡터 또는 행렬을 다룰 때 유용할 수 있다.


LoRA는 LLM 파인튜닝 기법 중 하나로, 기존 행렬의 가중치를 동결하고 추가로 학습할 때 증분( 가중치 행렬에 더해지는 작은 변화량 Δ𝑊 여기서는 A와 B)만을 따로 학습하면서, 이 증분 행렬을 두 개의 저차원 행렬 곱으로 표현하여 전체 모델의 용량을 줄이는 방법이다.

논문이 제안한 QLoRA 기법은 LoRA와 관련된 기법 중 하나로, 다음과 같이 진행된다.

1. 가중치 저장: QLoRA에서는 LoRA의 가중치들을 4비트로 양자화하여 저장한다. 이때, 일반 FP4가 아닌 NormalFloat라는 특수한 자료형을 사용한다.

NF4라고 불리는데, NF는 quantile quantization(분위 양자화, 정확한 한국어 번역이 없어 직역함)라는 개념에 기반한다. 

위 양자화는 양자화 대상이 되는 데이터 집합의 누적 분포를 통해 해당 분포의 분위를 구하고 이를 양자화 구간으로 간주한다. 이후 데이터를 양자화했을 때, 표의 Normalisation 과 같은 각 양자화 구간에 할당되는 데이터의 개수가 동일하도록 보장하는 양자화 방식이 바로 분위 양자화이다. 즉, 해당 데이터의 분포는 균등하기에, 비선형적인 구간을 가질 확률이 높고, 데이터의 분포가 달라질 때마다 이를 새로 계산해줘야하는 비용이 높은 양자화 방식이다.

 

그런데, 만약 모든 데이터 집합의 분포가 동일하게 고정된다면? 앞서 말한 확률 분포 추정, 분위 분할 등의 과정을 모두 생략하고 이미 정해진 양자화 구간에 데이터 포인트들을 맵핑만하면 된다. 이것이 Normalisation 를 보완한 NF4 데이터 타입이 경량화의 핵심인 이유이다.

  • 8비트 이상의 실수를 NF4 타입의 양자로 맵핑하는 과정을 크게 보면 다음과 같다.
    1.  분포의  분위수를 추정한다.
      (4비트 양자화라면 16개의 양자 후보가 필요하다. 사분위수가 3개이듯, 16개의 양자가 필요한 경우 17분위수를 구해야한다.)
      NF4 타입의 4비트 양자화에서는 16개의 양자화 구간이 필요합니다. 이를 위해서는 16개의 양자화 값을 만들어야 하며, 이 값들을 구간으로 나누기 위해 17개의 분위수를 추정합니다.
    2. 위 1번의 데이터 타입을 [-1, 1]로 정규화한다.
    3. 입력 텐서의 를 정규화에 사용될 양자화 상수로 삼아 양자화를 수행한다.
  • 그러나 위 과정을 그대로 이행하기에는 한 가지 문제가 있다. 16개의 양자 후보가 구간 내 대칭으로 분포하는 경우 중심, 즉 0에 대한 명확한 표현을 가질 수 없다는 것이다(e.g., [-1.0, -0.6, -0.2, 0.2, 0.6, 1.0]). 정확한 0의 부재가 양자화를 불가능하게 만드는 것은 아니나, 0값의 텐서를 입력받았을 때 quantization loss를 유의미하게 증가시킬 수 있다. Task를 막론하고 입력 텐서에는 zero-padding 등 처리해야하는 0이 꽤 많이 존재하기 때문에, 정확한 0의 표현은 무시할 수 없는 부분이다.
  • 이 때문에 구간을 대칭으로 16등분하는 것이 아닌, 음수부를 7등분, 양수부를 8등분한 후 0을 중심으로 구간을 concatenate하게 된다. 이렇게 만들어진 16개 분위수를 [-1, 1] 구간으로 정규화하면 이것이 바로 NF4의 양자화 구간이 되며, 입력된 텐서의 값에 해당하는 양자로 맵핑시킬 수 있게 된다.

Double Quantization

  • Double quantization은 말 그대로 한 번 양자화된 값을 한 번 더 양자화해 추가적인 메모리 이득을 보기 위한 기법이다. 대신 NF4로 양자화된 가중치 자체를 한 번 더 양자화한다는 것은 아니고, 가중치를 양자화할 때 추가적으로 발생하는 값인 '양자화 상수(quantization constant)'를 양자화하는 것이다. 위

k-bit Block-wise Quantization

  • 위의 NF4 양자화는 [-1, 1] 내로 구간을 정규화하는 과정이 포함된다. 그렇다는 말은, 입력 텐서의 absmax 값으로 나머지 값들을 나누는 과정이 필요하다는 뜻이다. 이때 사용되는 absmax가 바로 양자화 상수이다. QLoRA에서 PLM은 얼어있지만 LoRA는 업데이트가 필요하고, 이때의 LoRA 어댑터는 4비트가 아닌 16비트 실수형으로 존재한다. 때문에 4비트 모델과 16비트 어댑터 간 연산을 위해 quantization-dequantization을 반복해야 하므로, 이 양자화 상수는 모델 밖에 따로 저장되어야 한다.
  • 꽤 간단하지만 조금 더 세부적으로 들어가보자면, 그렇다면 이때의 absmax는 전체 가중치의 absmax인가? 하는 의문이 들 수 있다. LLM의 크기가 billion을 넘어 trillion으로 가고 있는 상황에서, 전체 가중치의 absmax를 상수로 고른다면 아무리 zero-centered 정규 분포라 한들 매우 큰 상수가 선정될테고, 대부분의 가중치를 0에 수렴시키는 사태를 발생시킬 것이다. 이를 방지하면서 어느정도의 계산 효율을 보장하기 위해 QLoRA는 k-bit block-wise quantization 방식을 채택한다. 일정 개수의 가중치를 묶어 하나의 블록으로 간주하고, 해당 블록의 absmax를 양자화 상수로써 채택하는 방식이다. 여기서 말한 블록이 k-bit block이다.

Quantize Quantization Constant(양자화 상수)

  • 위 이미지는 16비트 가중치를 4비트와 0.5비트로 양자화하는 과정을 보여주는데, Double quantization은 바로 이 양자화 상수를 양자화하는 것이다. 이전에 언급했듯이, 양자화 상수는 모델 밖에 따로 저장되는 형태이고, 단순히 4비트 모델의 메모리 외에도 이 양자화 상수를 저장하기 위한 약간의 추가적인 메모리 오버헤드가 발생한다. 비율로 보자면 크지 않을 수 있지만 모델이 커질수록 양자화 상수의 절대적 메모리 용량도 커지기 때문에, 여기에도 추가적인 양자화를 가하면서 최대한의 효율을 가져가는 것이다.
    1. Weight (16-bit):
      • 원본 가중치 데이터는 16비트 부동 소수점 값으로 저장됩니다.
    2. 4-bit Quantization:
      • 원본 16비트 가중치를 4비트로 양자화합니다. 이를 위해 16비트 가중치 값을 1st Quantization Constant로 나누어 4비트 정수로 변환합니다.
      • 예: 16비트 가중치 값 2.1을 4비트로 양자화하면 2.1을 0.5비트 양자화 상수로 나누어 4비트 정수 값으로 변환합니다.
    3. 1st Quantization Constant (absmax):
      • k-bit block-wise quantization은 블록의 크기가 클 수록 양자화 상수의 개수는 줄지만 정밀도가 낮아지는 trade-off가 존재한다. 이를 적정 선에서 맞추기 위해 QLoRA는 64의 블록 크기를 채택했다. 즉, 64개의 가중치가 1개의 양자화 상수(64개 가중치의 absmax)를 가지고 양자화된다는 뜻이다. 양자화 상수는 dequantization을 위해 필요한 것이니 FP32로 저장되고, 이를 가중치 1개 당 추가 메모리로 계산하면 32-bit / 64-block = 0.5비트이다. 다시 말해, NF4는 모델 크기의 측면에서 엄밀하게 말하자면 4비트가 아니라 4.5비트인 셈이다.
      • 이때 추가되는 비트를 0.5비트보다 작게 만들기 위해 양자화 상수에 quantization을 가한다. 마찬가지로 block-wise quantization을 사용하며, BF16을 NF4로 만들었던 가중치 양자화와 다르게 양자화 상수 양자화는 FP32를 FP8로 만든다. 이때 가중치 양자화를 위해 생성되었던 양자화 상수를 1차 양자화 상수(first quantization constant), 이 1차 양자화 상수를 한 번 더 양자화할 때 생성되는 양자화 상수를 2차 양자화 상수(second quantization constant)라고 한다.
      • 첫 번째 양자화 상수는 absmax (최대 절대값)로, 0.5비트로 표현됩니다. 이 상수는 원본 가중치의 최대값을 기준으로 합니다.
    4. 2nd Quantization Constant (absmax):
      • 두 번째 양자화 상수도 absmax (최대 절대값)로, 0.002비트로 표현됩니다. 이는 4비트로 양자화된 값의 최대 절대값을 기준으로 합니다.
      • FP8, 8-bit로의 양자화이므로 정밀도를 조금 더 보전할 수 있는 덕에 블록의 크기를 조금 더 늘릴 수 있으며, 256개의 1차 양자화 상수를 하나의 블록으로 보고 1개의 2차 양자화 상수를 생성한다. 결과적으로, 256개의 1차 양자화 상수는 64 * 256개의 가중치에서 도출되므로, 2차 양자화 상수는 가중치 1개 당 32-bit / (64 * 256-block) = 약 0.002비트의 메모리를 추가한다. 이 결과 생성된 '양자화된 1차 양자화 상수'는 기존의 '1차 양자화 상수'와 여전히 동일한 개수로 존재하지만, FP32에서 FP8이 되었으니 가중치 1개 당 비트는 0.5비트에서 8-bit / 64-block = 0.125비트로 감소하게 된다. 즉, 전체적으로 추가되는 비트는 0.5비트에서 (0.125 + 0.002) = 0.127비트가 되어, 가중치 한 개당 0.373비트를 감소(기존 NF4 대비)시켜 저장할 수 있다는 것이다.
    5. 결과:
      • 최종적으로, 이중 양자화(Double Quantization) 과정을 통해 16비트 가중치를 4.127비트로 줄입니다.
      • 메모리 절감 효과는 Weight 하나 당 약 0.373비트입니다.

QLoRA!

  • 상술했듯이 QLoRA는 LoRA에 새로운 양자화 기법을 추가한 느낌이 강한 기술이라, 기존 LoRA와의 차이점으로 QLoRA를 요약 설명해보겠다.
    1. FP32 PLM + FP32 LoRA 조합에서 NF4 PLM + BF16 LoRA 조합으로
    2. 양자화가 없는 타입에서 블록 양자화 상수를 한 번 더 양자화 한 이중 양자화
    3. VRAM 초과시 OOM 에러 발생에서 paged optimization을 통한 CPU RAM 활용으로
  • 바로 다음 섹션에서 설명되는 내용이지만, 번외로 바뀐 점 중 기존의 LoRA가 트랜스포머의  에만 adapted 되었던 것이 '모든 트랜스포머 내 선형 변환'에 붙는다는 점도 있다.

실험 결과

vs Standard Finetuning

기존 LoRA와 Alpaca의 Undertuning 이슈

  • QLoRA의 연구진들은 기존 LoRA가 더 좋은 성능을 낼 수 있었다고 주장했다. 기존 LoRA는 트랜스포머 모듈의 query와 value 벡터를 처리하는  에만 붙어 adaptation(모델 특정부분 학습,조정)을 수행했다. 그러나 위 그래프(LLaMA 7B의 Alpaca 데이터셋 finetuning 결과)에서 볼 수 있듯이 attention 레이어 뿐만 아니라 feed-forward network(FFN)까지 LoRA를 adapt했을 때 더 좋은 성능을 낸 결과를 보였다. 게다가 Alpaca 또한 다소 undertuned(과소 조정)된 부분이 있음을 확인했는데, finetuning 시 학습률과 배치 사이즈에서 재탐색을 수행한 결과 더 높은 RougeL 점수를 도출하는 모습을 보였다.

NF4 vs FP4 / int4

 

  • 위 그래프는 여러 4비트 데이터 타입의 LLaMA가 수행한 벤치마크 결과이며, 두번째 표는 125M~65B의 O 4비트 모델들의 Pile-CC 데이터셋에 대한 perplexity이다. PPL은 '헷갈리는 정도로 낮을 수록 성능이 좋은것임. NF4의 성능이 FP4나 int4 데이터 타입들에 비해 유의미하게 높은 성능을 보이며, double quantization까지 추가되는 경우 성능 저하 없이 메모리 사용량의 감소를 구현할 수 있음을 확인할 수 있다.

k-bit QLoRA vs 16-bit Full-Finetuning / 16-bit LoRA

  • 위 표는 125M~3B 크기 RoBERTa, T5 모델의 데이터 타입 별 GLUE와 Super-NaturalInstructions 벤치마크 점수이다. 두 벤치마크 모두에서 16, 8, 4비트의 어댑터가 16비트 full-finetuning 베이스라인 성능을 거의 재현 가능함을 확인할 수 있다. 저자들은 이 결과를 두고 adapter finetuning after quantization이 quantization loss를 거의 복구할 수 있음을 시사한다고 설명했다.

  • 위 표는 LoRA-adapted 된 7B~65B 크기 LLaMA의 데이터 타입 별 5-shot MMLU 정확도이다. LLaMA의 경우 33B부터는 리소스 한계로 full-finetuning을 할 수 없어 애초에 LoRA로 비교했다고 한다. 두 instruction following 데이터셋을 각각 학습한 모델들을 5-shot MMLU task로 검증했는데, NF4 + DQ 조합의 LLaMA + QLoRA가 BF16 LLaMA + LoRA의 성능을 거의 완벽히 따라잡을 수 있다는 결과가 도출되었다. 반면 FP4 버전은 BF16 버전보다 평균적으로 1% 정도 성능 하락이 있음을 확인하였다


결론적으로 QLoRA 기법은 FP4 또는 INT4로 단순하게 양자화된 모델 성능을 능가하며, 16비트 LoRA 또는 풀 파인튜닝에 준하는 성능을 달성할 수 있다고 한다.