Hugging Face에 대해 알아보자 (4) Quantization

서론

지난 포스팅에서는 Hugging Face를 활용한 generator에 대해 알아보았다. 이번 포스팅에서는 Quantizaion에 대해서 알아보겠다. Quantization의 역할이 뭐고, 어떻게 구현하는지에 대해서 차근차근 서술할 계획이다. 그리고 이번 포스팅을 위해서는 Llama 3.1에 접근할 수 있어야 되는데, 이를 위해서는 https://huggingface.co/meta-llama/Meta-Llama-3.1-8B에 들어가서 로그인하고 access를 요청해야된다. 신청하면 몇분만에 바로 되지는 않고, 두시간 정도 뒤에 될 것이다. 최소 나는 그랬음.

---

Quantization이란?

Quantizaion은 메모리 로드 시 32-bit에서 8-bit 내지 4-bit으로 가중치의 precision을 낮추는 것이다. 이를 통해 메모리 사용량을 줄이고(32->8  or 32->4[bit]), precision을 낮춤으로써 더 빠른 연산을 가능하게 하는 것이다. 혹, 32-bit에서 8-bit이 된거면 성능(Accuracy)이 4배 줄어드는거 아님?할 수도 있다. Trade off 관계는 맞지만,  Accuracy의 줄어듦은 줄어든 메모리 사용량과 비례하지 않는다. 오히려 Accuracy가 적게 줄어든데 비해 메모리 사용량이 크게 늘어서 그만한 가치가 있는 행위다.

---

Quantization을 사용해보자

우선 코랩을 켜서 런타임을 T4로 설정한다. API도 있어야됨. 이 과정이 모른다면 여기로 가면 된다. 차근차근 설명해놨음. 아무튼 설정이 됐다면 아래를 선언해서 양자화된 모델 로딩을 위한 주 패키지를 설치한다. 그리고 그 밑에 있는 블록을 실행해서 패키지를 import 할거 하고, API가 잘 작동 되는지 확인한다. 

!pip install -q requests torch bitsandbytes transformers sentencepiece accelerate

from google.colab import userdata
from huggingface_hub import login
from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer, BitsAndBytesConfig
import torch
import gc

hf_token = userdata.get('HF_TOKEN')
login(hf_token, add_to_git_credential=True)

 

서론에서 얘기했던 Llama에 접근이 승인 됐다면, 아래와 같이 선언한다. 승인이 되면 메일이 온다. 안오면 에러메세지 뜰거고 승인이 된 후면 잠잠히 있을 것이다.

LLAMA = "meta-llama/Meta-Llama-3.1-8B-Instruct"

---

Quantization Config

quant_config = BitsAndBytesConfig(
    load_in_4bit=True,  #1
    bnb_4bit_use_double_quant=True,#2
    bnb_4bit_compute_dtype=torch.bfloat16,#3
    bnb_4bit_quant_type="nf4"#4
)

위에서 선언한 Qantization config 함수에 대해 알아보고자 한다. 이 함수를 통해 model을 메모리에 로드하고, 메모리 사용을 줄일 수 있다. 이 포스팅의 핵심인 것임. 하나하나 뜯어보도록 하겠다. 1,2,3,4에 맞춰서 설명할거임. 우선 첫번째 load_in_4bit = True 라고 적혀있는데, 이건 기존 32-bit에서 4bit로 바꾼다는 의미다. 제일 초반에 4비트 혹은 8비트라고 적었는데 8비트로도 바꿀 수 있다. 그 경우는 load_in_8bit = True라고 적으면 됨. 나머지 2,3,4 모두 8비트로 선언하면 되니까 이제 이거에 관련된건 생략하도록 하겠다. 2번 parameter는 모든 weight를 두번 quantize하는 것이다. 이를 통해 메모리를 조금 더 saving할 수 있다고 한다. 3번의 의미는 4비트로 컴퓨트(계산)를 하는데 b torch 16 type으로 하겠다는 뜻이다. 마지막 4번의 의미는 4비트로 줄일 때 4bit numbers를 어떻게 처리(interpret)할 지 정하는 것이다. nf4는 Normal Float 4의 약자로 분포 기반 4bit 표현 방식이다. 정밀하고 효과적이라고 함. 덜 정교한건 fp4라고 한다. 쉽게 말해 4비트 해석할 방식 중 하나를 정한거임. 여기서 알아야 될건, 이 함수의 중요도인 메모리가 줄어드는 수준은 1>>>>>>2>>>3,4 라는 것이다. 사실 2부터 미미해진다.

---

Tokenizer 선언

Quantization Config를 선언했다면 tokenizer를 선언해야된다. 근데 이 다음 코드블록에 message가 들어가야돼서 messages를 먼저 선언해야된다.

messages = [
    {"role": "system", "content": "You are a helpful assistant"},
    {"role": "user", "content": "Tell me a joke for Data Scientists"}
  ]

tokenizer = AutoTokenizer.from_pretrained(LLAMA) #1
tokenizer.pad_token = tokenizer.eos_token #2
inputs = tokenizer.apply_chat_template(messages, return_tensors="pt").to("cuda") #3

