From 75aeeeb20a989bded00a7fbcd9155d30f9657bad Mon Sep 17 00:00:00 2001
From: Christian Medina <37550954+cmedinasoriano@users.noreply.github.com>
Date: Tue, 30 Jun 2026 18:45:14 -0400
Subject: [PATCH] fix: use torchrun for distributed training, load model on CPU

---
 train-on-this-server.sh   | 3 ++-
 training/scripts/train.py | 5 +++--
 2 files changed, 5 insertions(+), 3 deletions(-)

diff --git a/train-on-this-server.sh b/train-on-this-server.sh
index 6f4c654..72de205 100644
--- a/train-on-this-server.sh
+++ b/train-on-this-server.sh
@@ -68,7 +68,8 @@ echo "Training will take 6-24 hours depending on GPU."
 echo "Press Ctrl+C to stop (model will be saved at checkpoint)."
 echo ""
 
-python3 training/scripts/train.py --config training/configs/llama2-7b-lora.yaml
+# Use torchrun for distributed training (2 GPUs)
+torchrun --nproc_per_node=2 training/scripts/train.py --config training/configs/llama2-7b-lora.yaml
 
 echo ""
 echo "=== Training Complete ==="
diff --git a/training/scripts/train.py b/training/scripts/train.py
index 64183fc..3d6932d 100755
--- a/training/scripts/train.py
+++ b/training/scripts/train.py
@@ -36,10 +36,11 @@ def train(config_path):
     
     # Load model - let the model's own quantization config handle it
     # (Ornith uses CompressedTensors, not BitsAndBytes)
+    # Load on CPU first, then DeepSpeed will distribute
     model = AutoModelForCausalLM.from_pretrained(
         config["base_model"],
-        torch_dtype=torch.bfloat16,
-        device_map="auto",  # Let transformers distribute across GPUs
+        dtype=torch.bfloat16,
+        device_map="cpu",  # Load on CPU, DeepSpeed distributes
     )
     
     # Add LoRA