File size: 2,876 Bytes
bca45b0 | 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
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
UMM_ROOT="${UMM_ROOT:-$(cd "${SCRIPT_DIR}/.." && pwd)}"
cd "${UMM_ROOT}"
unset http_proxy https_proxy HTTP_PROXY HTTPS_PROXY all_proxy ALL_PROXY
export PYTHONPATH="${UMM_ROOT}/code/umm/runtime:${UMM_ROOT}/code/umm:${UMM_ROOT}/code/umm/stage2_covt:${UMM_ROOT}/third_party/transformers/src:${UMM_ROOT}/third_party/diffusers/src:${PYTHONPATH:-}"
export HF_HOME="${HF_HOME:-${UMM_ROOT}/cache/huggingface}"
export TORCH_HOME="${TORCH_HOME:-${UMM_ROOT}/cache/torch}"
export TMPDIR="${TMPDIR:-${UMM_ROOT}/tmp}"
export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0,1,2,3,4,5,6,7}"
export TOKENIZERS_PARALLELISM=false
export OMP_NUM_THREADS="${OMP_NUM_THREADS:-4}"
PYTHON_BIN="${PYTHON_BIN:-${UMM_ROOT}/.env/unvideo/bin/python}"
CHECKPOINT_STEP="${CHECKPOINT_STEP:-2000}"
EXPERIMENT_NAME="${EXPERIMENT_NAME:-stage2_covt_sam_vggt_covt4_v2}"
MODEL_PATH="${MODEL_PATH:-${UMM_ROOT}/models/ckpts/${EXPERIMENT_NAME}/checkpoint-${CHECKPOINT_STEP}-merged}"
MANIFEST="${MANIFEST:-${UMM_ROOT}/data/cvbench_500/validation.jsonl}"
OUTPUT_DIR="${OUTPUT_DIR:-${UMM_ROOT}/outputs/evals/cvbench/${EXPERIMENT_NAME}-checkpoint-${CHECKPOINT_STEP}}"
MASTER_PORT="${MASTER_PORT:-29620}"
SWANLAB_PROJECT="${SWANLAB_PROJECT:-unified-model}"
SWANLAB_RUN_NAME="${SWANLAB_RUN_NAME:-H3_LoRA-ckpt${CHECKPOINT_STEP}-CVBench500}"
PROCESS_NAME="${PROCESS_NAME:-H3_EVAL_${CHECKPOINT_STEP}}"
SWANLAB_KEY_FILE="${SWANLAB_KEY_FILE:-${UMM_ROOT}/secrets/swanlab_key.txt}"
if [[ ! -x "${PYTHON_BIN}" ]]; then
echo "Missing Python environment: ${PYTHON_BIN}" >&2
exit 1
fi
if [[ ! -f "${MODEL_PATH}/config.json" ]]; then
echo "Merged model is not ready: ${MODEL_PATH}" >&2
exit 1
fi
if [[ ! -s "${MANIFEST}" ]]; then
echo "CV-Bench manifest is missing: ${MANIFEST}" >&2
exit 1
fi
if [[ -z "${SWANLAB_API_KEY:-}" ]]; then
if [[ ! -s "${SWANLAB_KEY_FILE}" ]]; then
echo "SwanLab key file is missing: ${SWANLAB_KEY_FILE}" >&2
exit 1
fi
IFS= read -r SWANLAB_API_KEY < "${SWANLAB_KEY_FILE}" || true
if [[ -z "${SWANLAB_API_KEY}" ]]; then
echo "SwanLab key file is empty: ${SWANLAB_KEY_FILE}" >&2
exit 1
fi
export SWANLAB_API_KEY
fi
CMD=(
"${PYTHON_BIN}" -m torch.distributed.run
--standalone
--nnodes=1
--nproc-per-node=8
"--master-port=${MASTER_PORT}"
"${UMM_ROOT}/code/umm/stage2_covt/eval_cvbench.py"
--model-path "${MODEL_PATH}"
--manifest "${MANIFEST}"
--output-dir "${OUTPUT_DIR}"
--dtype bfloat16
--attn-implementation flash_attention_2
--max-new-tokens 64
--checkpoint-step "${CHECKPOINT_STEP}"
--process-name "${PROCESS_NAME}"
--swanlab
--swanlab-project "${SWANLAB_PROJECT}"
--swanlab-run-name "${SWANLAB_RUN_NAME}"
)
CMD+=("$@")
if [[ "${DRY_RUN:-0}" == "1" ]]; then
printf '%q ' "${CMD[@]}"
printf '\n'
exit 0
fi
exec "${CMD[@]}"
|