[llm-compressor] AWQ: 활성화 인식 가중치 양자화 구현

2026년 4월 13일수정: 2026년 4월 13일

들어가며

AWQ(Activation-aware Weight Quantization)는 GPTQ와 같은 시기에 제안된 4비트 양자화 기법이다. GPTQ가 "2차 정보로 보정"하는 접근이라면, AWQ는 **"활성화 크기가 큰 입력 채널에 대응하는 가중치 채널을 스케일로 보호"**하는 접근이다. 직관적으로 "중요한 채널의 가중치는 양자화 손실을 덜 입게 만들자"는 아이디어다. 놀랍게도 이 단순한 전략이 LLaMA-1 70B에서 GPTQ와 비슷한 정확도를 달성한다. 이 글은 AWQ 논문과 llm-compressor의 src/llmcompressor/modifiers/awq/base.py 구현을 분석한다.

공식 문서

논문: AWQ: Activation-aware Weight Quantization for LLM Compression and Acceleration
llm-compressor 예제: examples/awq/

논문 핵심 내용

AWQ는 Salient Weight Hypothesis로 시작한다. 논문의 핵심 관찰은 "LLM 가중치의 일부 채널(약 1%)이 출력에 지배적 영향을 미친다. 그 채널들은 큰 활성화 크기와 연관되어 있다"는 것이다. 이를 보호하기 위해 per-channel scaling을 제안한다.

주어진 Linear 레이어 $y = Wx$에 대해, 스케일 $s$를 도입해 다음과 같이 변환한다.

$$ y = W x = (W \cdot \text{diag}(s)) \cdot (\text{diag}(s)^{-1} x) = W' x' $$

여기서 $W' = W \cdot \text{diag}(s)$는 양자화할 새 가중치이고, $x' = \text{diag}(s)^{-1} x$는 새 입력이다. 중요한 트릭은 큰 활성화 채널(s 큼)에 해당하는 가중치가 W'에서 커지므로 양자화 해상도가 높아진다는 것이다. 반대로 x의 해당 채널은 나눗셈으로 작아지는데, 활성화는 양자화하지 않고 FP16/BF16에 두므로 손실이 없다.

s를 어떻게 결정하는가가 알고리즘의 핵심이다. AWQ는 활성화 통계 $|\mathbf{a}|_c$ (각 입력 채널 $c$의 평균 크기)를 사용해 $s_c = |\mathbf{a}|_c^\alpha$ 형태로 구하고, 소수 $\alpha$를 grid search(일반적으로 [0, 1]에서 20단계 정도)로 최적화한다.

벤치마크 (논문 기준, LLaMA-2 7B → W4A16)

항목	수치
지원 비트	3, 4
LLaMA-2 7B PPL (FP16)	5.47
LLaMA-2 7B PPL (RTN 4bit)	6.14
LLaMA-2 7B PPL (GPTQ 4bit)	5.83
LLaMA-2 7B PPL (AWQ 4bit)	5.60
A100 추론 속도 향상	2~3배 (대비 FP16)
재학습 필요	없음

AWQ는 GPTQ보다 정확도가 약간 높고, 무엇보다 역행렬 계산이 없어 속도가 빠르다. 헤시안 대신 활성화 통계만 쓰므로 메모리도 적게 든다.

핵심 구조/코드 분석

`AWQModifier` 생성자 파라미터

class AWQModifier(Modifier, QuantizationMixin):
    """
    Implements AWQ from https://arxiv.org/abs/2306.00978
    """
    # 스케일 탐색 파라미터
    duo_scaling: bool = True                  # 가중치와 활성화 둘 다에 근거한 스케일 계산
    num_grid_samples: int = 20                # 탐색 해상도 (alpha 20 단계)

    # 매핑: "어느 활성화 관측자가 어느 Linear 에 영향을 주는가"
    mappings: list[AWQMapping] | None = None
    dynamic_mappings: bool = True             # 모델 아키텍처에서 자동 매핑 추론

파라미터	기본값	의미
`duo_scaling`	True	weight_max * activation_scale 두 요인을 조합해 스케일 결정
`num_grid_samples`	20	alpha 그리드 탐색 수 (0 → 1 사이 20단계)
`mappings`	None	사용자 정의 매핑 (일반적으로 사용 안 함)
`dynamic_mappings`	True	모델 구조에서 자동으로 매핑 추론 (권장)

