83 lines
2.9 KiB
Bash
83 lines
2.9 KiB
Bash
#!/usr/bin/env bash
|
|
set -euo pipefail
|
|
|
|
if [[ $# -ne 4 ]]; then
|
|
echo "Usage: $0 METHOD MODEL_PATH GPU BENCHMARK" >&2
|
|
exit 2
|
|
fi
|
|
|
|
METHOD="$1"
|
|
MODEL_PATH="$2"
|
|
GPU="$3"
|
|
BENCHMARK="$4"
|
|
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
|
RECIPE="$(cd "$SCRIPT_DIR/.." && pwd)"
|
|
REPO="$(cd "$RECIPE/.." && pwd)"
|
|
EVAL_ROOT="${EVAL_ROOT:-$REPO/outputs/eval}"
|
|
OUT="$EVAL_ROOT/$METHOD"
|
|
LOG="$EVAL_ROOT/logs/$METHOD"
|
|
TOKENIZER="${VANILLA_DIR:-$REPO}"
|
|
HELMET_ROWS="$RECIPE/data/eval_inputs/helmet_icl_8k_n50_per_config.jsonl"
|
|
MRCR_ROWS="$RECIPE/data/eval_inputs/mrcr_8k_2_4_8needle_n10_per_config.jsonl"
|
|
RULER_LIMIT="${RULER_LIMIT:-20}"
|
|
BABILONG_LIMIT="${BABILONG_LIMIT:-50}"
|
|
AHA_GATE_HARD_THRESHOLD="${AHA_GATE_HARD_THRESHOLD:-0.5}"
|
|
mkdir -p "$OUT" "$LOG"
|
|
|
|
if [[ -f "$OUT/${BENCHMARK}.DONE" ]]; then
|
|
exit 0
|
|
fi
|
|
|
|
export CUDA_VISIBLE_DEVICES="$GPU"
|
|
export TOKENIZERS_PARALLELISM=false
|
|
export PYTHONUNBUFFERED=1
|
|
export AHA_SPARSITY_STATS_PATH="$OUT/sparsity_${BENCHMARK}.json"
|
|
export AHA_GATE_HARD_THRESHOLD
|
|
export AHA_FORCE_FULL_DECODE=0
|
|
export AHA_FORCE_FULL_PREFILL_TAIL=0
|
|
export AHA_FORCE_FULL_HEADS=
|
|
export AHA_FORCE_LOW_ALPHA_FULL_HEADS=0
|
|
export AHA_DUO_PREFILL_FULL=0
|
|
|
|
model_args="pretrained=${MODEL_PATH},trust_remote_code=True,dtype=bfloat16,max_length=16384,attn_implementation=sdpa"
|
|
split_a="niah_single_1,niah_single_3,niah_multikey_2,niah_multiquery,ruler_vt,ruler_fwe,ruler_qa_hotpot"
|
|
split_b="niah_single_2,niah_multikey_1,niah_multikey_3,niah_multivalue,ruler_cwe,ruler_qa_squad"
|
|
|
|
case "$BENCHMARK" in
|
|
ruler_a|ruler_b)
|
|
tasks="$split_a"
|
|
[[ "$BENCHMARK" == ruler_b ]] && tasks="$split_b"
|
|
python -m lm_eval \
|
|
--model hf --model_args "$model_args" \
|
|
--tasks "$tasks" --metadata '{"max_seq_lengths":[8192]}' \
|
|
--batch_size 1 --limit "$RULER_LIMIT" --log_samples \
|
|
--output_path "$OUT/ruler8k_splits/${BENCHMARK#ruler_}/lm_eval" \
|
|
>"$LOG/${BENCHMARK}.log" 2>&1
|
|
;;
|
|
babilong)
|
|
python -m lm_eval \
|
|
--model hf --model_args "$model_args" \
|
|
--tasks babilong_longctx --metadata '{"max_seq_lengths":"8k"}' \
|
|
--num_fewshot 2 --batch_size 1 --limit "$BABILONG_LIMIT" --log_samples \
|
|
--output_path "$OUT/babilong8k_qa1_qa5_n50" \
|
|
>"$LOG/babilong.log" 2>&1
|
|
;;
|
|
helmet)
|
|
python "$SCRIPT_DIR/eval_qwen_external_longctx_pilot.py" \
|
|
--benchmark helmet_icl --model-path "$MODEL_PATH" --tokenizer-path "$TOKENIZER" \
|
|
--attn-implementation sdpa --rows "$HELMET_ROWS" \
|
|
--method "$METHOD" --output "$OUT/helmet_icl8k_n50.jsonl" \
|
|
>"$LOG/helmet.log" 2>&1
|
|
;;
|
|
mrcr)
|
|
python "$SCRIPT_DIR/eval_qwen_external_longctx_pilot.py" \
|
|
--benchmark mrcr --model-path "$MODEL_PATH" --tokenizer-path "$TOKENIZER" \
|
|
--attn-implementation sdpa --rows "$MRCR_ROWS" \
|
|
--method "$METHOD" --output "$OUT/mrcr_8k_2_4_8needle_n10.jsonl" \
|
|
>"$LOG/mrcr.log" 2>&1
|
|
;;
|
|
*) echo "Unknown benchmark: $BENCHMARK" >&2; exit 2 ;;
|
|
esac
|
|
|
|
touch "$OUT/${BENCHMARK}.DONE"
|