293 lines
12 KiB
Bash
293 lines
12 KiB
Bash
#!/usr/bin/env bash
|
|
#SBATCH --job-name=model-eval
|
|
#SBATCH --partition=mllm_safety
|
|
#SBATCH --quotatype=spot
|
|
#SBATCH --nodes=1
|
|
#SBATCH --ntasks-per-node=8
|
|
#SBATCH --gres=gpu:8
|
|
#SBATCH --cpus-per-task=8
|
|
#SBATCH --time=20:00:00
|
|
#SBATCH --output=logs/%x-%j.out
|
|
#SBATCH --error=logs/%x-%j.err
|
|
#SBATCH --requeue
|
|
|
|
# ============================================================
|
|
# Unified Evaluation Configuration + Execution Script
|
|
# ============================================================
|
|
|
|
|
|
# ------------------------------------------------------------
|
|
# Declare associative arrays
|
|
# ------------------------------------------------------------
|
|
declare -A eval_llada_base_configs
|
|
declare -A eval_llada_instruct_configs
|
|
|
|
declare -A eval_dream_base_configs
|
|
declare -A eval_dream_instruct_configs
|
|
|
|
declare -A eval_bert_configs
|
|
|
|
|
|
# ============================================================
|
|
# ==================== LLaDA CONFIGS ========================
|
|
# ============================================================
|
|
# eval_llada_configs["<dataset>"]="num_fewshot|max_new_tokens|steps|block_length|seed|mc_num|cfg"
|
|
# ============================================================
|
|
|
|
# ---------- Base Generation ----------
|
|
eval_llada_base_configs["gsm8k"]="8|1024|1024|32|1234|128|0.0"
|
|
eval_llada_base_configs["bbh"]="3|1024|1024|32|1234|128|0.0"
|
|
eval_llada_base_configs["minerva_math"]="4|1024|1024|32|1234|128|0.0"
|
|
eval_llada_base_configs["humaneval"]="0|1024|1024|32|1234|128|0.0"
|
|
eval_llada_base_configs["mbpp"]="3|1024|1024|32|1234|128|0.0"
|
|
|
|
# ---------- Base Likelihood ----------
|
|
eval_llada_base_configs["gpqa_main_n_shot"]="5|1024|1024|1024|1234|128|0.5"
|
|
eval_llada_base_configs["truthfulqa_mc2"]="0|1024|1024|1024|1234|128|2.0"
|
|
eval_llada_base_configs["arc_challenge"]="0|1024|1024|1024|1234|128|0.5"
|
|
eval_llada_base_configs["hellaswag"]="0|1024|1024|1024|1234|128|0.5"
|
|
eval_llada_base_configs["winogrande"]="5|1024|1024|1024|1234|128|0.0"
|
|
eval_llada_base_configs["piqa"]="0|1024|1024|1024|1234|128|0.5"
|
|
eval_llada_base_configs["mmlu"]="5|1024|1024|1024|1234|1|0.0"
|
|
eval_llada_base_configs["cmmlu"]="5|1024|1024|1024|1234|1|0.0"
|
|
eval_llada_base_configs["ceval-valid"]="5|1024|1024|1024|1234|1|0.0"
|
|
|
|
# ---------- Instruct Generation ----------
|
|
eval_llada_instruct_configs["gsm8k_cot"]="8|1024|1024|32|1234|1|0.0"
|
|
eval_llada_instruct_configs["bbh"]="3|1024|1024|32|1234|1|0.0"
|
|
eval_llada_instruct_configs["minerva_math"]="4|1024|1024|32|1234|1|0.0"
|
|
eval_llada_instruct_configs["humaneval_instruct"]="0|1024|1024|32|1234|1|0.0"
|
|
eval_llada_instruct_configs["mbpp_llada_instruct"]="3|1024|1024|32|1234|1|0.0"
|
|
|
|
eval_llada_instruct_configs["mmlu_generative"]="0|3|3|3|1234|1|0.0"
|
|
eval_llada_instruct_configs["mmlu_pro"]="0|256|256|256|1234|1|0.0"
|
|
eval_llada_instruct_configs["hellaswag_gen"]="0|3|3|3|1234|1|0.0"
|
|
eval_llada_instruct_configs["arc_challarc_challenge_chatenge"]="0|5|5|5|1234|1|0.0"
|
|
eval_llada_instruct_configs["gpqa_n_shot_gen"]="5|32|32|32|1234|1|0.0"
|
|
|
|
# ============================================================
|
|
# ==================== DREAM CONFIGS ========================
|
|
# ============================================================
|
|
# eval_dream_configs["<dataset>"]="num_fewshot|max_new_tokens|steps|temperature|top_p|seed|mc_num"
|
|
# ============================================================
|
|
|
|
# ---------- Base Generation ----------
|
|
eval_dream_base_configs["humaneval_dream"]="0|512|512|0.2|0.95|1234|1"
|
|
eval_dream_base_configs["gsm8k_cot"]="8|256|256|0.0|0.95|1234|1"
|
|
eval_dream_base_configs["mbpp"]="3|512|512|0.2|0.95|1234|1"
|
|
eval_dream_base_configs["minerva_math"]="4|512|512|0.0|0.95|1234|1"
|
|
eval_dream_base_configs["bbh"]="3|512|512|0.0|0.95|1234|1"
|
|
|
|
# ---------- Base Likelihood ----------
|
|
eval_dream_base_configs["mmlu"]="5|512|512|0.0|0.95|1234|128"
|
|
eval_dream_base_configs["arc_easy"]="0|512|512|0.0|0.95|1234|128"
|
|
eval_dream_base_configs["arc_challenge"]="0|512|512|0.0|0.95|1234|128"
|
|
eval_dream_base_configs["hellaswag"]="0|512|512|0.0|0.95|1234|128"
|
|
eval_dream_base_configs["piqa"]="0|512|512|0.0|0.95|1234|128"
|
|
eval_dream_base_configs["gpqa_main_n_shot"]="5|512|512|0.0|0.95|1234|128"
|
|
eval_dream_base_configs["winogrande"]="5|512|512|0.0|0.95|1234|128"
|
|
eval_dream_base_configs["race"]="0|512|512|0.0|0.95|1234|128"
|
|
|
|
# ---------- Instruct Generation ----------
|
|
eval_dream_instruct_configs["mmlu_generative"]="4|128|128|0.1|0.9|1234|1"
|
|
eval_dream_instruct_configs["mmlu_generative_dream"]="4|128|128|0.1|0.9|1234|1"
|
|
eval_dream_instruct_configs["mmlu_pro"]="4|128|128|0.1|0.9|1234|1"
|
|
eval_dream_instruct_configs["gsm8k_cot"]="0|256|256|0.1|0.9|1234|1"
|
|
eval_dream_instruct_configs["minerva_math"]="0|512|512|0.1|0.9|1234|1"
|
|
eval_dream_instruct_configs["gpqa_main_n_shot"]="5|128|128|0.0|1.0|1234|1"
|
|
eval_dream_instruct_configs["humaneval_instruct"]="0|768|768|0.1|0.9|1234|1"
|
|
eval_dream_instruct_configs["mbpp_instruct"]="0|1024|1024|0.1|0.9|1234|1"
|
|
eval_dream_instruct_configs["mbpp_instruct_dream"]="0|1024|1024|0.1|0.9|1234|1"
|
|
eval_dream_instruct_configs["ifeval"]="0|1280|1280|0.1|0.9|1234|1"
|
|
|
|
# ============================================================
|
|
# ==================== BERT CONFIGS =========================
|
|
# ============================================================
|
|
# eval_bert_configs["<dataset>"]="num_fewshot|max_new_tokens|steps|block_length|seed|mc_num"
|
|
# ============================================================
|
|
|
|
eval_bert_configs["mmlu"]="5|512|512|32|1234|128"
|
|
eval_bert_configs["ceval-valid"]="5|1024|1024|32|1234|128"
|
|
eval_bert_configs["cmmlu"]="5|1024|1024|32|1234|128"
|
|
eval_bert_configs["hellaswag"]="0|1024|1024|1024|1234|128"
|
|
eval_bert_configs["winogrande"]="0|128|128|128|1234|128"
|
|
|
|
eval_bert_configs["gsm8k_bert"]="8|256|256|32|1234|128"
|
|
eval_bert_configs["minerva_math"]="4|256|256|32|1234|128"
|
|
eval_bert_configs["humaneval"]="0|256|256|32|1234|128"
|
|
eval_bert_configs["bbh"]="3|256|256|32|1234|128"
|
|
|
|
|
|
eval_bert_configs["hellaswag_gen"]="0|128|128|128|1234|1"
|
|
eval_bert_configs["mmlu_generative"]="0|128|128|128|1234|1"
|
|
eval_bert_configs["mmlu_pro"]="0|256|256|256|1234|1"
|
|
eval_bert_configs["arc_challenge_chat"]="0|128|128|128|1234|1"
|
|
|
|
# ============================================================
|
|
# ====================== END CONFIGS ========================
|
|
# ============================================================
|
|
|
|
|
|
# ============================================================
|
|
# ====================== END CONFIGS ========================
|
|
# ============================================================
|
|
|
|
|
|
# ===== Derived variables =====
|
|
NUM_NODES=${SLURM_NNODES}
|
|
GPUS_PER_NODE=$(echo "${SLURM_JOB_GPUS}" | tr ',' '\n' | wc -l)
|
|
WORLD_SIZE=$((NUM_NODES * GPUS_PER_NODE))
|
|
MASTER_PORT=$((20000 + SLURM_JOB_ID % 10000))
|
|
NODELIST=($(scontrol show hostnames "${SLURM_JOB_NODELIST}"))
|
|
MASTER_ADDR=${NODELIST[0]}
|
|
TRAIN_NODES=("${NODELIST[@]}")
|
|
|
|
echo "============================"
|
|
echo "JOB NAME: ${SLURM_JOB_NAME}"
|
|
echo "JOB ID: ${SLURM_JOB_ID}"
|
|
echo "NUM_NODES: ${NUM_NODES}"
|
|
echo "WORLD_SIZE: ${WORLD_SIZE}"
|
|
echo "MASTER: ${MASTER_ADDR}:${MASTER_PORT}"
|
|
echo "============================"
|
|
|
|
# ===== Environment =====
|
|
export PYTHONBREAKPOINT=0
|
|
export NCCL_ASYNC_ERROR_HANDLING=1
|
|
export NCCL_DEBUG=warn
|
|
export TORCH_DISTRIBUTED_DEBUG=DETAIL
|
|
export PYTHONPATH=.:$PYTHONPATH
|
|
export HF_ALLOW_CODE_EVAL=1
|
|
export HF_DATASETS_TRUST_REMOTE_CODE=1 # For cmmlu dataset
|
|
export MASTER_ADDR MASTER_PORT WORLD_SIZE
|
|
|
|
|
|
MODEL_CLASS=${1,,} # "llada" or "dream"
|
|
TASK=${2:-"gsm8k"} # dataset name
|
|
MODEL_NAME=${3} # model path or name (required)
|
|
INSTRUCT=${4:-"False"} # whether to evaluate instruct model
|
|
BATCH_SIZE=${5:-"1"} # control batchsize
|
|
USE_LOG=${6:-"False"} # optional: enable logging
|
|
LIMIT=${7:-"None"} # optional: limit number of test samples (default None)
|
|
|
|
|
|
if [[ -z "${MODEL_NAME}" ]]; then
|
|
echo "❌ Missing model name/path argument!"
|
|
echo "Usage: sbatch eval_model.sh <model_class> <task> <model_name_or_path> [instruct] [batch_size]"
|
|
exit 1
|
|
fi
|
|
|
|
if [[ "${MODEL_NAME}" == /* ]]; then
|
|
MODEL_PATH="${MODEL_NAME}"
|
|
else
|
|
MODEL_PATH="${BASE_MODELS_DIR}/${MODEL_NAME}"
|
|
fi
|
|
|
|
case "${MODEL_CLASS}" in
|
|
llada)
|
|
if [[ "${INSTRUCT,,}" == "true" ]]; then
|
|
CONFIG="${eval_llada_instruct_configs[$TASK]}"
|
|
CONFIG_SET="instruct"
|
|
else
|
|
CONFIG="${eval_llada_base_configs[$TASK]}"
|
|
CONFIG_SET="base"
|
|
fi
|
|
|
|
if [[ -z "${CONFIG}" ]]; then
|
|
echo "❌ Unknown task '${TASK}' for LLaDA (${CONFIG_SET} mode)."
|
|
echo "Available tasks (base): ${!eval_llada_base_configs[@]}"
|
|
echo "Available tasks (instruct): ${!eval_llada_instruct_configs[@]}"
|
|
exit 1
|
|
fi
|
|
|
|
IFS="|" read -r NUM_FEWSHOT MAX_NEW_TOKENS STEPS BLOCK_LENGTH SEED MC_NUM CFG <<< "${CONFIG}"
|
|
|
|
MODEL_TYPE="llada"
|
|
SCRIPT_PATH="dllm/pipelines/llada/eval.py"
|
|
MODEL_ARGS="pretrained=${MODEL_PATH},is_check_greedy=False,mc_num=${MC_NUM},max_new_tokens=${MAX_NEW_TOKENS},steps=${STEPS},block_length=${BLOCK_LENGTH},cfg=${CFG}"
|
|
;;
|
|
|
|
dream)
|
|
if [[ "${INSTRUCT,,}" == "true" ]]; then
|
|
CONFIG="${eval_dream_instruct_configs[$TASK]}"
|
|
CONFIG_SET="instruct"
|
|
else
|
|
CONFIG="${eval_dream_base_configs[$TASK]}"
|
|
CONFIG_SET="base"
|
|
fi
|
|
|
|
if [[ -z "${CONFIG}" ]]; then
|
|
echo "❌ Unknown task '${TASK}' for Dream (${CONFIG_SET} mode)."
|
|
echo "Available tasks (base): ${!eval_dream_base_configs[@]}"
|
|
echo "Available tasks (instruct): ${!eval_dream_instruct_configs[@]}"
|
|
exit 1
|
|
fi
|
|
|
|
IFS="|" read -r NUM_FEWSHOT MAX_NEW_TOKENS STEPS TEMPERATURE TOP_P SEED MC_NUM <<< "${CONFIG}"
|
|
|
|
MODEL_TYPE="dream"
|
|
SCRIPT_PATH="dllm/pipelines/dream/eval.py"
|
|
MODEL_ARGS="pretrained=${MODEL_PATH},mc_num=${MC_NUM},max_new_tokens=${MAX_NEW_TOKENS},steps=${STEPS},temperature=${TEMPERATURE},top_p=${TOP_P},add_bos_token=true,escape_until=true"
|
|
;;
|
|
|
|
bert)
|
|
CONFIG="${eval_bert_configs[$TASK]}"
|
|
if [[ -z "${CONFIG}" ]]; then
|
|
echo "❌ Unknown task '${TASK}' for BERT."
|
|
echo "Available tasks: ${!eval_bert_configs[@]}"
|
|
exit 1
|
|
fi
|
|
|
|
IFS="|" read -r NUM_FEWSHOT MAX_NEW_TOKENS STEPS BLOCK_LENGTH SEED MC_NUM <<< "${CONFIG}"
|
|
|
|
MODEL_TYPE="bert"
|
|
SCRIPT_PATH="dllm/pipelines/bert/eval.py"
|
|
MODEL_ARGS="pretrained=${MODEL_PATH},is_check_greedy=False,mc_num=${MC_NUM},max_new_tokens=${MAX_NEW_TOKENS},steps=${STEPS},block_length=${BLOCK_LENGTH}"
|
|
;;
|
|
|
|
*)
|
|
echo "❌ Invalid model_class '${MODEL_CLASS}'. Must be 'llada' or 'dream' or 'bert'."
|
|
exit 1
|
|
;;
|
|
esac
|
|
|
|
|
|
[[ "${INSTRUCT}" == "True" ]] && APPLY_CHAT_TEMPLATE_ARG="--apply_chat_template True" || APPLY_CHAT_TEMPLATE_ARG=""
|
|
[[ "${LIMIT}" == "None" ]] && LIMIT_ARG="" || LIMIT_ARG="--limit ${LIMIT}"
|
|
[[ "${USE_LOG}" == "True" ]] && \
|
|
LOG_ARG="--log_samples --output_path ./logs/${MODEL_CLASS}_${TASK}_${SLURM_JOB_ID}_samples.json" \
|
|
|| LOG_ARG="--output_path ./logs/${MODEL_CLASS}_${TASK}_${SLURM_JOB_ID}_samples.json"
|
|
|
|
echo -e "\nLaunching ${MODEL_CLASS} on ${TASK} using ${MODEL_PATH}"
|
|
echo "============================"
|
|
echo "Few-shot: ${NUM_FEWSHOT}"
|
|
echo "Seed: ${SEED}"
|
|
echo "Batch size: ${BATCH_SIZE}"
|
|
echo "Use chat template: ${USE_CHAT_TEMPLATE}"
|
|
echo "============================"
|
|
|
|
RUN_CMD="accelerate launch \
|
|
--num_processes ${WORLD_SIZE} \
|
|
--num_machines ${NUM_NODES} \
|
|
--main_process_ip ${MASTER_ADDR} \
|
|
--main_process_port ${MASTER_PORT} \
|
|
--machine_rank ${SLURM_PROCID} \
|
|
${SCRIPT_PATH} \
|
|
--num_fewshot ${NUM_FEWSHOT} \
|
|
--batch_size ${BATCH_SIZE} \
|
|
--model ${MODEL_TYPE} \
|
|
--model_args \"${MODEL_ARGS}\" \
|
|
--tasks ${TASK} \
|
|
--seed ${SEED} \
|
|
${LOG_ARG} \
|
|
--confirm_run_unsafe_code \
|
|
${LIMIT_ARG} \
|
|
${APPLY_CHAT_TEMPLATE_ARG}"
|
|
|
|
if [[ "${NUM_NODES}" -eq 1 ]]; then
|
|
echo "Single-node execution"
|
|
eval ${RUN_CMD}
|
|
else
|
|
echo "Multi-node execution"
|
|
srun --nodes="${NUM_NODES}" --ntasks="${NUM_NODES}" --nodelist="${SLURM_JOB_NODELIST}" ${RUN_CMD}
|
|
fi
|