cmedina/agenx-lora-training
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Christian Medina
2026-06-30 18:45:14 -04:00
parent 5ad5518ca6
commit 75aeeeb20a
2 changed files with 5 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
train-on-this-server.sh
View File

@@ -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 "Press Ctrl+C to stop (model will be saved at checkpoint)."
echo "" 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 ""
echo "=== Training Complete ===" echo "=== Training Complete ==="

5
training/scripts/train.py
View File

@@ -36,10 +36,11 @@ def train(config_path):
# Load model - let the model's own quantization config handle it # Load model - let the model's own quantization config handle it
# (Ornith uses CompressedTensors, not BitsAndBytes) # (Ornith uses CompressedTensors, not BitsAndBytes)
# Load on CPU first, then DeepSpeed will distribute
model = AutoModelForCausalLM.from_pretrained( model = AutoModelForCausalLM.from_pretrained(
config["base_model"], config["base_model"],
torch_dtype=torch.bfloat16, dtype=torch.bfloat16,
device_map="auto", # Let transformers distribute across GPUs device_map="cpu", # Load on CPU, DeepSpeed distributes
) )
# Add LoRA # Add LoRA