nit

allenai · Oct 24, 2024 · b37cf0c · b37cf0c
1 parent 553123f
commit b37cf0c
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/scripts/eval/oe-eval.sh b/scripts/eval/oe-eval.sh
@@ -100,8 +100,9 @@ TASKS=(
 MODEL_TYPE="--model-type vllm"
 BATCH_SIZE_VLLM=10000
 BATCH_SIZE_OTHER=1
-GPU_COUNT=1
-GPU_COUNT_OTHER=2
+# Set GPU_COUNT and GPU_COUNT_OTHER based on NUM_GPUS
+GPU_COUNT="$NUM_GPUS"
+GPU_COUNT_OTHER=$((NUM_GPUS * 2)) 
 MODEL_TYPE_OTHER=""
 
 for TASK in "${TASKS[@]}"; do
@@ -113,7 +114,7 @@ for TASK in "${TASKS[@]}"; do
     else
         BATCH_SIZE=$BATCH_SIZE_VLLM
         MODEL_TYPE="--model-type vllm"
-        GPU_COUNT=1
+        GPU_COUNT=$GPU_COUNT
     fi
 
     python oe-eval-internal/oe_eval/launch.py --model "$MODEL_NAME" --beaker-workspace "ai2/tulu-3-results" --beaker-budget ai2/oe-adapt --task "$TASK" $MODEL_TYPE --batch-size "$BATCH_SIZE" --model-args "{\"model_path\":\"${MODEL_LOCATION}\", \"max_length\": ${MAX_LENGTH}}" ${HF_UPLOAD_ARG} --gpus "$GPU_COUNT" --gantry-args '{"env-secret": "OPENAI_API_KEY=openai_api_key"}' ${REVISION_ARG}