이번에도 하나하나 뜯어서 보자면, 1번은 토크나이저 Llama로 위에서 설정한걸로 설정하는거고, 두번째는 안해도 되는건데 안하면 오류가 뜬다고해서 하는게 좋다고 한다. pad_token이라는게 보일것이다. 이건 tokenizer에 있는 prompt나 Neural Network에 fed할 때 fill-up하는 패딩(padding) 토큰을 어떻게 설정할지 결정하는 것이다. 패딩 토큰을 tokenizer의 eos_token으로 하는것임. eos token은 end-of-sequence의 줄임말인데 정리하면 2번 선언은 종료 토큰을 사용해 padding하겠다는 것이다. 3번은 llm을 하면 익숙한 것이다. 이 선언의 의미는, 메세지를 위에서 선언한 tokenizaion으로 token화한뒤, cuda(=GPU)로 파싱하겠다(넘겨주겠다)는 의미다.

model = AutoModelForCausalLM.from_pretrained(LLAMA, device_map="auto", quantization_config=quant_config)

이제 위에서 선언한 quantization_config를 통해 Llama모델의 quantizaiton을 진행한다. (device_map은 사용 가능한 gpu있을시 사용하는거고 없을시 사용안하는 parameter임.) 이후 print(model.get_memory_footprint()/ 1e6)를 하면 메모리가 얼마나 사용됐는 지 알 수 있다. 모델을 선언하면 다음과 같이 output이 나온다.

LlamaForCausalLM(
  (model): LlamaModel(
    (embed_tokens): Embedding(128256, 4096) #1
    (layers): ModuleList( #2
      (0-31): 32 x LlamaDecoderLayer(
        (self_attn): LlamaAttention( #attentionlayers. 
          (q_proj): Linear4bit(in_features=4096, out_features=4096, bias=False)
          (k_proj): Linear4bit(in_features=4096, out_features=1024, bias=False)
          (v_proj): Linear4bit(in_features=4096, out_features=1024, bias=False)
          (o_proj): Linear4bit(in_features=4096, out_features=4096, bias=False)
        )
        (mlp): LlamaMLP( # multi-layer- perceptron layers : 
          (gate_proj): Linear4bit(in_features=4096, out_features=14336, bias=False)
          (up_proj): Linear4bit(in_features=4096, out_features=14336, bias=False)
          (down_proj): Linear4bit(in_features=14336, out_features=4096, bias=False)
          (act_fn): SiLU() #3
        )
        (input_layernorm): LlamaRMSNorm((4096,), eps=1e-05)
        (post_attention_layernorm): LlamaRMSNorm((4096,), eps=1e-05)
      )
    )
    (norm): LlamaRMSNorm((4096,), eps=1e-05)
    (rotary_emb): LlamaRotaryEmbedding()
  )
  (lm_head): Linear(in_features=4096, out_features=128256, bias=False) # 4
)

 

첫번째는 단어들의 dimension이다. 두번째는 모듈 리스트로 각각은 neural Network Layer를 의미한다. 세번째는 activation함수로  llmam3.1 model에 사용된 함수다. SiLU()함수가 사용됐다. Relu를 닮은 함수인데 일단 그게 주는 아니니까 넘어가고, 네번째는 output dimension이다. 첫번째 input dimension과 매치 되는걸 확인하면 되겠다.

---

Run Model

앞에서 quantizaion을 한 model이 만들어졌고, 그 위에서는inputs으로 tokenizer를 선언했다. 이 둘을 결합하여 output을 만든다. 그리고 이렇게 사용했으면 다음 코드블록을 사용해 메모리를 청소해야된다. 그렇지 않으면 나중에 조금만 뭐해도 바로 다운되고 문제가 생긴다고 한다. 

outputs = model.generate(inputs, max_new_tokens=80)
print(tokenizer.decode(outputs[0]))

del model, inputs, tokenizer, outputs
gc.collect()
torch.cuda.empty_cache()

 

---

끝으로

끝으로 위에 있는 모든 걸 wrapping한 함수를 적겠다. 중간에 streamer는 output을 streaming을 하는 코드다.대답할때 한번에 안나오고 한글자씩 나오는 그거 맞음. 맨 아래와 같이 선언하면 위에 선언했던 메세지에 대한 대답이 나온다. 이상으로 이번 포스팅은 마치겠다.

def generate(model, messages):
  tokenizer = AutoTokenizer.from_pretrained(model)
  tokenizer.pad_token = tokenizer.eos_token
  inputs = tokenizer.apply_chat_template(messages, return_tensors="pt", add_generation_prompt=True).to("cuda")
  streamer = TextStreamer(tokenizer) # streaming
  model = AutoModelForCausalLM.from_pretrained(model, device_map="auto", quantization_config=quant_config)
  outputs = model.generate(inputs, max_new_tokens=80, streamer=streamer)
  del model, inputs, tokenizer, outputs, streamer
  gc.collect()
  torch.cuda.empty_cache()
  
  generate(LLAMA, messages)