diff --git a/tools/gemma_train/pyproject.toml b/tools/gemma_train/pyproject.toml
index c8d7f25..44b6f7a 100644
--- a/tools/gemma_train/pyproject.toml
+++ b/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",
diff --git a/tools/gemma_train/src/gemma_train/__main__.py b/tools/gemma_train/src/gemma_train/__main__.py
index cc92719..7b2c8b0 100644
--- a/tools/gemma_train/src/gemma_train/__main__.py
+++ b/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,