`AWQMapping`: 활성화 → 가중치 매핑

AWQ가 어려운 이유는 "어느 활성화 통계를 어느 Linear의 스케일에 써야 하는가"를 정해야 하기 때문이다. 간단한 q_proj라면 입력 x가 그대로 들어가지만, RMSNorm → q/k/v projection의 경우 세 projection이 같은 입력을 공유한다. 이 때 스케일도 공유되어야 한다.

src/llmcompressor/modifiers/awq/mappings.py는 이를 위한 자료구조를 정의한다.

class AWQMapping(BaseModel):
    smooth_layer: str           # 입력의 RMSNorm 같은 "smooth 대상" 레이어 이름
    balance_layers: list[str]   # 공유 입력을 받는 양자화 대상 Linear 이름 목록

예시:

AWQMapping(
    smooth_layer="input_layernorm",     # RMSNorm
    balance_layers=[                     # 이 세 Linear 가 같은 입력을 받음
        "self_attn.q_proj",
        "self_attn.k_proj",
        "self_attn.v_proj",
    ],
)

smooth_layer는 "이 레이어의 weight를 diag(s)^-1로 곱해서 입력 스케일 조정을 흡수"하는 레이어다. 보통 바로 앞의 RMSNorm이다. balance_layers는 그 스케일로 양자화되는 Linear들의 목록이다.

`dynamic_mappings`: 아키텍처 자동 감지

src/llmcompressor/modifiers/awq/dynamic_mappings.py는 모델 클래스를 보고 적절한 매핑을 생성한다.

_ARCHITECTURE_TO_MAPPINGS = {
    "LlamaForCausalLM": [
        AWQMapping(
            smooth_layer="input_layernorm",
            balance_layers=["self_attn.q_proj", "self_attn.k_proj", "self_attn.v_proj"],
        ),
        AWQMapping(
            smooth_layer="self_attn.v_proj",
            balance_layers=["self_attn.o_proj"],
        ),
        AWQMapping(
            smooth_layer="post_attention_layernorm",
            balance_layers=["mlp.gate_proj", "mlp.up_proj"],
        ),
        AWQMapping(
            smooth_layer="mlp.up_proj",
            balance_layers=["mlp.down_proj"],
        ),
    ],
    "Qwen2ForCausalLM": [...],
    "MistralForCausalLM": [...],
    ...
}

dynamic_mappings=True이면 AWQModifier가 모델 클래스 이름을 보고 자동으로 이 목록을 사용한다. LLaMA·Mistral·Qwen2 같은 주요 모델은 모두 사전 정의되어 있다. 커스텀 모델의 경우 사용자가 직접 mappings 리스트를 제공해야 한다.

`on_start` 또는 `on_sequential_epoch_end`에서 스케일 탐색

def _search_scale_for_mapping(self, state, mapping):
    """Grid search for the best alpha for this mapping"""
    # 1) 이 매핑의 balance_layers 에서 활성화 통계 모으기
    x_stats = self._activation_stats[mapping.balance_layers[0]]  # 첫 레이어의 입력

    # 2) 가중치 최대값 수집
    w_max = torch.cat([
        state.model.get_submodule(name).weight.abs().amax(dim=0)
        for name in mapping.balance_layers
    ], dim=0).amax(dim=0)

    # 3) alpha grid search
    best_alpha = 0.0
    best_loss = float("inf")
    for i in range(self.num_grid_samples):
        alpha = i / self.num_grid_samples
        if self.duo_scaling:
            s = (x_stats.pow(alpha) / w_max.pow(1 - alpha)).clamp(min=1e-4)
        else:
            s = x_stats.pow(alpha)

        # 4) 이 s 로 가중치 변환 + 가짜 양자화 후 오차 측정
        loss = self._measure_quant_loss(mapping, s, state)
        if loss < best_loss:
            best_loss = loss
            best_alpha = alpha
            best_s = s

    # 5) 최적 s 를 가중치와 smooth_layer 에 적용
    self._apply_scale(mapping, best_s, state)

핵심은 **duo_scaling**이다. True일 때 공식은

