Fixes for 12B

tools/gemma_train/pyproject.toml

@@ -8,17 +8,17 @@ version = "0.1.0"
 requires-python = ">=3.8"
 
 dependencies = [
-    "accelerate==1.4.0",
-    "bitsandbytes==0.45.3",
+    "accelerate==1.2.0",
+    "bitsandbytes>=0.45.0",
     "datasets==3.3.2",
     "evaluate==0.4.3",
     "llm_engine_utils @ file:///root/sia/lib/llm_engine_utils",
-    "peft==0.14.0",
+    "peft==0.13.2",
     "python-dotenv>=1.0.0",
     "sentencepiece>=0.2.0",
-    "torch>=2.4.0",
-    "transformers>=4.51.3",
-    "trl==0.15.2",
+    "torch==2.4.1",
+    "transformers==4.51.3",
+    "trl>=0.17.0",
 ]
 
 [project.scripts]

tools/gemma_train/src/gemma_train/__main__.py

@@ -19,6 +19,7 @@ def train(config: Config):
         config.model,
         token=config.api_key,
         trust_remote_code=True,
+        use_fast=False,
     )
     tokenizer.save_pretrained(config.output_dir/"tokenizer")
 
@@ -81,12 +82,13 @@ def train(config: Config):
 def merge(config: Config):
     adapted_model = AutoPeftModelForCausalLM.from_pretrained(
         config.output_dir/"lora_adapter",
-        torch_dtype=torch.float16,  # Use float16 for better compatibility
-        device_map="auto",
-        offload_folder="offload",
+        torch_dtype=torch.float16,
         token=config.api_key,
+        load_in_4bit=False,
+        load_in_8bit=False,
     )
     merged_model = adapted_model.merge_and_unload()
+
     merged_model.save_pretrained(
         config.output_dir/"merged",
         safe_serialization=True