fix: use torchrun for distributed training, load model on CPU

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 ==="

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