feat: use DeepSpeed ZeRO-3 for proper model sharding across GPUs
This commit is contained in:
@@ -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!"
|
||||
Reference in New Issue
Block a user