feat: use DeepSpeed ZeRO-3 for proper model sharding across GPUs

This commit is contained in:
Christian Medina
2026-07-01 22:08:37 -04:00
parent 9d8c2b6cf0
commit a655040db2
3 changed files with 31 additions and 10 deletions

View File

@@ -18,9 +18,9 @@ pip install --upgrade pip
echo "Installing training dependencies..."
pip install transformers datasets trl peft accelerate bitsandbytes deepspeed
# Run training with CUDA memory optimization
echo "Starting training with accelerate..."
# Run training with DeepSpeed ZeRO-3
echo "Starting training with DeepSpeed ZeRO-3..."
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
accelerate launch --multi_gpu --num_processes 2 train.py --config training/configs/ornith-35b-lora.yaml
torchrun --nproc_per_node=2 train.py --config training/configs/ornith-35b-lora.yaml
echo "Training completed!"