[llm-compressor] AutoRound: 부호 경사 하강법으로 라운딩 최적화
들어가며
AutoRound는 Intel이 2023년에 발표한 양자화 알고리즘이다. 핵심 아이디어는 "반올림 방향을 고정하지 말고, 학습 가능한 변수로 만들어 부호 경사 하강법(SignSGD)으로 최적화하자"는 것이다. 기존 RTN(Round-to-Nearest)은 각 가중치를 가장 가까운 정수로 반올림하는데, AutoRound는 "올림/내림 중 어느 쪽이 출력 손실을 최소화하는가"를 경사 정보로 결정한다. 결과적으로 3비트 같은 극단 양자화에서 GPTQ 이상의 정확도를 달성한다. AutoRound 논문과 src/llmcompressor/modifiers/autoround/base.py를 분석한다.
공식 문서
- 논문: Optimize Weight Rounding via Signed Gradient Descent for the Quantization of LLMs
- AutoRound 레포: intel/auto-round
- 예제: examples/autoround/
논문 핵심 내용
RTN 양자화는 다음과 같다.
$$ \hat{w} = s \cdot \text{round}\left(\frac{w}{s}\right) $$
여기서 $\text{round}(\cdot)$는 결정적이다. AutoRound는 이를 학습 가능한 offset $V$로 대체한다.
$$ \hat{w} = s \cdot \text{round}\left(\frac{w}{s} + V\right) $$
$V \in [-0.5, 0.5]$는 "이 가중치를 어느 쪽으로 반올림할지"를 연속적으로 제어한다. 이 $V$를 학습 가능한 파라미터로 두고, 출력 손실 $|\hat{W}x - Wx|^2$를 최소화하도록 SignSGD로 업데이트한다.
SignSGD는 일반 SGD와 달리 경사의 부호만 사용한다.
$$ V \leftarrow V - \eta \cdot \text{sign}(\nabla_V L) $$
이는 경사의 크기가 중요하지 않고 방향만 중요한 상황에서 효과적이다. $V$는 어차피 $[-0.5, 0.5]$로 clamp되므로, 작은 학습률로 여러 번 부호 방향으로 이동하면서 최적점에 수렴한다.
벤치마크 (논문 기준)
| 항목 | 수치 |
|---|---|
| 지원 비트 | 2, 3, 4 |
| LLaMA-2 7B W3A16 정확도 (GPTQ 대비) | 약간 우수 |
| LLaMA-2 7B W2A16 정확도 (GPTQ 대비) | 크게 우수 |
| 캘리브레이션 시간 (LLaMA 7B) | ~1 GPU 시간 |
| 재학습 필요 | 없음 (경량 PTQ) |
핵심 구조/코드 분석
AutoRoundModifier 생성자 파라미터
class AutoRoundModifier(Modifier, QuantizationMixin):
"""
Implements AutoRound from https://arxiv.org/abs/2309.05516
"""
nsamples: int = 128 # 캘리브레이션 샘플 수
iters: int = 200 # SignSGD 반복 횟수
lr: float = 0.005 # 학습률
minmax_lr: float | None = None # min/max 학습률 (None 이면 lr 과 동일)
seqlen: int = 2048 # 샘플 sequence 길이
enable_quanted_input: bool = True # 양자화된 입력 사용 (quant-aware)
amp: bool = True # AMP 자동 혼합 정밀도
minmax_shrink: float = 1.0 # min/max 값에 곱할 shrink factor
| 파라미터 | 기본값 | 의미 |
|---|---|---|
nsamples |
128 | 논문 권장치. 128이면 보통 충분 |
iters |
200 | SignSGD 반복. 300 이상은 diminishing returns |
lr |
0.005 | $V$ 업데이트 학습률. SignSGD라 비교적 작게 설정 |
minmax_lr |
None | min/max 자체도 학습할 수 있게 별도 학습률 |
seqlen |
2048 | 캘리브레이션 샘플 시퀀스 길이 |
enable_quanted_input |
True | quant-aware 모드 — 이전 레이어의 양자화 결과를 입력으로 사용 |
amp |
True | FP16/BF16 혼합 정밀도로 속도 개선 |
minmax_shrink |
1.0 | min/max 초기값을 이 factor 로 축소 |
on_initialize: Layer-by-layer 최적화 준비
def on_initialize(self, state: State, **kwargs) -> bool:
if not QuantizationMixin.has_config(self):
raise ValueError("AutoRoundModifier requires quantization config")
QuantizationMixin.initialize_quantization(self, state.model)
return True
여기까지는 Quantization Base와 동일하다. 차이는 on_sequential_epoch_end에서 나타난다.
on_event with SEQUENTIAL_EPOCH_END: 한 서브그래프 양자화
def on_event(self, state: State, event: Event, **kwargs):
if event.type_ == EventType.SEQUENTIAL_EPOCH_END:
subgraph = kwargs.get("subgraph")
# 이 서브그래프에 속한 각 Linear 에 대해 AutoRound 최적화 수행
for module_name in subgraph.consumed_names:
module = state.model.get_submodule(module_name)
if isinstance(module, torch.nn.Linear):
self._auto_round_module(module)
def _auto_round_module(self, module):
W = module.weight.data.float()
scale, zp = self._init_scale_zp(W) # 초기 스케일/제로포인트 (Min-Max 또는 MSE)
# 학습 가능한 V 초기화: shape == W.shape, 값은 [-0.5, 0.5]
V = torch.zeros_like(W, requires_grad=True)
# 캘리브레이션 입력/출력 캡처 (intermediates_cache 참조)
x_batches, y_true_batches = self._get_calib_batches(module)
# SignSGD 루프
for step in range(self.iters):
total_loss = 0.0
for x, y_true in zip(x_batches, y_true_batches):
# 양자화된 가중치 계산
W_q = self._fake_quantize_with_offset(W, scale, zp, V)
y_hat = torch.nn.functional.linear(x, W_q, module.bias)
loss = (y_hat - y_true).pow(2).mean()
loss.backward()
total_loss += loss.item()
# SignSGD 업데이트 — 경사의 부호만 사용
with torch.no_grad():
V.data -= self.lr * V.grad.sign()
V.data.clamp_(-0.5, 0.5)
V.grad.zero_()
# 최종 양자화 적용
W_final = self._fake_quantize_with_offset(W, scale, zp, V).data
module.weight.data = W_final
module.weight_scale = scale
module.weight_zero_point = zp
_fake_quantize_with_offset는 AutoRound의 핵심이다.
def _fake_quantize_with_offset(self, W, scale, zp, V):
"""
Equation from paper:
W_q = scale * clip(round(W/scale + zp + V), qmin, qmax)
V is in [-0.5, 0.5] and is learnable.
"""
q = torch.round(W / scale + zp + V).clamp(qmin, qmax)
return scale * (q - zp)
V가 [-0.5, 0.5]에 clamp되므로 "반올림 방향을 연속적으로 제어"한다. 경계값 ±0.5에서 round가 정수를 바꾸므로 $V$가 양자화된 정수값을 결정하게 된다.
enable_quanted_input 모드
if self.enable_quanted_input:
# 이전 레이어의 양자화 결과를 입력으로 사용
x = prev_layer_quantized_output
else:
x = original_input
이 플래그는 "quant-aware calibration" 모드다. True일 때 각 레이어가 "이전 레이어가 양자화되었다는 사실"을 인지하고 학습한다. 이는 양자화 오차가 누적되는 현실을 반영해 최종 출력 품질을 개선한다. 단점은 Sequential Pipeline이 필수이며 메모리 사용량이 증가.
왜 이 설계인가
1. $V$를 학습 가능 변수로. 기존 RTN은 반올림 방향을 결정적으로 정한다. AutoRound는 이를 "양자화 손실을 최소화하는 방향"으로 학습한다. 단순한 아이디어지만 극단 양자화(2~3비트)에서 큰 효과.
2. SignSGD 사용. 일반 SGD보다 단순하고 빠르다. 경사 크기가 양자화 오차의 스케일에 민감하므로 부호만 쓰는 것이 더 robust하다. 또한 $V$가 bounded variable이라 SignSGD가 자연스럽다.
3. Layer-by-layer 최적화. 전체 모델을 한 번에 학습하지 않고, Sequential Pipeline의 SEQUENTIAL_EPOCH_END 훅에서 한 레이어씩 처리한다. GPTQ와 같은 구조.
4. 작은 학습률 lr=0.005. SignSGD 때문이다. 각 스텝은 부호 방향으로만 이동하므로 큰 학습률은 진동을 만든다. 0.005~0.01이 경험적으로 안정적이다.
5. iters=200 기본값. 논문 실험에서 200~400 사이가 최적이라고 보고한다. 200은 속도와 정확도의 균형점. 필요하면 더 늘릴 수 있다.
마무리
AutoRound는 극단 양자화(2~3비트) 시나리오에서 GPTQ를 능가하는 경우가 많다. 학습 기반이지만 재훈련 없이 200 이터레이션으로 끝나므로 비교적 가볍다. 다음 글은 남은 Logarithmic Equalization을 본다.
참고 자료
관련 포스트
llm-compressor 의 다른글
- 이전글 [llm-compressor] SmoothQuant: 활성화→가중치 양자화 난이도 이동
- 현재글 : [llm-compressor] AutoRound: 부호 경사 하강법으로 라운딩 최적화
- 다음글 [llm-compressor] Logarithmic Equalization: 로그 스케일 채널 균등화
댓글