File size: 2,380 Bytes
53e66de | 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 | #!/usr/bin/env bash
set -euo pipefail
# Run this script inside an allocated/interactive GPU session.
# It does not request SLURM resources.
PROJECT_DIR="${PROJECT_DIR:-/public/home/scnb9biwet/jiangqq/CodonTransformer-main}"
HF_HOME="${HF_HOME:-/public/home/scnb9biwet/.cache/huggingface}"
CONDA_ENV="${CONDA_ENV:-struct-evo}"
TOKENIZER_PATH="${TOKENIZER_PATH:-${PROJECT_DIR}/model/src/CodonTransformerTokenizer.json}"
TRAIN_DATA_PATH="${TRAIN_DATA_PATH:-${PROJECT_DIR}/scripts/data/processed/pretrain_data.json}"
CHECKPOINT_DIR="${CHECKPOINT_DIR:-${PROJECT_DIR}/weight/checkpoints/pretrain}"
BATCH_SIZE="${BATCH_SIZE:-6}"
MAX_EPOCHS="${MAX_EPOCHS:-5}"
NUM_WORKERS="${NUM_WORKERS:-5}"
ACCUMULATE_GRAD_BATCHES="${ACCUMULATE_GRAD_BATCHES:-1}"
NUM_GPUS="${NUM_GPUS:-16}"
LEARNING_RATE="${LEARNING_RATE:-0.00005}"
WARMUP_FRACTION="${WARMUP_FRACTION:-0.1}"
SAVE_INTERVAL="${SAVE_INTERVAL:-5}"
SEED="${SEED:-123}"
DEBUG="${DEBUG:-0}"
OFFLINE="${OFFLINE:-1}"
cd "${PROJECT_DIR}"
mkdir -p "${CHECKPOINT_DIR}"
export HF_HOME
export PYTHONPATH="${PROJECT_DIR}/model:${PYTHONPATH:-}"
export CUBLAS_WORKSPACE_CONFIG=:4096:2
export NCCL_DEBUG="${NCCL_DEBUG:-INFO}"
export PYTHONFAULTHANDLER=1
if [[ "${OFFLINE}" == "1" ]]; then
export HF_HUB_OFFLINE=1
export TRANSFORMERS_OFFLINE=1
fi
if [[ -n "${CONDA_ENV}" ]] && command -v conda >/dev/null 2>&1; then
# shellcheck disable=SC1091
source "$(conda info --base)/etc/profile.d/conda.sh"
conda activate "${CONDA_ENV}"
fi
if [[ ! -f "${TOKENIZER_PATH}" ]]; then
echo "Missing TOKENIZER_PATH: ${TOKENIZER_PATH}" >&2
exit 1
fi
if [[ ! -f "${TRAIN_DATA_PATH}" ]]; then
echo "Missing TRAIN_DATA_PATH: ${TRAIN_DATA_PATH}" >&2
echo "Create it first with CodonTransformer.CodonData.prepare_training_data." >&2
exit 1
fi
args=(
"${PROJECT_DIR}/scripts/pretrain.py"
--tokenizer_path "${TOKENIZER_PATH}"
--train_data_path "${TRAIN_DATA_PATH}"
--checkpoint_dir "${CHECKPOINT_DIR}"
--batch_size "${BATCH_SIZE}"
--max_epochs "${MAX_EPOCHS}"
--num_workers "${NUM_WORKERS}"
--accumulate_grad_batches "${ACCUMULATE_GRAD_BATCHES}"
--num_gpus "${NUM_GPUS}"
--learning_rate "${LEARNING_RATE}"
--warmup_fraction "${WARMUP_FRACTION}"
--save_interval "${SAVE_INTERVAL}"
--seed "${SEED}"
)
if [[ "${DEBUG}" == "1" ]]; then
args+=(--debug)
fi
python "${args[@]}"
|