본문으로 건너뛰기

[llm-compressor] SmoothQuant: 활성화→가중치 양자화 난이도 이동

들어가며

SmoothQuant는 LLM의 W8A8 (8비트 가중치, 8비트 활성화) 양자화를 실용화한 기법이다. 대부분의 LLM 활성화는 outlier channel을 가진다 — 몇 개의 채널이 나머지보다 수십 배 큰 값을 가진다. 이 outlier 때문에 activation의 per-tensor 양자화가 실패한다(스케일이 outlier에 맞춰지면 정상 값들의 해상도가 사라진다). SmoothQuant의 아이디어는 "가중치에 per-channel scale을 곱해서 outlier의 '난이도'를 가중치 쪽으로 옮기자"는 것이다. SmoothQuant 논문과 llm-compressor의 src/llmcompressor/modifiers/smoothquant/base.py를 분석한다.

공식 문서

논문 핵심 내용

SmoothQuant의 핵심 관찰은 **"LLM 활성화는 양자화하기 어렵고, 가중치는 양자화하기 쉽다"**는 것이다. Activation outlier가 특정 채널에 집중되는 반면, 가중치는 비교적 균일하다. 이 비대칭을 이용해, activation의 난이도 일부를 가중치로 "이동"시킨다.

주어진 Linear 레이어 $Y = XW$에서 per-channel scale $s$를 도입한다.

$$ Y = X W = (X \cdot \text{diag}(s)^{-1}) \cdot (\text{diag}(s) \cdot W) = \hat{X} \hat{W} $$

여기서 $\hat{X} = X / s$는 outlier가 완화된 smoothed activation, $\hat{W} = s W$는 더 불균일해진 가중치다. 이 변환 후에는 $\hat{X}$가 per-tensor activation 양자화에 친화적이고, $\hat{W}$는 여전히 per-channel 양자화로 처리 가능하다.

$s$는 각 채널의 "migration strength" $\alpha$로 제어된다.

$$ s_c = \max(|X_c|)^\alpha / \max(|W_c|)^{1-\alpha} $$

  • $\alpha = 0$: 난이도 이동 없음 (원본 유지)
  • $\alpha = 1$: 모든 outlier를 가중치로 완전히 넘김 (가중치 양자화 실패)
  • $\alpha = 0.5$: 균형점 (논문 권장값은 LLaMA에서 0.8)

벤치마크 (논문 기준, OPT-175B)

항목 수치
지원 양자화 W8A8 (INT8 가중치 + INT8 활성화)
OPT-175B perplexity (FP16) 10.13
OPT-175B perplexity (naive W8A8) 실패 (PPL 수백)
OPT-175B perplexity (SmoothQuant W8A8) 10.14 (손실 무시할 만함)
추론 속도 향상 1.56배
메모리 절감 2배

SmoothQuant는 활성화까지 양자화한다는 점이 GPTQ/AWQ와 다르다. 이 덕에 GEMM 연산 전체가 INT8로 실행되어 추론 속도가 크게 개선된다.

핵심 구조/코드 분석

SmoothQuantModifier 파라미터

class SmoothQuantModifier(Modifier):
    """
    Implements SmoothQuant from https://arxiv.org/abs/2211.10438
    """
    smoothing_strength: float = 0.5         # alpha — 0.0(원본) ~ 1.0(완전 이동)
    mappings: list[SmoothMapping] | None = None    # smooth layer → balance layer 매핑
    ignore: list[str] = field(default_factory=list)   # 제외할 모듈 패턴
파라미터 기본값 의미
smoothing_strength 0.5 $\alpha$. 논문 기본값은 0.5지만 LLaMA에서는 0.8이 권장되기도 함
mappings None 사용자 정의 smooth-balance 매핑
ignore [] 제외할 모듈

SmoothMapping: AWQ와 동일한 구조

src/llmcompressor/modifiers/smoothquant/utils.pyAWQ와 유사한 매핑 구조를 정의한다.

@dataclass
class SmoothMapping:
    smooth_layer: str           # 스케일을 흡수할 레이어 (RMSNorm 등)
    balance_layers: list[str]   # 스케일을 곱할 Linear 레이어들

실제 매핑 목록도 AWQ와 거의 같다. 예를 들어 LLaMA의 경우:

[
    SmoothMapping(
        smooth_layer="input_layernorm",
        balance_layers=["self_attn.q_proj", "self_attn.k_proj", "self_attn.v_proj"],
    ),
    SmoothMapping(
        smooth_layer="post_attention_layernorm",
        balance_layers=["mlp.gate_proj", "mlp.up_proj"],
    ),
]

AWQ와 SmoothQuant의 차이는 동일 매핑 구조에 다른 스케일 공식을 적용한다는 점이다. AWQ는 "가중치를 양자화 친화적으로 만들기" 위함이고, SmoothQuant는 "활성화를 양자화 친화적으로 만들기" 위함이다.

on_start: 활성화 통계 수집 훅

