#!/bin/bash
# Training script for Cyron LoRA on this server (2x RTX 5090)

set -e

# Unload AI model to free GPU memory
echo "Unloading AI model..."
ai none

echo "=== Cyron LoRA Training Setup ==="

# Create venv
python3 -m venv venv
source venv/bin/activate
pip install --upgrade pip

# Install training dependencies (PyTorch already installed)
echo "Installing training dependencies..."
pip install transformers datasets trl peft accelerate bitsandbytes deepspeed

# Run training with DeepSpeed ZeRO-3
echo "Starting training with DeepSpeed ZeRO-3..."
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
torchrun --nproc_per_node=2 train.py --config training/configs/ornith-35b-lora.yaml

echo "Training completed!"