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

2026-06-30 18:45:14 -04:00
parent 5ad5518ca6
commit 75aeeeb20a
2 changed files with 5 additions and 3 deletions
--- 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