#!/usr/bin/env bash
set -euo pipefail

export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
export LD_LIBRARY_PATH=$CONDA_PREFIX/lib:$CONDA_PREFIX/targets/x86_64-linux/lib:$CONDA_PREFIX/lib/python3.12/site-packages/nvidia/cu13/lib:$LD_LIBRARY_PATH

MODEL_NAME="Qwen/Qwen3.5-9B"
TRAIN_FILE="data/train_dpo_clean.jsonl"
DEV_FILE="data/dev_dpo_clean.jsonl"
OUTPUT_DIR="outputs/qwen35_dpo_run1"
REF_LOGPS_CACHE_DIR="cache/ref_logps_qwen35"

torchrun --nproc_per_node=2 scripts/train_dpo.py \
  --model-name "${MODEL_NAME}" \
  --train-file "${TRAIN_FILE}" \
  --dev-file "${DEV_FILE}" \
  --output-dir "${OUTPUT_DIR}" \
  --ref-logps-cache-dir "${REF_LOGPS_CACHE_DIR}" \
  --learning-rate 5e-7 \
  --num-train-epochs 1.0 \
  --per-device-train-batch-size 1 \
  --per-device-eval-batch-size 1 \
  --gradient-accumulation-steps 16 \
  --max-length 6144 \
  --max-prompt-length 5632 \
  --precompute-ref-log-probs \
  --eval-steps 100 \
  --save-steps 100 \
  --logging-steps 10 \
  --beta 0.1 \
  --lora-r 64 \
  --lora-alpha 128 \
  --lora-dropout 0.05 \
  --bf16