wip deepseek train

sia/__main__.py

@@ -67,7 +67,7 @@ class Main:
                 config.deepseek_model,
                 config.deepseek_temperature,
                 config.deepseek_token_limit,
-                config.hf_api_key,  # Use the existing HF API key
+                config.hf_api_key,
             )
 
         if not self._llms:

sia/llm_engine/deepseek_llm_engine.py

@@ -1,6 +1,6 @@
 from typing import Callable, Iterator, Optional
 import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
+from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer, BitsAndBytesConfig
 from threading import Thread
 from pathlib import Path
 
@@ -43,17 +43,28 @@ class DeepSeekLlmEngine(LlmEngine):
         # Set padding token to avoid warnings
         if self._tokenizer.pad_token is None:
             self._tokenizer.pad_token = self._tokenizer.eos_token
-            
-        # Load model with 4-bit quantization by default
+        
+        # Configure 4-bit quantization with CPU offloading
+        quantization_config = BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_compute_dtype=torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16,
+            bnb_4bit_use_double_quant=True,
+            bnb_4bit_quant_type="nf4",
+            llm_int8_enable_fp32_cpu_offload=True
+        )
+        
+        # Configure device map for efficient memory usage
+        # "auto" with the proper quantization config will handle the memory constraints
         self._device_map = "auto"
         
+        # Load model with quantization config
         self._model = AutoModelForCausalLM.from_pretrained(
             self._model_path,
             return_dict=True,
             low_cpu_mem_usage=True,
             trust_remote_code=True,
             device_map=self._device_map,
-            load_in_4bit=True, 
+            quantization_config=quantization_config,
             torch_dtype=torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16,
             token=api_key,
         )