$$ s = \left( \frac{|\mathbf{x}|_c^\alpha}{\text{max}(|W_c|)^{1-\alpha}} \right) $$

이는 활성화 크기와 가중치 최대값을 모두 고려한 balanced scale이다. 논문의 경험적 관찰에 따르면 duo가 pure activation scale보다 안정적이다.

`_apply_scale`: 실제 가중치 변환

def _apply_scale(self, mapping, s, state):
    """
    For each balance_layer, multiply weight columns by s.
    For smooth_layer, divide by s to cancel out.
    """
    # balance layers: W ← W * diag(s)
    for layer_name in mapping.balance_layers:
        module = state.model.get_submodule(layer_name)
        module.weight.data *= s.unsqueeze(0)    # 컬럼(입력 채널)에 s 곱

    # smooth layer: W ← W / diag(s)  (출력 채널에 1/s)
    smooth = state.model.get_submodule(mapping.smooth_layer)
    if hasattr(smooth, "weight"):
        smooth.weight.data /= s.unsqueeze(-1)
    if hasattr(smooth, "bias") and smooth.bias is not None:
        smooth.bias.data /= s

중요한 수학적 대칭: 수식 $y = W'(s^{-1} x) = (W \cdot s)(s^{-1} x)$가 성립하려면 s^-1이 어딘가에서 x에 곱해져야 한다. 그 "어딘가"가 바로 smooth_layer다. RMSNorm의 weight를 s로 나누면, RMSNorm의 출력이 s^-1로 스케일된 상태가 되어 수학적 동치가 유지된다. 활성화 양자화 없이 가중치 양자화만 개선되는 마법이다.

AWQModifier가 QuantizationMixin과 협업하는 방식

AWQModifier는 실제 양자화를 자기가 하지 않는다. 스케일링이 끝나면 QuantizationMixin.on_end가 (또는 후속 양자화 Modifier가) 이미 변환된 가중치를 일반 방식으로 양자화한다. AWQ의 기여는 "양자화 전에 가중치를 더 양자화 친화적인 분포로 만드는" 전처리 단계다.

왜 이 설계인가

1. 활성화 통계만 사용. GPTQ의 헤시안과 달리 AWQ는 $x$의 평균 크기만 계산하면 된다. 메모리와 시간 복잡도가 O(n²) 대신 O(n). 70B 모델도 빠르게 처리된다.

2. Duo scaling 기본값. 순수 activation scale은 때때로 가중치 분포를 왜곡시켜 불안정하다. Duo는 두 요인을 balancing해 안정성을 높인다. 논문 experiment에서도 duo가 권장된다.

3. 자동 매핑 감지. 사용자가 매번 "어느 RMSNorm이 어느 projection에 연결되는지"를 손으로 쓰면 지루하다. dynamic_mappings=True가 모델 클래스 이름으로 자동 추론해 사용자 경험을 개선한다.

4. 수학적 동치 변환. W ← W*s, RMSNorm.weight ← RMSNorm.weight/s 두 변환이 모델 출력을 그대로 유지한다 (양자화 전까지). 이 덕에 AWQ는 "모델을 바꾸지 않고 양자화 친화적으로 재배치"한다고 할 수 있다.

5. Grid search의 단순성. 연속 최적화 대신 20단계 grid search를 쓴다. 복잡한 최적화기가 없고, 각 단계의 계산이 병렬이라 GPU에서 빠르다. 정확도도 충분히 좋다.

마무리

AWQ는 llm-compressor의 두 번째 간판 양자화 Modifier다. GPTQ와 비슷한 W4A16 정확도를 훨씬 빠르게 달성한다. 다음 글은 활성화 분포 자체를 가중치에 "넘기는" SmoothQuant를 본다.

참고 자료

논문: AWQ: Activation-aware Weight Quantization for LLM Compression and Acceleration
공식 문서: examples/awq
소스 코드:
관련 포스트:

llm-compressor 의 다른글

이전글 [llm-compressor] GPTQ: 2차 정보 기반 후훈련 양자화 구현
현재글 : [llm-compressor] AWQ: 활성화 인식 가중치 양자화 구현
다음글 [llm-compressor] SmoothQuant: 활성화→가중치 양자화 난이도 이동

댓글

관련 포스트

llm-compressor 의 다른글