Switch to Vulkan

tools/gemma_infer/pyproject.toml

@@ -9,7 +9,7 @@ requires-python = ">=3.8"
 
 dependencies = [
     "blobfile>=3.0.0",
-    "llama-cpp-python @ git+https://github.com/abetlen/llama-cpp-python.git#egg=llama-cpp-python&env=CMAKE_ARGS=-DLLAMA_BUILD=OFF",
+    "llama-cpp-python @ git+https://github.com/abetlen/llama-cpp-python.git@v0.3.16#egg=llama-cpp-python&env=CMAKE_ARGS=-DLLAMA_BUILD=OFF",
     "llm_engine_utils @ file:///root/sia/lib/llm_engine_utils",
     "protobuf>=6.0.0",
     "python-dotenv>=1.0.0",

tools/gemma_infer/src/gemma_infer/gemma_llm_engine.py

@@ -3,13 +3,15 @@ os.environ["LLAMA_CPP_LIB_PATH"] = "/usr/local/lib"
 os.environ["LD_LIBRARY_PATH"] += ":/usr/local/lib"
 os.chdir("/usr/local/lib")
 
-from llama_cpp import Llama, LogitsProcessorList
+from llama_cpp import Llama, LogitsProcessorList, llama_cpp
 from llm_engine_utils import LlmEngine
 from pathlib import Path
 from transformers import AutoTokenizer
 from typing import Iterator
 from xml_schema_validator import LlamaCppLogitsProcessor
 
+llama_cpp._lib.ggml_backend_load_all()
+
 class GemmaLlmEngine(LlmEngine):
     def __init__(
             self,

tools/gemma_train/src/gemma_train/__main__.py

@@ -1,7 +1,7 @@
 from llm_engine_utils.dataset import Dataset
 from pathlib import Path
 from peft import LoraConfig, AutoPeftModelForCausalLM
-from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, TrainingArguments
+from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments
 from trl import SFTTrainer
 import os
 import sys
@@ -23,19 +23,13 @@ def train(config: Config):
     )
     tokenizer.save_pretrained(config.output_dir/"tokenizer")
 
-    bnb_config = BitsAndBytesConfig(
-        load_in_4bit=True,
-        bnb_4bit_use_double_quant=True,
-        bnb_4bit_quant_type="nf4",
-        bnb_4bit_compute_dtype=torch.bfloat16
-    )
-
     model = AutoModelForCausalLM.from_pretrained(
         config.model,
-        quantization_config=bnb_config,
-        device_map="auto",
+        torch_dtype=torch.float32,
+        device_map="cpu",
         token=config.api_key,
         attn_implementation='eager',
+        low_cpu_mem_usage=True,
     )
 
     dataset = Dataset(config.config_path)
@@ -57,14 +51,15 @@ def train(config: Config):
         warmup_steps=1,
         max_steps=1,
         learning_rate=1e-3,
-        fp16=True,
+        fp16=False,
         logging_steps=1,
         save_strategy="steps",
         save_steps=1,
         output_dir=config.output_dir/"lora",
-        optim="paged_adamw_8bit",
+        optim="adamw_torch",
         seed=42,
         group_by_length=True,
+        use_cpu=True,
     )
 
     trainer = SFTTrainer(