Quantization - Giảm dung lượng RAM cho LLM

Giới thiệu

Các mô hình ngôn ngữ lớn (LLM) đã trở nên phổ biến trong những năm gần đây do khả năng tạo văn bản, dịch ngôn ngữ và trả lời câu hỏi. Tuy nhiên, các mô hình này thường đòi hỏi rất nhiều bộ nhớ RAM và VRAM (GPU), hơn nhiều máy tính thông dụng.

Trong bài viết này, tôi sẽ trình bày cách sử dụng Quantization để giảm dung lượng bộ nhớ của LLM. Trên mạng tôi tìm thấy rất nhiều mô hình LLM muốn sử dụng, nhưng trên máy tính không có đủ VRAM. Tôi đã nghiên cứu cách để có thể chạy được LLM trên máy tính cá nhân để tiện lập trình. Tôi sẽ so sánh hiệu suất của mô hình chưa được quantize với mô hình đã được lượng hóa để xem quantize có tác động gì đến độ chính xác.

Lưu ý: bài viết này hướng dẫn cách quantize model cho những LLM mà chưa có phiên bản quantize (4bit, 8bit). Nếu như LLM có phiên bản quantize rồi, thì nên sử dụng trực tiếp phiên bản đó. Ví dụ: nếu muốn dùng Alpaca quantize 4bit, nên dùng trực tiếp  https://huggingface.co/4bit/alpaca-7b-native-4bit

Cài đặt

Để thực hiện theo bài viết này, ngoài pytorch, bạn cần cài đặt các gói sau:

pip install modelscope
pip install -q -U bitsandbytes
pip install -q -U git+https://github.com/huggingface/transformers.git
pip install -q -U git+https://github.com/huggingface/peft.git
pip install -q -U git+https://github.com/huggingface/accelerate.git
pip install sentencepiece

Kiểm tra GPU, nếu có thì LLM sẽ chạy nhanh hơn:

import torch
print(f"Is Cuda Available: {torch.cuda.is_available()}")    
for device_index in range(torch.cuda.device_count()):
    print(f"Device Index: {device_index}")
    print(f"Device Name: {torch.cuda.get_device_name(device_index)}")
    
# Is Cuda Available: True
# Device Index: 0
# Device Name: NVIDIA GeForce GTX 1060 with Max-Q Design -> Có GPU
    

Quantization

Model chúng ta sẽ sử dụng là Bloomz, một họ các mô hình ngôn ngữ lớn được phát triển bởi BigScience dựa vào model đa ngôn ngữ Bloom. Nó có khả năng thực hiện mệnh lệnh của con người với hàng chục ngôn ngữ. Bloomz có nhiều kích thước, từ 300 triệu tham số đến 176 tỷ tham số, chúng ta sẽ dùng model 1.1 tỷ tham số.

Quantization được sử dụng để giảm dung lượng bộ nhớ của LLM. Quantization liên quan đến việc giảm số bit được sử dụng để lưu trữ trọng số và độ lệch của mô hình. Quantize thông dụng hiện giờ là 4bit và 8bit.  Quantize bằng 4bit có thể làm giảm kích thước mô hình lên đến 4 lần, nhưng nó cũng có thể làm giảm độ chính xác của mô hình, như sẽ thấy ở phần thử nghiệm.

Sau đây là code để tạo model Bloomz kèm theo comment về setting Quantize:

from transformers import BloomForCausalLM, AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

device = torch.device("cuda:0" nếu torch.cuda.is_available() else "cpu")

