본문으로 건너뛰기

[vLLM] AWQ: 활성화 인식 가중치 양자화

들어가며

AWQ(Activation-aware Weight Quantization)는 MIT Han Lab에서 제안한 양자화 기법으로, GPTQ와는 다른 철학을 가진다. GPTQ가 헤시안 행렬을 통해 최적의 양자화를 찾는다면, AWQ는 훨씬 단순한 관찰에서 출발한다: 가중치의 모든 채널이 동등하게 중요하지 않으며, 활성화(activation)가 큰 채널에 연결된 가중치가 더 중요하다. 이 채널들을 보호하면 전체 양자화 품질이 크게 향상된다.

공식 문서

vLLM 공식 문서: AutoAWQ Quantization

핵심 구조/코드 분석

AWQConfig: 설정과 제약

vllm/model_executor/layers/quantization/awq.py에 구현된 AWQConfig는 GPTQ보다 훨씬 간결하다.

class AWQConfig(QuantizationConfig):
    """Config class for AWQ.
    Reference: https://arxiv.org/abs/2306.00978
    """
    def __init__(
        self,
        weight_bits: int,
        group_size: int,
        zero_point: bool,
        modules_to_not_convert: list[str] | None = None,
    ) -> None:
        self.weight_bits = weight_bits
        self.group_size = group_size
        self.zero_point = zero_point
        self.modules_to_not_convert = modules_to_not_convert or []
        if self.weight_bits != 4:
            raise ValueError(
                "Currently, only 4-bit weight quantization is supported for "
                f"AWQ, but got {self.weight_bits} bits."
            )
        self.pack_factor = 32 // self.weight_bits  # 8

GPTQ와 달리 AWQ는 현재 4비트만 지원한다. pack_factor는 8로, INT32 하나에 4비트 가중치 8개가 패킹된다. modules_to_not_convert는 양자화를 건너뛸 모듈을 지정하는데, 보통 임베딩 레이어나 LM 헤드가 여기에 포함된다. 또한 최소 GPU 요구사항이 Compute Capability 75(Turing)이다.

AWQLinearMethod: 가중치 레이아웃

GPTQ와 AWQ의 가장 큰 차이점은 가중치 패킹 방향이다:

class AWQLinearMethod(LinearMethodBase):
    def create_weights(self, layer, input_size_per_partition, ...):
        qweight = PackedvLLMParameter(
            data=torch.empty(
                input_size_per_partition,
                output_size_per_partition // self.quant_config.pack_factor,
                dtype=torch.int32,
            ),
            input_dim=0, output_dim=1,
            packed_dim=1,  # 출력 차원에서 패킹!
            packed_factor=self.quant_config.pack_factor,
        )

GPTQ는 packed_dim=0(입력 차원)에서 패킹하지만, AWQ는 packed_dim=1(출력 차원)에서 패킹한다. 이 차이는 커널 최적화 전략과 직결된다.

스케일과 제로포인트도 그룹 단위로 관리된다:

        num_groups = input_size_per_partition // group_size
        qzeros = PackedvLLMParameter(
            data=torch.empty(num_groups, output_size_per_partition // pack_factor,
                             dtype=torch.int32),
            ...
        )
        scales = GroupQuantScaleParameter(
            data=torch.empty(num_groups, output_size_per_partition,
                             dtype=params_dtype),
            ...
        )

추론: 동적 커널 선택

AWQ의 apply 메서드에는 흥미로운 휴리스틱이 있다:

def apply(self, layer, x, bias=None):
    FP16_MATMUL_HEURISTIC_CONDITION = x.shape[:-1].numel() >= 256

    if FP16_MATMUL_HEURISTIC_CONDITION or envs.VLLM_BATCH_INVARIANT:
        out = ops.awq_dequantize(qweight, scales, qzeros, 0, 0, 0)
        out = torch.matmul(reshaped_x, out)
    else:
        out = ops.awq_gemm(reshaped_x, qweight, scales, qzeros, pack_factor)

토큰 수가 256 이상이면 역양자화 후 FP16 행렬곱을 수행하고, 그 미만이면 융합 AWQ GEMM 커널을 사용한다. 배치 크기가 클 때는 cuBLAS의 FP16 GEMM이 더 효율적이기 때문이다. VLLM_BATCH_INVARIANT 환경변수가 설정되면 항상 역양자화 경로를 택하는데, 이는 Triton 오버라이드를 위한 배치 불변 모드이다.

MoE 지원: Marlin 폴백

AWQ는 FusedMoE 레이어에 대해 두 가지 경로를 제공한다:

