Update libra/model/builder.py

libra/model/builder.py

@@ -24,6 +24,9 @@ from libra.constants import DEFAULT_IMAGE_PATCH_TOKEN, DEFAULT_IM_START_TOKEN, D
 
 
 def load_pretrained_model(model_path, model_base, model_name, device="cpu"):
+    quantization_config = BitsAndBytesConfig(
+        load_in_8bit=True
+    )
     device_map = {"": device}
     kwargs = {
         "device_map": device_map,
@@ -81,14 +84,14 @@ def load_pretrained_model(model_path, model_base, model_name, device="cpu"):
             model.load_state_dict(mm_projector_weights, strict=False)
         else:
             tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False)
-            model = LibraLlamaForCausalLM.from_pretrained(model_path, low_cpu_mem_usage=True, **kwargs)
+            model = LibraLlamaForCausalLM.from_pretrained(model_path, quantization_config=quantization_config, low_cpu_mem_usage=True, **kwargs)
     else:
         # Load language model
         if model_base is not None:
             # PEFT model
             from peft import PeftModel
             tokenizer = AutoTokenizer.from_pretrained(model_base, use_fast=False)
-            model = AutoModelForCausalLM.from_pretrained(model_base, torch_dtype=torch.bfloat16, low_cpu_mem_usage=True, device_map="auto")
+            model = AutoModelForCausalLM.from_pretrained(model_base, quantization_config=quantization_config, torch_dtype=torch.bfloat16, low_cpu_mem_usage=True, device_map="auto")
             print(f"Loading LoRA weights from {model_path}")
             model = PeftModel.from_pretrained(model, model_path)
             print(f"Merging weights")
@@ -98,7 +101,7 @@ def load_pretrained_model(model_path, model_base, model_name, device="cpu"):
         else:
             use_fast = False
             tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False)
-            model = AutoModelForCausalLM.from_pretrained(model_path, low_cpu_mem_usage=True, **kwargs)
+            model = AutoModelForCausalLM.from_pretrained(model_path, quantization_config=quantization_config, low_cpu_mem_usage=True, **kwargs)
 
     image_processor = None