def create_model(quantize_4bit=False, use_fast=False):
    # Model BloomZ. Có nhiều model lớn tốt hơn nhưng sẽ không đủ bộ nhớ thử được trên máy
    model_path = "bigscience/bloomz-1b1"

    # Thực hiện quantize 4-bit nếu được yêu cầu
    if quantize_4bit:
        # Cấu hình các cài đặt lượng hóa
        quantization_config = BitsAndBytesConfig(
            load_in_4bit=True,  # Tải trọng số được lượng hóa trước theo định dạng 4 bit
            bnb_4bit_quant_type="nf4",  # Sử dụng loại lượng hóa "nf4" cho trọng số 4 bit
            bnb_4bit_compute_dtype=torch.bfloat16,  # Sử dụng torch.bfloat16 cho các phép tính trung gian
            bnb_4bit_use_double_quant=True,  # Sử dụng độ chính xác kép để lượng hóa kích hoạt
        )

        # Áp dụng lượng hóa cho mô hình
        print("Đang lượng hóa mô hình xuống 4 bit...")
        model = AutoModelForCausalLM.from_pretrained(
            model_path,
            quantization_config=quantization_config,
            device_map="auto", # Có thể một phần model RAM, một phần VRAM
            trust_remote_code=True,
        )
        print("Quantize đã hoàn thành.")
    else:
        # Sử dụng mô hình gốc mà không cần lượng hóa
        model = AutoModelForCausalLM.from_pretrained(
            model_path,
            device_map="auto",
            trust_remote_code=True,
        )

    # Chuyển đổi kích thước mô hình sang gigabyte
    print(f"Kích thước mô hình: {round(model.get_memory_footprint() / 1e9, 2)} GB")

    # Đặt mô hình vào chế độ đánh giá
    model.eval()

    # Chuyển mô hình sang thiết bị được xác định trước đó (GPU hoặc CPU) 
    if not quantize_4bit:
        model.to(device)

    # Trả về mô hình và tokenizer đã tạo
    return model, tokenizer

Tạo Prompt

Chúng ta viết code dưới để tạo prompt cho LLM để dịch ngôn ngữ này sang ngôn ngữ khác:

prompt = """

Here is the input paragraph in {input_lang}:
```
{input}
```

Translate input paragraph to {output_lang}:
"""

Còn đây là một số câu (dựa theo kịch phim) bằng tiếng Anh để thử dịch sang tiếng Việt :

paragraph1 = """
I am a software developer.
"""


paragraph2 = """
Bedevere: What makes you think she is a witch?
"""

paragraph3 = """
King Arthur: I am your king.
Peasant Woman: Well, I didn't vote for you.
King Arthur: You don't vote for kings.
"""
input_docs = [paragraph1, paragraph2, paragraph3]

Thử nghiệm

Dưới là code để sử dụng model dịch văn bản

import time

def translate(model, tokenizer, input_docs, input_lang, output_lang, quantize_4bit=False):

  output_docs = []  # Danh sách các văn bản đã dịch
  translation_times = []  # Danh sách thời gian dịch cho từng văn bản

  for index, input_doc in enumerate(input_docs):  
    start_time = time.time()

    # Tạo prompt cho đầu vào dịch, bao gồm thông tin về văn bản đầu vào, ngôn ngữ đầu vào và ngôn ngữ mong muốn
    input_doc = prompt.format(input=input_doc, input_lang=input_lang, output_lang=output_lang)

    # Chuyển đổi văn bản đầu vào thành tensor bằng tokenizer
    inputs = tokenizer(input_doc, return_tensors="pt")

    if not quantize_4bit:
        inputs = inputs.to(device)

    # Dịch văn bản bằng mô hình Transformer
    
    generate_ids = model.generate(
        inputs['input_ids'],
        # attention_mask=inputs['attention_mask'],
        # Giữ 3 kết quả dịch khác nhau để chọn kết quả dịch có xác xuất cao nhất
        num_beams=3,
        # Chiều dài tối đa của đoạn dịch là 1.25 lần chiều dài đầu vào
        max_length=len(input_doc) * 1.25,
        # Dừng sớm nếu tìm thấy bản dịch hợp lý và chỉ trả về một bản dịch

        early_stopping=True,
        num_return_sequences=1,
        do_sample=False,
    )

    # Chuyển đổi tensor thành văn bản đã dịch
    decoded = tokenizer.decode(generate_ids[0], skip_special_tokens=True)

    output_docs.append(f"{decoded}")

    end_time = time.time()

    # Tính toán thời gian dịch cho văn bản này
    translation_times.append(end_time - start_time)

  return output_docs, translation_times

Model không dùng quantize

import numpy as np
normal_model, tokenizer = create_model(quantize_4bit=False, use_fast=False)
# Model size: 4.26 GB

Để ý: RAM đang dùng là 4.26GB. Bây giờ chúng ta sẽ dịch tiếng anh sang tiếng Việt

