run_static_sweep_vllm.sh

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

# Static branch sweep using vLLM backend.
# Usage:
#   bash run_static_sweep_vllm.sh CUDA_DEVICE MODEL_PATH SAMPLES "DATASETS"
# Example:
#   bash run_static_sweep_vllm.sh 2 Qwen/Qwen3-4B 100 "aime_2024 aime_2025"

DEVICE=${1:-0} # CUDA_VISIBLE_DEVICES
MODEL_PATH=${2:-Qwen/Qwen3-4B}
SAMPLES=${3:-100}
DATASETS=${4:-"aime_2024 aime_2025"}

BRANCH_SET=(16 32 64)

echo "Model: ${MODEL_PATH}"
echo "Datasets: ${DATASETS}"
echo "Samples per dataset: ${SAMPLES}"
echo "Branches: ${BRANCH_SET[*]}"

cd "$(dirname "$0")"/ || exit 1

for B in "${BRANCH_SET[@]}"; do
  echo "=== Running static sweep (vLLM): branches=${B} ==="
  CUDA_VISIBLE_DEVICES=${DEVICE} \
  PYTHONPATH=src \
  python -u src/run_math_benchmarks.py \
    --model-path "${MODEL_PATH}" \
    --model-type vllm \
    --datasets ${DATASETS} \
    --max-samples ${SAMPLES} \
    --mode static \
    --branches ${B} \
    --max-parallel-branches ${B} \
    --max-new-tokens 32768 \
    --temperature 0.7 \
    --top-p 0.95 \
    --verbose 2>&1 | tee -a iclr_results/static_sweep_vllm_b${B}.log
done

echo "Static vLLM sweep completed. See iclr_results/ for organized outputs."