AWQ

AWQ (활성화 인지 가중치 양자화)

AWQ(Activation-aware Weight Quantization)는 LLM 성능에 중요한 가중치의 작은 일부를 보존해서, 모델을 4-bit로 압축하면서도 성능 저하를 최소화하는 양자화 기법이에요. 이 가이드에서는 Transformers에서 autoawq로 양자화된 모델을 로딩하고 활용하는 법을 정리해요.

출처: AWQ - Hugging Face Transformers 공식문서

개요

Activation-aware Weight Quantization (AWQ)는 LLM 성능에 중요한 가중치의 작은 비율을 보존해서 모델을 4-bit로 압축하며, 성능 저하를 최소화합니다.

AWQ 알고리즘으로 모델을 양자화하는 라이브러리는 llm-awq, autoawq, optimum-intel 등 여러 가지가 있어요. Transformers는 llm-awq와 autoawq 라이브러리로 양자화된 모델 로딩을 지원합니다. 이 가이드는 autoawq로 양자화된 모델 로딩을 보여주지만, llm-awq 양자화 모델도 과정은 비슷합니다.

아래 명령으로 autoawq를 설치합니다.

pip install autoawq

[!WARNING] AutoAWQ는 Transformers를 4.47.1 버전으로 다운그레이드해요. AutoAWQ로 추론하려면 설치 후 Transformers 버전을 다시 설치해야 할 수 있습니다.

AWQ 양자화 모델을 식별하려면 모델의 config.json 파일에서 quant_method 키를 확인하세요.

{
  "_name_or_path": "/workspace/process/huggingfaceh4_zephyr-7b-alpha/source",
  "architectures": [
    "MistralForCausalLM"
  ],
  ...
  ...
  ...
  "quantization_config": {
    "quant_method": "awq",
    "zero_point": true,
    "group_size": 128,
    "bits": 4,
    "version": "gemm"
  }
}

AWQ 양자화 모델을 from_pretrained()로 로딩합니다. 기본적으로 성능을 위해 다른 가중치는 fp16으로 설정됩니다. dtype 파라미터로 이 가중치들을 다른 형식으로 로딩할 수 있어요.

모델이 CPU에 로딩됐다면 device_map 파라미터로 가속기로 옮기세요.

from transformers import AutoModelForCausalLM, AutoTokenizer
from accelerate import Accelerator
import torch

device = Accelerator().device

model = AutoModelForCausalLM.from_pretrained(
  "TheBloke/zephyr-7B-alpha-AWQ",
  dtype=torch.float32,
  device_map=device
)

attn_implementation으로 FlashAttention2를 켜 추론을 더 가속할 수 있어요.

from transformers import AutoModelForCausalLM, AutoTokenizer
from accelerate import Accelerator

model = AutoModelForCausalLM.from_pretrained(
  "TheBloke/zephyr-7B-alpha-AWQ",
  attn_implementation="flash_attention_2",
  device_map=Accelerator().device
)

융합 모듈 (Fused modules)

융합 모듈(fused modules)은 정확도와 성능을 개선합니다. Llama와 Mistral 아키텍처의 AWQ 모듈은 기본 지원되지만, 미지원 아키텍처의 AWQ 모듈도 융합할 수 있어요.

[!WARNING] 융합 모듈은 FlashAttention2 같은 다른 최적화 기법과 함께 쓸 수 없습니다.

AwqConfig를 만들고 fuse_max_seq_len과 do_fuse=True 파라미터를 설정해 융합 모듈을 켭니다. fuse_max_seq_len은 전체 시퀀스 길이로, 컨텍스트 길이와 예상 생성 길이를 포함해야 합니다. 안전하게 더 큰 값으로 설정하세요.

아래 예시는 TheBloke/Mistral-7B-OpenOrca-AWQ 모델의 AWQ 모듈을 융합합니다.

import torch
from transformers import AwqConfig, AutoModelForCausalLM

quantization_config = AwqConfig(
    bits=4,
    fuse_max_seq_len=512,
    do_fuse=True,
)
model = AutoModelForCausalLM.from_pretrained(
  "TheBloke/Mistral-7B-OpenOrca-AWQ",
  quantization_config=quantization_config
).to(0)

TheBloke/Mistral-7B-OpenOrca-AWQ 모델은 융합 전후로 batch_size=1에서 벤치마크됐어요.

Unfused module

Fused module

융합·비융합 모듈의 속도와 처리량은 optimum-benchmark 라이브러리로도 테스트됐어요.

융합 모듈을 지원하지 않는 아키텍처에서는 AwqConfig를 만들고 modules_to_fuse에 커스텀 융합 매핑을 정의해 어떤 모듈을 융합할지 정합니다.

아래 예시는 TheBloke/Yi-34B-AWQ 모델의 AWQ 모듈을 융합합니다.

import torch
from transformers import AwqConfig, AutoModelForCausalLM

quantization_config = AwqConfig(
    bits=4,
    fuse_max_seq_len=512,
    modules_to_fuse={
        "attention": ["q_proj", "k_proj", "v_proj", "o_proj"],
        "layernorm": ["ln1", "ln2", "norm"],
        "mlp": ["gate_proj", "up_proj", "down_proj"],
        "use_alibi": False,
        "num_attention_heads": 56,
        "num_key_value_heads": 8,
        "hidden_size": 7168
    }
)

model = AutoModelForCausalLM.from_pretrained(
  "TheBloke/Yi-34B-AWQ",
  quantization_config=quantization_config
).to(0)

modules_to_fuse 파라미터는 다음 키를 포함해야 합니다.

• "attention": 융합할 어텐션 레이어 이름을 query, key, value, output projection 순서로. 융합하지 않으려면 빈 리스트.

• "layernorm": 커스텀 융합 LayerNorm으로 교체할 모든 LayerNorm 레이어 이름. 융합하지 않으려면 빈 리스트.

• "mlp": 단일 MLP 레이어로 융합할 MLP 레이어 이름. 순서는 (gate (dense, layer, post-attention) / up / down).

• "use_alibi": 모델이 ALiBi 위치 임베딩을 쓰는지 여부.

• "num_attention_heads": 어텐션 헤드 수.

• "num_key_value_heads": Grouped Query Attention (GQA) 구현에 사용할 key value 헤드 수.

  파라미터 값	어텐션
  num_key_value_heads=num_attention_heads	Multi-Head Attention
  num_key_value_heads=1	Multi-Query Attention
  num_key_value_heads=...	Grouped Query Attention

• "hidden_size": hidden 표현의 차원.

ExLlamaV2

ExLlamaV2 커널은 더 빠른 prefill과 decode를 지원합니다. 아래 명령으로 ExLlamaV2를 지원하는 최신 autoawq를 설치하세요.

pip install git+https://github.com/casper-hansen/AutoAWQ.git

AwqConfig에서 version="exllama"를 설정하면 ExLlamaV2 커널이 켜집니다.

[!TIP] ExLlamaV2는 AMD GPU에서 지원됩니다.

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, AwqConfig

quantization_config = AwqConfig(version="exllama")

model = AutoModelForCausalLM.from_pretrained(
    "TheBloke/Mistral-7B-Instruct-v0.1-AWQ",
    quantization_config=quantization_config,
    device_map="auto",
)

리소스

AWQ 데모 노트북에서 모델 양자화, 허브에 양자화 모델 푸시 등 더 많은 예시를 볼 수 있어요.

더 알아보기 (Learn more)

• 양자화 개요: Quantization Overview
• 다른 양자화 방법: bitsandbytes, GPTQ, GGUF