File size: 2,971 Bytes
46b9eea | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 | #!/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"
|