Improved gemma training

tools/gemma_train/pyproject.toml

@@ -12,6 +12,7 @@ dependencies = [
     "bitsandbytes>=0.45.0",
     "datasets==3.3.2",
     "evaluate==0.4.3",
+    "kernels>=0.11.1",
     "llm_engine_utils @ file:///root/sia/lib/llm_engine_utils",
     "peft==0.13.2",
     "python-dotenv>=1.0.0",

tools/gemma_train/src/gemma_train/__main__.py

@@ -37,8 +37,8 @@ def train(config: Config):
     dataset = dataset.to_transformers_dataset(tokenizer)
 
     lora_config = LoraConfig(
-        r=4,
-        lora_alpha=4,
+        r=12,
+        lora_alpha=24,
         target_modules=["q_proj", "o_proj", "k_proj", "v_proj", "gate_proj", "up_proj", "down_proj"],
         lora_dropout=0.05,
         bias="none",
@@ -49,12 +49,12 @@ def train(config: Config):
         per_device_train_batch_size=1,
         gradient_accumulation_steps=4,
         warmup_steps=1,
-        max_steps=1,
+        num_train_epochs=2,
         learning_rate=1e-3,
         fp16=False,
-        logging_steps=1,
+        logging_steps=5,
         save_strategy="steps",
-        save_steps=1,
+        save_steps=10,
         output_dir=config.output_dir/"lora",
         optim="adamw_torch",
         seed=42,