fix: quantization

Author: zeroday0619

HakaseCore/llm/llama3.py

@@ -2,7 +2,7 @@
 import os.path
 
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 
 
 class LLama3(object):
@@ -24,11 +24,14 @@ def __init__(self, accelerate_engine: str = "cuda", debug: bool = False) -> None
                         f"{accelerate_engine} is not a valid accelerate_engine"
                     )
 
-        model = AutoModelForCausalLM.from_pretrained(self.model_id, device_map="auto")
-        self.model = torch.quantization.quantize_dynamic(
-            model,
-            {torch.nn.Linear},
-            dtype=torch.qint8,
+        bnb_config = BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_quant_type="nf4",
+            bnb_4bit_use_double_quant=True,
+            bnb_4bit_compute_dtype=torch.bfloat16,
+        )
+        self.model = AutoModelForCausalLM.from_pretrained(
+            self.model_id, quantization_config=bnb_config, device_map="auto"
         )
         self.tokenizer = AutoTokenizer.from_pretrained(
             self.model_id, add_special_tokens=True

requirements.txt

@@ -3,4 +3,5 @@ torch==2.2.2
 torchvision==0.17.2
 torchaudio==2.2.2
 transformers==4.40.1
-accelerate==0.30.0
+accelerate==0.30.0
+bitsandbytes==0.43.1