본문으로 건너뛰기

[llm-compressor] AutoRound: 부호 경사 하강법으로 라운딩 최적화

들어가며

AutoRound는 Intel이 2023년에 발표한 양자화 알고리즘이다. 핵심 아이디어는 "반올림 방향을 고정하지 말고, 학습 가능한 변수로 만들어 부호 경사 하강법(SignSGD)으로 최적화하자"는 것이다. 기존 RTN(Round-to-Nearest)은 각 가중치를 가장 가까운 정수로 반올림하는데, AutoRound는 "올림/내림 중 어느 쪽이 출력 손실을 최소화하는가"를 경사 정보로 결정한다. 결과적으로 3비트 같은 극단 양자화에서 GPTQ 이상의 정확도를 달성한다. AutoRound 논문src/llmcompressor/modifiers/autoround/base.py를 분석한다.

공식 문서

논문 핵심 내용

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 PipelineSEQUENTIAL_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 의 다른글