Speed up inference

sia/llm_engine/qwq_llm_engine.py

@@ -1,9 +1,9 @@
 from pathlib import Path
 from threading import Thread
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, TextIteratorStreamer, BitsAndBytesConfig
+from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer, pipeline, BitsAndBytesConfig
 from typing import Callable, Iterator, Optional
 from xml_schema_validator import XmlLogitsProcessor
-import json
+import os
 import torch
 
 from . import LlmEngine
@@ -26,31 +26,40 @@ class QwQLlmEngine(LlmEngine):
             xml_schema_text: Optional XML schema to validate against
         """
         self._temperature = temperature
-
+                
+        # Configure 4-bit quantization for massive memory savings
         quantization_config = BitsAndBytesConfig(
             load_in_4bit=True,
             bnb_4bit_compute_dtype=torch.bfloat16,
-            bnb_4bit_quant_type="nf4",
-            bnb_4bit_use_double_quant=True
-        )
-
-        model = AutoModelForCausalLM.from_pretrained(
-            model_path,
-            return_dict=True,
-            device_map="auto",
-            use_cache=True,
-            quantization_config=quantization_config,
+            bnb_4bit_use_double_quant=True,
+            bnb_4bit_quant_type="nf4"
         )
 
+        # Load tokenizer first - this uses minimal memory
         self._tokenizer = AutoTokenizer.from_pretrained(
             model_path,
+            padding_side="left",
+            trust_remote_code=True,
         )
-
+        
+        # Load model with 4-bit quantization
+        self._model = AutoModelForCausalLM.from_pretrained(
+            model_path,
+            device_map="auto",
+            quantization_config=quantization_config,
+            torch_dtype=torch.bfloat16,
+            low_cpu_mem_usage=True,
+            trust_remote_code=True,
+        )
+        
+        # Create inference pipeline with memory-efficient settings
         self._pipeline = pipeline(
             "text-generation",
-            model=model,
+            model=self._model,
             tokenizer=self._tokenizer,
             return_full_text=False,
+            device_map="auto",
+            torch_dtype=torch.bfloat16,
         )
 
         if xml_schema_text:
@@ -95,6 +104,7 @@ class QwQLlmEngine(LlmEngine):
             "temperature": self._temperature,
             "max_new_tokens": self.token_limit(),
             "streamer": streamer,
+            "use_cache": True,
         }
         
         if self._logits_processor:
@@ -135,4 +145,4 @@ class QwQLlmEngine(LlmEngine):
         return len(self._tokenizer.encode(prompt))
 
     def token_limit(self) -> int:
-        return self._pipeline.model.config.max_position_embeddings
+        return self._pipeline.model.config.max_position_embeddings