bfloat16

third_party/VideoLLaMA2/videollama2/model/__init__.py

@@ -53,7 +53,7 @@ VLLMConfigs = {
 
 
 def load_pretrained_model(model_path, model_base, model_name, load_8bit=False, load_4bit=False, device_map="auto", device="cuda", use_flash_attn=False, **kwargs):
-    print("00000000000000000000000000: ", device)
+    print("00000000000000000000000000: ", device, use_flash_attn)
     if 'token' in kwargs:
         token = kwargs['token']
     else:
@@ -181,7 +181,7 @@ def load_pretrained_model(model_path, model_base, model_name, load_8bit=False, l
         elif model_type in ['videollama2_mixtral']:
             model = Videollama2MixtralForCausalLM.from_pretrained(model_path, low_cpu_mem_usage=True, config=config, **kwargs)
         elif model_type in ['videollama2_qwen2']:
-            model = Videollama2Qwen2ForCausalLM.from_pretrained(model_path, low_cpu_mem_usage=True, config=config, **kwargs)
+            model = Videollama2Qwen2ForCausalLM.from_pretrained(model_path, low_cpu_mem_usage=True, torch_dtype=torch.bfloat16, config=config, **kwargs)
         elif model_type in ['videollama2_gemma2']:
             model = Videollama2Gemma2ForCausalLM.from_pretrained(model_path, low_cpu_mem_usage=True, config=config, **kwargs)
         elif model_type in ['videollama2_phi3']: