Update libra/model/builder.py

libra/model/builder.py

@@ -26,14 +26,14 @@ 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_4bit=True,  
-    bnb_4bit_compute_dtype=torch.bfloat16,  
+    bnb_4bit_compute_dtype=torch.float16,  
     bnb_4bit_use_double_quant=True,  
     bnb_4bit_quant_type='nf4' 
     )
     device_map = {"": device}
     kwargs = {
         "device_map": device_map,
-        "torch_dtype": torch.bfloat16 
+        "torch_dtype": torch.float16 
     }
 
         
@@ -83,7 +83,7 @@ def load_pretrained_model(model_path, model_base, model_name, device="cpu"):
             model = LibraLlamaForCausalLM.from_pretrained(model_base, low_cpu_mem_usage=True, config=cfg_pretrained, **kwargs)
 
             mm_projector_weights = torch.load(os.path.join(model_path, 'mm_projector.bin'), map_location='cpu')
-            mm_projector_weights = {k: v.to(torch.bfloat16) for k, v in mm_projector_weights.items()}
+            mm_projector_weights = {k: v.to(torch.float16) for k, v in mm_projector_weights.items()}
             model.load_state_dict(mm_projector_weights, strict=False)
         else:
             tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False)
@@ -94,13 +94,13 @@ def load_pretrained_model(model_path, model_base, model_name, device="cpu"):
             # PEFT model
             from peft import PeftModel
             tokenizer = AutoTokenizer.from_pretrained(model_base, use_fast=False)
-            model = AutoModelForCausalLM.from_pretrained(model_base, quantization_config=quantization_config, 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.float16, 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")
             model = model.merge_and_unload()
             print('Convert to FP16...')
-            model.to(torch.bfloat16)
+            model.to(torch.float16)
         else:
             use_fast = False
             tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False)