local reward model

1 files changed, 31 insertions, 0 deletions

diff --git a/scripts/test_armo_reward.sh b/scripts/test_armo_reward.sh
new file mode 100644
index 0000000..18004c7
--- /dev/null
+++ b/scripts/test_armo_reward.sh
@@ -0,0 +1,31 @@
+#!/bin/bash
+#SBATCH --job-name=test_armo
+#SBATCH --account=bfqt-delta-gpu
+#SBATCH --partition=gpuA100x4
+#SBATCH --nodes=1
+#SBATCH --ntasks=1
+#SBATCH --cpus-per-task=8
+#SBATCH --gres=gpu:nvidia_a100:1
+#SBATCH --mem=64G
+#SBATCH --time=0:15:00
+#SBATCH --output=test_armo_%j.out
+#SBATCH --error=test_armo_%j.err
+
+# Test ArmoRM reward model
+# Requires ~16GB GPU memory for 8B model
+
+cd /projects/bfqt/users/yurenh2/ml-projects/personalization-user-model
+source /u/yurenh2/miniforge3/etc/profile.d/conda.sh
+conda activate eval
+
+export HF_HOME=/projects/bfqt/users/yurenh2/hf_cache/huggingface
+export PYTHONPATH="${PWD}/src:${PYTHONPATH}"
+
+echo "=== Testing ArmoRM Reward Model ==="
+echo "GPU: $(nvidia-smi --query-gpu=name --format=csv,noheader)"
+echo ""
+
+python scripts/test_armo_reward.py --device cuda
+
+echo ""
+echo "=== Test Complete ==="