65 lines
2.0 KiB
Bash
65 lines
2.0 KiB
Bash
|
#!/usr/bin/env bash
|
||
|
|
set -euo pipefail
|
||
|
|
|
||
|
|
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||
|
|
RECIPE="$(cd "$SCRIPT_DIR/.." && pwd)"
|
||
|
|
REPO="$(cd "$RECIPE/.." && pwd)"
|
||
|
|
OUTPUT_ROOT="${OUTPUT_ROOT:-$REPO/outputs}"
|
||
|
|
EVAL_ROOT="${EVAL_ROOT:-$OUTPUT_ROOT/eval}"
|
||
|
|
VANILLA_MODEL="${VANILLA_DIR:-$REPO}"
|
||
|
|
AHA_MODEL="${AHA_MODEL:-$OUTPUT_ROOT/aha/stage2/checkpoint-25}"
|
||
|
|
L2A_MODEL="${L2A_MODEL:-$OUTPUT_ROOT/l2a_style/stage2/checkpoint-25}"
|
||
|
|
GPU_LIST="${GPU_LIST:-0}"
|
||
|
|
THRESHOLDS=(0.45 0.50 0.525 0.55 0.575 0.60 0.625 0.65)
|
||
|
|
BENCHMARKS=(ruler_a ruler_b babilong helmet mrcr)
|
||
|
|
IFS=',' read -r -a GPUS <<< "$GPU_LIST"
|
||
|
|
declare -A SLOT_PIDS=()
|
||
|
|
mkdir -p "$EVAL_ROOT"
|
||
|
|
|
||
|
|
for model in "$VANILLA_MODEL" "$AHA_MODEL" "$L2A_MODEL"; do
|
||
|
|
[[ -f "$model/config.json" ]] || { echo "Missing model: $model" >&2; exit 2; }
|
||
|
|
done
|
||
|
|
|
||
|
|
jobs=()
|
||
|
|
for benchmark in "${BENCHMARKS[@]}"; do
|
||
|
|
jobs+=("vanilla|$VANILLA_MODEL|0.5|$benchmark")
|
||
|
|
done
|
||
|
|
for threshold in "${THRESHOLDS[@]}"; do
|
||
|
|
slug="${threshold/./}"
|
||
|
|
for benchmark in "${BENCHMARKS[@]}"; do
|
||
|
|
jobs+=("token_kv_head_t${slug}|$AHA_MODEL|$threshold|$benchmark")
|
||
|
|
jobs+=("token_t${slug}|$L2A_MODEL|$threshold|$benchmark")
|
||
|
|
done
|
||
|
|
done
|
||
|
|
|
||
|
|
launch() {
|
||
|
|
local gpu="$1" method="$2" model="$3" threshold="$4" benchmark="$5"
|
||
|
|
EVAL_ROOT="$EVAL_ROOT" RULER_LIMIT=20 BABILONG_LIMIT=50 \
|
||
|
|
AHA_GATE_HARD_THRESHOLD="$threshold" \
|
||
|
|
"$SCRIPT_DIR/eval_cell.sh" "$method" "$model" "$gpu" "$benchmark" &
|
||
|
|
LAST_PID="$!"
|
||
|
|
}
|
||
|
|
|
||
|
|
for i in "${!jobs[@]}"; do
|
||
|
|
slot=$((i % ${#GPUS[@]}))
|
||
|
|
if [[ -n "${SLOT_PIDS[$slot]:-}" ]]; then
|
||
|
|
wait "${SLOT_PIDS[$slot]}"
|
||
|
|
fi
|
||
|
|
IFS='|' read -r method model threshold benchmark <<< "${jobs[$i]}"
|
||
|
|
launch "${GPUS[$slot]}" "$method" "$model" "$threshold" "$benchmark"
|
||
|
|
SLOT_PIDS[$slot]="$LAST_PID"
|
||
|
|
done
|
||
|
|
|
||
|
|
status=0
|
||
|
|
for pid in "${SLOT_PIDS[@]}"; do
|
||
|
|
wait "$pid" || status=1
|
||
|
|
done
|
||
|
|
[[ "$status" == 0 ]] || exit "$status"
|
||
|
|
|
||
|
|
python "$SCRIPT_DIR/summarize_qwen1p7b_router_granularity_20260714.py" \
|
||
|
|
--repo "$REPO" \
|
||
|
|
--eval-root "$EVAL_ROOT" \
|
||
|
|
--baseline-root "$EVAL_ROOT/vanilla" \
|
||
|
|
--input-dir "$RECIPE/data/eval_inputs" \
|
||
|
|
--output-dir "$OUTPUT_ROOT/summary"
|