Switch to Vulkan

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(