mirror of
https://github.com/RYDE-WORK/lnp_ml.git
synced 2026-07-25 15:35:25 +08:00
110 lines
5.1 KiB
Bash
110 lines
5.1 KiB
Bash
#!/usr/bin/env bash
|
||
set -uo pipefail
|
||
cd "$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
|
||
|
||
# ===== 环境（可复现）=====
|
||
export TRANSFORMERS_OFFLINE=1 HF_HUB_OFFLINE=1 PYTHONUNBUFFERED=1
|
||
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True TOKENIZERS_PARALLELISM=false
|
||
export PYTHONHASHSEED=${PYTHONHASHSEED:-42}
|
||
mkdir -p logs models/abl_full reports
|
||
|
||
# ===== 参数 =====
|
||
SEED=${SEED:-42}
|
||
BATCH=${BATCH:-16}; QBATCH=${QBATCH:-8}
|
||
N_OUTER=${N_OUTER:-5}; N_INNER=${N_INNER:-3}
|
||
N_TRIALS=${N_TRIALS:-20}; EPOCHS=${EPOCHS:-20}; PATIENCE=${PATIENCE:-5}
|
||
INPUT=${INPUT:-data/interim/internal.csv}
|
||
CHEM_CACHE=data/processed/chemeleon_embeddings.npz
|
||
UNIMOL_CACHE=data/processed/unimol_embeddings.npz
|
||
MOLT5=models/molt5-base; BIOT5=models/biot5-plus-base; QWEN=models/qwen2.5-7b-instruct
|
||
LLM_COMMON="--use-llm --use-rag --rag-top-k 4 --use-soft-prompt --no-llm-freeze"
|
||
|
||
GPU_POOL=${GPU_POOL:-0,1} # 我们可用的卡
|
||
MAX_RETRY=${MAX_RETRY:-5}
|
||
POLL=${POLL:-60} # 轮询/重试间隔(秒)
|
||
IFS=',' read -r -a POOL <<< "$GPU_POOL"
|
||
|
||
ts(){ awk '{ print strftime("%F %T"), $0; fflush() }'; }
|
||
|
||
record_env(){ local dir=$1; shift
|
||
{ echo "date: $(date -Is)"
|
||
echo "git_commit: $(git rev-parse HEAD 2>/dev/null)"
|
||
echo "git_dirty_files: $(git status --porcelain 2>/dev/null | wc -l)"
|
||
echo "seed: $SEED pythonhashseed: $PYTHONHASHSEED"
|
||
echo "cmd: $*"; } > "${dir}/repro.txt"
|
||
git diff > "${dir}/local_uncommitted.patch" 2>/dev/null
|
||
pip freeze > "${dir}/pip_freeze.txt" 2>/dev/null
|
||
}
|
||
|
||
wait_free(){ local gpu=$1 need=$2 free
|
||
while :; do
|
||
free=$(nvidia-smi --query-gpu=memory.free --format=csv,noheader,nounits -i "$gpu" 2>/dev/null)
|
||
[ "${free:-0}" -ge "$need" ] && return 0
|
||
echo "[$(date '+%F %T')] GPU${gpu} 剩 ${free}MiB(<${need}) 等待 ${POLL}s..." >&2
|
||
sleep "$POLL"
|
||
done
|
||
}
|
||
|
||
run_one(){ local gpu=$1 name=$2 need=$3; shift 3
|
||
local rundir="models/abl_full/${name}/seed${SEED}"
|
||
if ls "${rundir}"/summary.json "${rundir}"/*/summary.json >/dev/null 2>&1; then
|
||
echo "[$(date '+%F %T')] SKIP ${name}（已完成）"; return 0; fi
|
||
mkdir -p "${rundir}"; local n=1
|
||
while :; do
|
||
wait_free "$gpu" "$need"
|
||
echo "[$(date '+%F %T')] >>> ${name} try $n on GPU${gpu}" | tee -a "${rundir}/run.log"
|
||
record_env "${rundir}" "GPU=$gpu $*"
|
||
if CUDA_VISIBLE_DEVICES="$gpu" python -m lnp_ml.modeling.nested_cv_optuna \
|
||
--input-path "${INPUT}" \
|
||
--output-dir "models/abl_full/${name}" --resume-dir "${rundir}" \
|
||
--seed ${SEED} --device cuda \
|
||
--n-outer-folds ${N_OUTER} --n-inner-folds ${N_INNER} \
|
||
--n-trials ${N_TRIALS} --epochs-per-trial ${EPOCHS} --inner-patience ${PATIENCE} \
|
||
"$@" 2>&1 | ts | tee -a "${rundir}/run.log"; then
|
||
echo "[$(date '+%F %T')] <<< ${name} DONE" | tee -a "${rundir}/run.log"; return 0; fi
|
||
(( n > MAX_RETRY )) && { echo "[$(date '+%F %T')] !!! ${name} FAILED x${MAX_RETRY}" | tee -a "${rundir}/run.log"; return 1; }
|
||
echo "[$(date '+%F %T')] ${name} 失败, ${POLL}s 后断点续跑重试..." | tee -a "${rundir}/run.log"; sleep "$POLL"; ((n++))
|
||
done
|
||
}
|
||
|
||
# ===== 变体队列：名称|需显存MiB|flags =====
|
||
QUEUE=$(mktemp); LOCK="${QUEUE}.lock"
|
||
add(){ echo "$*" >> "$QUEUE"; }
|
||
# Study 2：encoder 消融（MoE/LLM 全关）
|
||
add "s2_chemeleon|8000|--batch-size ${BATCH} --use-chemeleon --chemeleon-cache ${CHEM_CACHE}"
|
||
add "s2_unimol|8000|--batch-size ${BATCH} --use-unimol --unimol-cache ${UNIMOL_CACHE}"
|
||
# Study 3：LLM backbone（同集成模式，T5→LoRA）
|
||
add "s3_molt5|10000|--batch-size ${BATCH} --use-mpnn ${LLM_COMMON} --llm-use-lora --llm-model-path ${MOLT5}"
|
||
add "s3_biot5|10000|--batch-size ${BATCH} --use-mpnn ${LLM_COMMON} --llm-use-lora --llm-model-path ${BIOT5}"
|
||
# Study 1：核心模块（LLM=Qwen-RAG，Qwen→QLoRA）
|
||
add "s1_baseline|8000|--batch-size ${BATCH} --use-mpnn"
|
||
add "s1_moe|8000|--batch-size ${BATCH} --use-mpnn --use-moe"
|
||
add "s1_llm|14000|--batch-size ${QBATCH} --use-mpnn ${LLM_COMMON} --llm-use-qlora --llm-model-path ${QWEN}"
|
||
add "s1_both|14000|--batch-size ${QBATCH} --use-mpnn --use-moe ${LLM_COMMON} --llm-use-qlora --llm-model-path ${QWEN}"
|
||
|
||
pop_job(){ exec 9>"$LOCK"; flock 9
|
||
local line; line=$(head -n1 "$QUEUE" 2>/dev/null)
|
||
[ -n "$line" ] && sed -i '1d' "$QUEUE"
|
||
flock -u 9; echo "$line"; }
|
||
|
||
worker(){ local gpu=$1
|
||
while :; do
|
||
local job; job=$(pop_job); [ -z "$job" ] && break
|
||
local name=${job%%|*}; local r=${job#*|}; local need=${r%%|*}; local flags=${r#*|}
|
||
run_one "$gpu" "$name" "$need" $flags
|
||
done
|
||
}
|
||
|
||
echo "[$(date '+%F %T')] START seed=${SEED} pool=${GPU_POOL}"
|
||
for g in "${POOL[@]}"; do worker "$g" & done
|
||
wait
|
||
|
||
# ===== 汇总 =====
|
||
latest(){ ls -dt "models/abl_full/$1/seed${SEED}"/summary.json "models/abl_full/$1/seed${SEED}"/*/summary.json 2>/dev/null | head -1 | xargs -r dirname; }
|
||
sum_args=()
|
||
for name in s1_baseline s1_moe s1_llm s1_both s2_chemeleon s2_unimol s3_molt5 s3_biot5; do
|
||
d=$(latest "$name"); [ -n "$d" ] && sum_args+=(--run "${name}=$d")
|
||
done
|
||
d=$(latest s1_baseline); [ -n "$d" ] && sum_args+=(--run "s2_mpnn=$d")
|
||
python scripts/summarize_ablation.py "${sum_args[@]}" --out reports/ablation_full_summary.csv 2>&1 | ts | tee logs/summary_full.log
|
||
echo "[$(date '+%F %T')] ALL DONE" |