feat: add PYTORCH_CUDA_ALLOC_CONF for memory optimization

train-on-this-server.sh

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