def on_start(self, state: State, event: Event, **kwargs):
    """Register forward pre-hooks to collect per-channel activation max"""

    # 각 smooth_layer 에 대해 forward pre-hook 등록
    self._hooks = []
    self._input_max = {}   # {smooth_layer_name: per-channel max 텐서}

    for mapping in self.resolved_mappings:
        smooth_module = state.model.get_submodule(mapping.smooth_layer)

        def _hook(mod, args, mapping_name=mapping.smooth_layer):
            x = args[0] if isinstance(args, tuple) else args
            if isinstance(x, tuple):
                x = x[0]
            if x is None:
                return
            # 채널별 최대 절댓값 수집
            cur_max = x.abs().detach().amax(dim=tuple(range(x.dim() - 1)))
            if mapping_name in self._input_max:
                self._input_max[mapping_name] = torch.maximum(
                    self._input_max[mapping_name], cur_max
                )
            else:
                self._input_max[mapping_name] = cur_max

        hook = smooth_module.register_forward_pre_hook(_hook)
        self._hooks.append(hook)

AWQ와 같이 forward pre-hook을 쓴다. 입력의 각 채널별 최대 절댓값(per-channel max absolute)을 누적한다. 이 값이 |X|_c 통계로 쓰인다.

on_sequential_epoch_end 또는 on_finalize: 스케일 계산과 적용

def on_finalize(self, state: State, **kwargs) -> bool:
    """After calibration, compute and apply smoothing scales"""

    for mapping in self.resolved_mappings:
        x_max = self._input_max[mapping.smooth_layer]    # |X|_c

        # balance_layers 의 가중치 최대값 (열 단위)
        w_max_list = []
        for layer_name in mapping.balance_layers:
            module = state.model.get_submodule(layer_name)
            # 가중치의 각 열(입력 채널)별 최대 절댓값
            w_max_list.append(module.weight.abs().amax(dim=0))
        w_max = torch.stack(w_max_list).amax(dim=0)     # 세 Linear 중 최대

        # 스케일 공식: s_c = max(|X|_c)^α / max(|W|_c)^(1-α)
        alpha = self.smoothing_strength
        s = (x_max.pow(alpha) / w_max.pow(1 - alpha)).clamp(min=1e-5)

        # Apply: W ← s * W, smooth_layer.weight ← smooth_layer.weight / s
        self._apply_smoothing(mapping, s, state)

    # 정리
    for hook in self._hooks:
        hook.remove()
    self._input_max.clear()
    return True


def _apply_smoothing(self, mapping, s, state):
    # balance layers: W ← diag(s) * W (입력 채널을 s 로 스케일)
    for layer_name in mapping.balance_layers:
        module = state.model.get_submodule(layer_name)
        module.weight.data *= s.unsqueeze(0)       # shape (out, in) * (1, in)

    # smooth layer: weight ← weight / s (bias 도 동일)
    smooth = state.model.get_submodule(mapping.smooth_layer)
    if hasattr(smooth, "weight"):
        smooth.weight.data /= s                     # RMSNorm weight shape == (hidden,)
    if hasattr(smooth, "bias") and smooth.bias is not None:
        smooth.bias.data /= s

이 코드가 SmoothQuant의 본질이다.

  1. 활성화 통계: forward pre-hook으로 모은 |X|_c
  2. 가중치 통계: balance layers의 각 입력 채널별 최대 절댓값. 세 projection(q/k/v)이 같은 입력을 공유하므로 셋의 최대를 취한다.
  3. 스케일 계산: s_c = |X|_c^α / |W|_c^(1-α)
  4. 변환 적용: 가중치에 $s$ 곱, smooth layer에 $1/s$ 곱. 수학적 등가성 보장.

변환 후 forward를 다시 돌리면 smooth_layer의 출력이 이미 $1/s$로 스케일된 상태라 balance layers의 입력이 작아지고, 그만큼 가중치가 커져 양자화 해상도가 개선된다.

왜 이 설계인가

1. α=0.5 기본값. 논문이 제안하는 "균형점"이다. 0.5는 "가중치와 활성화가 비슷하게 양자화 친화적"이 되는 스케일이다. LLaMA 같은 모델은 0.8이 더 좋기도 하지만, 일반론으로는 0.5가 안전한 선택이다.

2. forward pre-hook 단일 pass. 캘리브레이션 데이터를 한 번만 순회하면 per-channel max가 모두 모인다. GPTQ의 "여러 pass + 역행렬"과 비교해 극도로 가볍다. 70B 모델도 몇 분 안에 SmoothQuant 적용이 끝난다.

3. 동일 매핑 구조 재사용. AWQ와 같은 SmoothMapping 구조를 쓴다. 두 알고리즘이 개념적으로 쌍대적(dual)이기 때문이다. AWQ는 "가중치에 유리하게", SmoothQuant는 "활성화에 유리하게" 이동한다.

4. 수학적 등가성 보장. W *= s, smooth_layer.weight /= s 두 변환이 모델 출력을 정확히 보존한다. 양자화 전까지는 모델 동작이 바뀌지 않아, 사용자는 "이 변환이 모델을 어떻게 망가뜨릴지" 걱정할 필요 없다.

5. W8A8 특화. SmoothQuant는 활성화 양자화를 전제한다. W4A16 같은 weight-only 스킴에는 SmoothQuant를 쓸 이유가 없다. 이 Modifier의 존재 의미는 "activation을 8비트로 양자화하되 outlier 문제를 우회"하는 것.

마무리

SmoothQuant는 W8A8 양자화의 핵심 도구다. INT8 연산이 FP16보다 훨씬 빠르므로, 속도가 우선인 서빙 환경에서 필수다. 다음 글은 또 다른 양자화 알고리즘 AutoRound를 본다.

참고 자료

댓글

관련 포스트

llm-compressor 의 다른글