Fix quantized inference (#302)

Fixed possible mismatches caused by high version dependencies.
ymcui · Sep 21, 2023 · 09eadc6 · 09eadc6
1 parent ba4e228
commit 09eadc6
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 16 deletions.
diff --git a/scripts/inference/gradio_demo.py b/scripts/inference/gradio_demo.py
@@ -166,17 +166,21 @@ def setup():
             if args.lora_model is None:
                 args.tokenizer_path = args.base_model
         tokenizer = LlamaTokenizer.from_pretrained(args.tokenizer_path, legacy=True)
+        if args.load_in_4bit or args.load_in_8bit:
+            quantization_config = BitsAndBytesConfig(
+                load_in_4bit=args.load_in_4bit,
+                load_in_8bit=args.load_in_8bit,
+                bnb_4bit_compute_dtype=load_type,
+            )
 
         base_model = LlamaForCausalLM.from_pretrained(
             args.base_model,
             torch_dtype=load_type,
             low_cpu_mem_usage=True,
             device_map='auto',
-            quantization_config=BitsAndBytesConfig(
-                load_in_4bit=args.load_in_4bit,
-                load_in_8bit=args.load_in_8bit,
-                bnb_4bit_compute_dtype=load_type
-            )
+            load_in_4bit=args.load_in_4bit,
+            load_in_8bit=args.load_in_8bit,
+            quantization_config=quantization_config if (args.load_in_4bit or args.load_in_8bit) else None
         )
 
         model_vocab_size = base_model.get_input_embeddings().weight.size(0)

diff --git a/scripts/inference/inference_hf.py b/scripts/inference/inference_hf.py
@@ -110,17 +110,21 @@ def generate_prompt(instruction, system_prompt=DEFAULT_SYSTEM_PROMPT):
         tokenizer = LlamaTokenizer.from_pretrained(args.tokenizer_path, legacy=True)
     else:
         tokenizer = LlamaTokenizer.from_pretrained(args.tokenizer_path, legacy=True)
+        if args.load_in_4bit or args.load_in_8bit:
+            quantization_config = BitsAndBytesConfig(
+                load_in_4bit=args.load_in_4bit,
+                load_in_8bit=args.load_in_8bit,
+                bnb_4bit_compute_dtype=load_type,
+            )
 
         base_model = LlamaForCausalLM.from_pretrained(
             args.base_model,
             torch_dtype=load_type,
             low_cpu_mem_usage=True,
             device_map='auto',
-            quantization_config=BitsAndBytesConfig(
-                load_in_4bit=args.load_in_4bit,
-                load_in_8bit=args.load_in_8bit,
-                bnb_4bit_compute_dtype=load_type
-            )
+            load_in_4bit=args.load_in_4bit,
+            load_in_8bit=args.load_in_8bit,
+            quantization_config=quantization_config if (args.load_in_4bit or args.load_in_8bit) else None
             )
 
         model_vocab_size = base_model.get_input_embeddings().weight.size(0)

diff --git a/scripts/openai_server_demo/openai_api_server.py b/scripts/openai_server_demo/openai_api_server.py
@@ -69,17 +69,20 @@
     if args.lora_model is None:
         args.tokenizer_path = args.base_model
 tokenizer = LlamaTokenizer.from_pretrained(args.tokenizer_path, legacy=True)
-
+if args.load_in_4bit or args.load_in_8bit:
+    quantization_config = BitsAndBytesConfig(
+        load_in_4bit=args.load_in_4bit,
+        load_in_8bit=args.load_in_8bit,
+        bnb_4bit_compute_dtype=load_type,
+    )
 base_model = LlamaForCausalLM.from_pretrained(
     args.base_model,
     torch_dtype=load_type,
     low_cpu_mem_usage=True,
     device_map='auto' if not args.only_cpu else None,
-    quantization_config=BitsAndBytesConfig(
-        load_in_4bit=args.load_in_4bit,
-        load_in_8bit=args.load_in_8bit,
-        bnb_4bit_compute_dtype=load_type
-    )
+    load_in_4bit=args.load_in_4bit,
+    load_in_8bit=args.load_in_8bit,
+    quantization_config=quantization_config if (args.load_in_4bit or args.load_in_8bit) else None
 )
 
 model_vocab_size = base_model.get_input_embeddings().weight.size(0)