outputs, trans_times = translate(normal_model, tokenizer, input_docs, "English", "Vietnamese")
print_translations(input_docs, outputs, trans_times)
print(f"Mean translations times: {np.array(trans_times).mean()}")

Output sẽ là:

Here is the input paragraph in English:
```

I am a software developer.

```

Translate input paragraph to Vietnamese:
Tôi là một kỹ sư phần mềm.
>>> Translation time: 1.0163302421569824
--------------------------------------------------


Here is the input paragraph in English:
```

Bedevere: What makes you think she is a witch?

```

Translate input paragraph to Vietnamese:
Bedevere: Điều gì khiến bạn nghĩ cô ấy là một con ma?
>>> Translation time: 0.8176043033599854
--------------------------------------------------


Here is the input paragraph in English:
```

King Arthur: I am your king.
Peasant Woman: Well, I didn't vote for you.
King Arthur: You don't vote for kings.

```

Translate input paragraph to Vietnamese:
King Arthur: Tôi là vua của bạn. Nữ công chúa: Tôi không ủng hộ bạn. Vua Arthur: Bạn không ủng hộ vua.
>>> Translation time: 1.32578706741333
--------------------------------------------------
Mean translations times: 1.0532405376434326

Ta thấy model dịch được, nếu sử dụng GPU mất trung bình 1 giây. Tuy nhiên, vì model có tham số nhỏ, nên vẫn chưa chính xác. Ví dụ như "witch" dịch thành "ma" thay bằng phù thủy, "vote" dịch thành ủng hộ thay bằng "bầu (chọn)".

Model sử dụng quantize

Tiếp theo ta tạo model sử dụng phương pháp quantize

# Xóa model cũ
del normal_model
torch.cuda.empty_cache()

quantized_model, tokenizer = create_model(quantize_4bit=True, use_fast=True)
# WARNING:root:Some parameters are on the meta device device because they were offloaded to the .
# WARNING:root:Some parameters are on the meta device device because they were offloaded to the cpu/disk.
# Model size: 1.11 GB

Để ý dung lượng giảm từ 4.26GB thành 1.11GB, gần 1/4. Đồng thời một số tham số của model được lưu ở CPU RAM hoặc ổ cứng thay bằng chỉ ở GPU VRAM. Tiếp theo chúng ta dịch từ tiếng Anh sang tiếng Việt:


Here is the input paragraph in English:
```

I am a software developer.

```

Translate input paragraph to Vietnamese:
Tôi là một kỹ sư phần mềm.
>>> Translation time: 1.862919807434082
--------------------------------------------------


Here is the input paragraph in English:
```

Bedevere: What makes you think she is a witch?

```

Translate input paragraph to Vietnamese:
Bedevere: Cái gì khiến bạn nghĩ cô ấy là một con ma?
>>> Translation time: 3.2726874351501465
--------------------------------------------------


Here is the input paragraph in English:
```

King Arthur: I am your king.
Peasant Woman: Well, I didn't vote for you.
King Arthur: You don't vote for kings.

```

Translate input paragraph to Vietnamese:
King Arthur: Tôi là vua của bạn. Người phụ nữ nghèo: Tôi không bầu bạn với vua.
>>> Translation time: 5.169569969177246
--------------------------------------------------
Mean translations times: 3.435059070587158

Điểm một chúng ta thấy là thời gian dịch tăng lên. Sử dụng quantize đôi khi sẽ phải lưu một số tham số trong CPU RAM hoặc ổ cứng, và phải chuyển các thông  tin nên nó có thể chậm hơn. Nếu model rất phức tạp, quantize sau đó lại dequantize sẽ gây mất nhiều thời gian hơn.

Điểm thứ hai là ở đoạn văn thứ 3, model dịch thiếu câu thứ 3 hiếu của vua Arthur. Như vậy, độ chính xác của model sau khi dùng quantize cũng có thể giảm.

Kết luận

Từ thử nghiệm trên, quantization qua BitsAndBytes rất tốt nếu bạn muốn giảm RAM, VRAM mà model LLM sử dụng. Tuy nhiên cũng nên lưu ý một số điểm yếu của quantization như độ chính xác có thể giảm, và tùy model tốc độ cũng có thể giảm.

Tham khảo