Files
isokl_steering_calibration/scripts/queue_sweep.sh
T

70 lines
2.2 KiB
Bash

#!/usr/bin/env bash
# Queue all sweep cells via pueue. Each cell gets its own job with why/resolve.
set -euo pipefail
cd "$(dirname "$0")/.."
ROOT="$PWD"
MODELS=(
"Qwen/Qwen3.5-0.8B"
"Qwen/Qwen2.5-0.5B-Instruct"
"meta-llama/Llama-3.2-1B-Instruct"
"Qwen/Qwen3-4B-Instruct-2507"
"Qwen/Qwen3-4B-Thinking-2507"
)
METHODS=("mean_diff" "directional_ablation" "pca")
SEEDS=(0 1 2)
WINDOWS=(20 50)
# Priority: small models first so the figure starts populating; 4B last.
prio_for_model() {
case "$1" in
*0.8B*) echo 35 ;;
*0.5B*) echo 30 ;;
*1B*) echo 20 ;;
*4B*) echo 10 ;;
*) echo 5 ;;
esac
}
n=0
for model in "${MODELS[@]}"; do
prio=$(prio_for_model "$model")
for method in "${METHODS[@]}"; do
for seed in "${SEEDS[@]}"; do
for window in "${WINDOWS[@]}"; do
run_id="$(basename "$model")_${method}_s${seed}_w${window}"
if [ -f "outputs/${run_id}/calib.json" ]; then
if env -i HOME="$HOME" PATH="$ROOT/.venv/bin:/usr/bin:/usr/local/bin" "$ROOT/.venv/bin/python" - <<PY
import json
from pathlib import Path
run = Path("outputs/${run_id}")
traj = json.loads((run / "trajectory.json").read_text()) if (run / "trajectory.json").exists() else {}
pm = json.loads((run / "pmass.json").read_text()) if (run / "pmass.json").exists() else {}
ok = (
"per_prompt_per_t_kl" in traj
and "schema" in pm
and "4.0" in traj.get("per_t_p95_kl", {})
and "4.0" in pm.get("pmass", {})
)
raise SystemExit(0 if ok else 1)
PY
then
echo "skip ${run_id} (fresh)"
continue
fi
echo "rerun ${run_id} (stale outputs)"
fi
label="why: stability of iso-KL calib for ${method} on ${model} (seed ${seed}, T=${window}); resolve: include cell in figure1 if it converges, else flag as bracket-pinned"
pueue add -w "$ROOT" -o "$prio" -l "$label" -- \
uv run --extra all python scripts/run_cell.py \
--model "$model" --method "$method" --seed "$seed" --window "$window" \
--run-id "$run_id"
n=$((n+1))
done
done
done
done
echo "queued ${n} jobs"
echo "monitor: pueue status | head -40"
echo "aggregate after queue drains: just aggregate"