def get_quant_method(self, layer, prefix):
    if isinstance(layer, FusedMoE):
        if not check_moe_marlin_supports_layer(layer, self.group_size):
            # MoeWNA16 커널로 폴백
            return MoeWNA16Config.from_config(config).get_quant_method(layer, prefix)
        # AWQ Marlin MoE 커널 사용
        awq_marlin_config = AWQMarlinConfig.from_config(marlin_compatible_config_dict)
        return awq_marlin_config.get_quant_method(layer, prefix)

Marlin이 지원되면 AWQ Marlin MoE 커널을, 아니면 WNA16 커널을 사용한다. 이런 계층적 폴백 구조가 다양한 하드웨어에서의 호환성을 보장한다.

왜 이 설계인가

1. 4비트 전용: AWQ 논문에서 제안한 per-channel scaling trick은 4비트에 최적화되어 있다. 2비트나 8비트로의 확장은 별도의 커널 최적화가 필요하므로, 현재는 가장 실용적인 4비트만 지원한다.

2. 출력 차원 패킹: AWQ 커널은 출력 뉴런 단위로 병렬화되므로, 출력 차원에서 패킹하면 하나의 INT32 로드로 연속된 출력 채널의 가중치를 한 번에 가져올 수 있다.

3. 동적 커널 전환: 양자화 커널은 작은 배치에서 빛나지만, 큰 배치에서는 오히려 cuBLAS보다 느릴 수 있다. 256 토큰을 기준으로 전환하는 것은 실험적으로 검증된 휴리스틱이다.

4. modules_to_not_convert 자동 감지: maybe_update_config에서 safetensors 메타데이터를 분석하여, 양자화되지 않은 레이어를 자동으로 감지한다. 사용자가 일일이 지정하지 않아도 올바르게 동작한다.

논문 핵심 내용

AWQ 논문의 핵심 관찰은 가중치의 1% 미만이 활성화 분포 기준으로 "중요"하며, 이 소수의 중요 채널을 보호하면 전체 양자화 품질이 크게 향상된다는 것이다. 중요 채널을 FP16으로 유지하는 대신, per-channel scaling을 적용하여 양자화 오차를 줄이는 하드웨어 친화적 방식을 사용한다.

WikiText-2 Perplexity 비교 (INT4, group 128)

모델 FP16 RTN GPTQ AWQ
LLaMA-7B 5.68 5.96 6.22 5.78
LLaMA-13B 5.09 5.25 5.23 5.19
LLaMA-30B 4.10 4.23 4.24 4.21
LLaMA-65B 3.53 3.67 3.66 3.62
Llama-2-7B - 5.73 5.69 5.60
Llama-2-13B - 4.98 4.98 4.97
Llama-2-70B - 3.46 3.42 3.41

3비트 양자화 (INT3, group 128)

모델 FP16 RTN GPTQ-R AWQ
LLaMA-7B 5.68 7.01 6.53 6.35
LLaMA-13B 5.09 5.88 5.64 5.52
LLaMA-65B 3.53 4.24 4.21 3.95
Llama-2-70B - 3.98 3.86 3.74

AWQ는 4비트와 3비트 모두에서 RTN과 GPTQ보다 일관되게 낮은 perplexity를 달성한다. 특히 모델이 커질수록 FP16과의 차이가 줄어든다 -- Llama-2-70B의 경우 4비트 AWQ(3.41)가 FP16에 거의 근접한다.

수학/코딩 태스크

모델 태스크 FP16 RTN GPTQ AWQ
Llama-2-7B GSM8K 13.87 11.07 12.13 13.57
Llama-2-13B GSM8K 26.16 21.23 24.26 25.25
Llama-2-70B GSM8K 56.41 53.98 56.03 56.40
CodeLlama-7B MBPP pass@1 38.53 37.51 31.97 40.64

수학과 코딩처럼 정밀도에 민감한 태스크에서 AWQ의 우위가 더 두드러진다. CodeLlama-7B에서는 AWQ가 FP16보다도 높은 pass@1(40.64% vs 38.53%)을 기록했는데, 이는 양자화가 일종의 정규화 효과를 줄 수 있음을 시사한다.

멀티모달 모델(OpenFlamingo-9B)에서도 INT4-AWQ는 FP16에 근접한 성능을 유지했으며, TinyChat 구현 기준 FP16 대비 3.2-3.3배 속도 향상을 달성했다.

마무리

AWQ는 GPTQ보다 단순한 알고리즘이지만, 활성화 인식이라는 직관적 아이디어로 우수한 양자화 품질을 달성한다. vLLM에서는 동적 커널 선택과 Marlin 폴백을 통해 다양한 배치 크기와 모델 아키텍처에서 실용적인 성능을 제공한다.

댓글

관련 포스트

vLLM 의 다른글