diff --git a/training/configs/llama2-7b-lora.yaml b/training/configs/llama2-7b-lora.yaml
index a4100a0..3cd82b8 100644
--- a/training/configs/llama2-7b-lora.yaml
+++ b/training/configs/llama2-7b-lora.yaml
@@ -1,7 +1,7 @@
 # LoRA Training Configuration for Llama-2-7b
 # Dataset: cyron_summary_lora_dataset (20k examples)
 
-base_model: meta-llama/Llama-2-7b-hf
+base_model: deepreinforce-ai/Ornith-1.0-35B-FP8
 model_type: LlamaForCausalLM
 tokenizer_type: LlamaTokenizer
 
diff --git a/training/scripts/train.py b/training/scripts/train.py
index 5754c23..76d27e9 100755
--- a/training/scripts/train.py
+++ b/training/scripts/train.py
@@ -39,10 +39,13 @@ def train(config_path):
         use_nested_quant=config.get("use_nested_quant", False),
     )
     
+    # Use all available GPUs
+    device_map = "auto" if torch.cuda.device_count() == 1 else "balanced"
+    
     model = AutoModelForCausalLM.from_pretrained(
         config["base_model"],
         quantization_config=bnb_config,
-        device_map="auto",
+        device_map=device_map,
     )
     model = prepare_model_for_kbit_training(model)