[Quantization] 이론 이해 및 정리
Quantization에 대해 정리를 하고자 마음 먹은지 몇달은 된 것 같은데, 미루고 미루다 이제야 적게 되네요.
AI Researcher로 일을 하면서 GPU와 같은 하드웨어에 대한 지식은 깊이 알지 못했는데, 최근 들어 V100, A100, H100 등의 Nvidia GPU 출시, 모델 훈련 및 추론 시 활용되는 Flash Attention 등이 등장하면서 그 필요성이 점점 대두되는 것 같습니다.
빠르진 않더라도 착실하게 지식을 쌓아가고자 이렇게 글을 정리합니다!
우선 양자화를 이해하기 위해서는 Floating Point data type에 대해 짚고 넘어가야할 것 같네요. 이런 지식들은 특히 Nvidia 블로그 글에 잘 정리되어 있는데, 저 또한 해당 블로그들을 참고하면서 지식을 정리해보겠습니다. 아래는 quantization의 원리를 이해하기 위한 코드도 포함되어 있으니 참고 부탁드립니다! (제 github에서도 해당 내용을 확인하실 수 있습니다 ^_^)
1. Precision In Computing
모델사이즈는 parameter 수와 precision에 의해 결정되므로, 모델 경량화를 진행할 때 Precision을 어떻게 효율적으로 조정하는지가 가장 중요합니다.
<The IEEE Standard for Floating-Point Arithmetic>은 컴퓨터에서 숫자를 이진수로 표현하는 일반적인 규약입니다. double-precision 형식에서는 각 숫자가 64비트를 차지하고, single-precision 형식은 32비트, half-precision는 16비트만 사용합니다.
double-precision (64 bits):지수에 11비트를, 유효숫자에 52비트를 할당하여 표현할 수 있는 숫자의 범위와 크기를 크게 확장합니다.
single-precision (32 bits): 하나의 비트가 숫자가 양수인지 음수인지를 나타냅니다. 8비트는 지수에 할당되며, 이 지수는 (이진법이기 때문에) 2의 거듭제곱을 의미합니다. 나머지 23비트는 그 숫자를 구성하는 유효숫자(significand)를 표현하는 데 사용됩니다. 현재 AI 영역에서 다루고 있는 대부분의 모델 weights는 single-precision으로 저장합니다 (이를 domain에서는 full-precision이라고도 부릅니다).
half-precision (16 bits): 지수에 5비트, 유효숫자에 10비트만 할당하여 훨씬 작은 범위를 표현합니다.
Multi Precision: 필요한 경우 배정밀도를 사용하고, 다른 부분에서는 반정밀도나 단정밀도를 사용하는 것을 의미합니다.
Mixed Precision: 단일 연산 내에서 다양한 정밀도를 사용하여 효율성을 높이면서도 정확성을 유지합니다. 계산을 반정밀도 값으로 시작한 후 결과를 더 높은 정밀도로 저장하는 것을 예로 들 수 있습니다. 구체적인 예시로, 16비트 행렬을 곱할 때 결과는 32비트 크기로 저장하는 것이 있으며, 이를 통해 최종 계산 결과는 배정밀도로 계산한 것과 유사한 정확도를 가질 수 있습니다.
이외에도 최근 AI 패러다임이 발전하면서 등장한 BF16, TF32도 있습니다 (오른쪽 그림).
half-precision같은 경우 지수에 5bits밖에 사용하지 않기에, single-precision에 비해 나타낼 수 있는 실수 범위가 적어 overflow 문제가 자주 발생하곤 합니다. 쉽게 말하면 지나치게 큰 숫자를 표현하지 못하는 문제가 발생할 수 있는 문제이죠. 이는 모델을 훈련하거나 추론할 때 큰 성능 저하를 야기하곤 합니다. 그렇기에 BF16이 자주 사용됩니다. precision을 조금 잃는 대신, 지수부에 single-precision과 같은 8bits를 할당하여 해당 문제를 해결합니다.
TF32는 여기에 더 나아가, 지수부에 8bits, 유효숫자에 10bits, 부호 1bits를 사용하여 총 19bits로 표현합니다. 이는 내부적으로 특정한 연산을 수행할 때 사용된다고 합니다.
지루했습니다. 자 그럼 간단히 생각해보죠. GPT3가 처음 나왔을 때 175B parameter를 가져 큰 화제를 불러일으켰죠. optimizer 등과 같은 부가적인 요소를 제외하고, 위에서 배운 내용을 토대로 단순 모델 사이즈를 계산해볼 수 있습니다. 모델은 보통 32bits (single-precision)으로 저장되고, 32bits = 4bytes이므로:
175 * 10^9 * 4 = 700GB
말도 안 되게 크다는 것을 알 수 있고, 이 모델을 온전히 GPU에 올려 추론을 진행하기 위해서는 A100 80GB 10장 조금 덜 필요하죠. 대기업이나 기관이 아니고서야 개인이 이러한 모델을 사용할 수 있을리가 없습니다. 이러한 맥락에서 나온 것이 바로 quantization 기술입니다.
AI Acceleration에 대해서는 Quantization 외의 Compression, Compilation 등과 같은 다른 기법도 존재한다고 합니다. 저에게는 생소한 분야이므로, 이를 잘 설명하고 있는 이미지를 첨부합니다.
2. Quantization
[참조 자료]
Qualcomm, Enabling power-efficient AI through quantization(2020.04)
Louizos et al. 2019. Relaxed Quantization for discretized neural networks
Gupta et al. 2015 Deep Learning with Limited Numerical Precision
Bengio et al. 2013. Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation
Quantization이란 weights를 floating points로 저장하지 않고, integer와 scaling factor로 변환하는 방식을 의미합니다.
"Quantization is a linear mapping from floating point values to quantized integer values." 이란 문장이 quantization의 의미를 가장 잘 나타내는 듯 합니다.
이는 곧 기존의 floating point values로 되어 있을 weight tensor 혹은 activation tensor matrix를 integer weight matrix와 하나의 scaling factor로 저장할 수 있음을 의미합니다. 물론, 어느 정도의 precision 손실은 감소해야겠죠.
다음은 Quantization의 종류에 대해 살펴보겠습니다.
2.1. Type of Quantization
Quantization의 종류로는 Symmetric Quantization, Asymmetric Quantization이 있습니다. 정규화에 대해 아신다면 이와 매우 비슷한 성격을 지닌다는 걸 알 수 있습니다.
참고로, 이러한 구분 외에 다양한 구분이 존재합니다. 본 글에서는 Symmetric vs. Asymmetric Quantization에 대해서만 다룹니다.
- Per-tensor vs. Per-channel Quantization
- Static vs. Dynamic Quantization
- Post Training Quantization vs. Quantization Aware Training
- Uniform vs. Mixed Precision Quantization
Floating point value x_float이 있다고 가정하면, 이를 x_quantized로 전환하는 과정을 시뮬레이션 해보면 다음과 같습니다.
# b bits quantization
x_int = round((x_float - offset) / scaling_factor)
x_quantized = clamp(x_int, min=0, max=2^b - 1)
new_x_float = x_int * scaling_factor + offset
위 과정은 floating point value를 quantization하는 과정에 해당하며, 이 과정에서 이루어지는 모든 연산은 FP32로 이루어집니다.
물론, 아래 그림처럼 이러한 quantization은 어쩔 수 없는 computation overhead를 야기하기도 하며, 이는 어떤 종류의 quantization을 사용하느냐에 따라 그 정도가 달라집니다.
그렇다면, quantization 결과 모델의 성능은 어떻게 될까요? 아쉽게도 parameter size가 큰 모델에서는 큰 성능 하락이 없지만, 그렇지 않은 경우 심각한 성능하락이 발생한다고 합니다. 위에서 살펴본 것처럼 Post Training Quantization은 필연적으로 precision 문제를 가지고 있으며, training을 진행할 때 rounding, clamp 과정은 gradient를 소실시키기에 "Quantization을 의식"하지 못한 채 이루어진 경우가 많기 때문입니다. 이러한 문제를 해결하고자 등장한 것이 "Quantization Aware Training"입니다.
2.2. Quantization Aware Training
[Straight-Through Estimator (STE)]
round operation 은 아시다시피 미분가능하지 않으므로, rounded된 weights의 gradient는 0이 되거나 소실됩니다. 이는 필연적으로 역전파를 어렵게 만듭니다. 또한, clamp는 min, max range 안에 값이 놓이도록 자르는 과정인데, 마찬가지로 잘린 weights들은 gradient가 소실되므로 vanishing gradient 문제가 발생합니다. 이러한 문제를 해결하기 위해 등장한 것이 Straight-Through Estimator (STE)입니다. 아이디어는 간단한데, forward pass에서는 round, clamp operation을 그대로 유지하되, backward를 진행할 때는 quantization이 진행되지 않은 것처럼 gradient를 재정의하는 것을 의미합니다. 간단한 예시를 살펴보시죠.
import numpy as np
def quantize(x):
return np.round(x)
def straight_through_estimator(x):
x_quantized = quantize(x)
return x_quantized, x
def train_step(weights, input_data, learning_rate=0.01):
quantized_weights, _ = straight_through_estimator(weights)
output = np.dot(input_data, quantized_weights)
# Loss (just a simple example)
target = 1.0
loss = (output - target) ** 2
# Backward pass: Calculate gradient
gradient = 2 * (output - target) * input_data
# Apply STE: In practice, this means weights are updated using the original gradient
weights -= learning_rate * gradient
return weights, loss
np.random.seed(42)
weights = np.random.randn(3)
input_data = np.random.randn(3)
for i in range(5):
weights, loss = train_step(weights, input_data)
print(f"Step {i+1}, Weights: {weights}, Loss: {loss:.4f}")
위 과정을 살펴보면, 실제 gradient를 취하는 대신, loss function을 미분하여 gradient를 재정의하는 것을 보실 수 있습니다.
그러나, 눈치 빠르신 분들은 분명 이러한 의문도 생기실 겁니다. "quantized value와 real value 간의 mismatch 문제가 발생하지 않는가?"
네 맞습니다. 이를 해결하기 위한 또 다른 해결책들이 있으니, stochastic rounding과 relaxed quantization가 있습니다.
이외에도, min, max를 조정하기 위한 아래와 같은 기술도 있다고 합니다.
quantization의 효율성을 향상시키기 위한 연구는 현재도 매우 활발히 진행 중인 것으로 보입니다. DFQ (Data-Free quantization) 등과 같은 방법도 있는 것 같으니, 자세한 사항은 직접 논문을 통해 공부해보시길 바랍니다.
3. Conclusion
Quantization 기술에 대해서는 문외한이었는데, 모델 사이즈가 커지면서 저와 같은 AI Researcher들도 해당 기술에 대해 깊이 알아야할 니즈가 생긴 것 같습니다. 특히 위에서 언급했던 Quantization Aware Training (QAT)같은 기술은 최근 많은 LLM에서도 Training 단계에서부터 적용될 정도로 매우 핵심적인 기술이라고 생각합니다.
마지막으로 저도 quantization을 공부하면서 느낀 거지만, 직접 코드를 짜며 이해하는 게 가장 정확하고 빠른 방법인 것 같습니다. 제 github에 있는 quantization 공부 자료를 남기니, 관심있으신 분은 들어가셔서 참고하셔도 좋을 것 같습니다!
긴 글 읽어주셔서 감사합니다. 혹시 틀린 부분이 있거나 조언해주실 부분이 있다면 언제든 댓글 남겨주세요 ^^. 감사합니다!