File size: 2,985 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 | #!/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}"
BENCHMARK="${BENCHMARK:?Set BENCHMARK to cvbench, blink, or vstar}"
CHECKPOINT_STEP="${CHECKPOINT_STEP:-7000}"
MODEL_PATH="${MODEL_PATH:-${UMM_ROOT}/models/ckpts/stage2_covt_from3k_bcd_2k/merged/checkpoint-${CHECKPOINT_STEP}-merged}"
MANIFEST="${MANIFEST:?Set MANIFEST to a prepared benchmark manifest}"
OUTPUT_DIR="${OUTPUT_DIR:?Set OUTPUT_DIR to a new result directory}"
MASTER_PORT="${MASTER_PORT:-29700}"
PROCESS_NAME="${PROCESS_NAME:-H3_${BENCHMARK}_7K}"
MAX_NEW_TOKENS="${MAX_NEW_TOKENS:-256}"
# SWANLAB_PROJECT is a reserved structured setting in SwanLab 0.9.x. Accept
# it for backwards compatibility, then remove it from the Python environment.
SWANLAB_PROJECT_NAME="${SWANLAB_PROJECT_NAME:-${SWANLAB_PROJECT:-unified-model}}"
unset SWANLAB_PROJECT
SWANLAB_RUN_NAME="${SWANLAB_RUN_NAME:-H3_LoRA-ckpt${CHECKPOINT_STEP}-${BENCHMARK}}"
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 "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
exec "${PYTHON_BIN}" -m torch.distributed.run \
--standalone \
--nnodes=1 \
--nproc-per-node=8 \
"--master-port=${MASTER_PORT}" \
"${UMM_ROOT}/code/umm/stage2_covt/eval_multibench.py" \
--model-path "${MODEL_PATH}" \
--manifest "${MANIFEST}" \
--output-dir "${OUTPUT_DIR}" \
--benchmark "${BENCHMARK}" \
--checkpoint-step "${CHECKPOINT_STEP}" \
--max-new-tokens "${MAX_NEW_TOKENS}" \
--process-name "${PROCESS_NAME}" \
--swanlab \
--swanlab-project "${SWANLAB_PROJECT_NAME}" \
--swanlab-run-name "${SWANLAB_RUN_NAME}" \
"$@"
|