File size: 2,579 Bytes
17452c2 | 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 84 | : "${BASE_DIR:?Environment variable BASE_DIR not set}"
: "${DATA_DIR:?Environment variable DATA_DIR not set}"
# ============================================
# This script loads datasets from Hugging Face Hub (PUBLIC dataset)
# ============================================
# Dataset: zihaojing/MuMo-Pretraining (public)
#
# Note: --use_auth_token is NOT needed for public datasets
# Only add --use_auth_token if:
# - Your dataset is private
# - Your model is private
# - The dataset/model is gated and requires access approval
# ============================================
export PYTHONPATH=${BASE_DIR}
filename=$(basename "${BASH_SOURCE[0]}" .sh)
output_model=${DATA_DIR}/model/pretrain/${filename}
export WANDB_PROJECT="NeurIPS_Rebuttal"
export WANDB_DIR="${output_model}/wandb"
# Base config
DS_CONFIG=${BASE_DIR}/config/deepspeed/ds_config_zero2.json
MODEL_CONFIG=${BASE_DIR}/config/mumo/config_cls.json
# Keep this
SCRIPT_PATH="$(realpath "$0")"
if [ ! -d ${output_model} ];then
mkdir ${output_model}
fi
cp ${SCRIPT_PATH} ${output_model}
cp ${DS_CONFIG} ${output_model}
cp ${MODEL_CONFIG} ${output_model}/config.json
cp ${BASE_DIR}/train/pretrain.py ${output_model}
# export CUDA_HOME=/usr/local/cuda
# Deepspeed settings
MASTER_PORT=29500
GPUs=7
# Runner
deepspeed --master_port ${MASTER_PORT} --include localhost:${GPUs} ${BASE_DIR}/train/pretrain.py \
--run_name ${filename} \
--config_name ${MODEL_CONFIG} \
--tokenizer_name ${BASE_DIR}/smiles_tokenizer/mumo_tokenizer \
--use_fast_tokenizer false \
--output_dir ${output_model} \
--model_class MuMoPretrain \
--ddp_timeout 18000000 \
--dataset_name zihaojing/MuMo-Pretraining \
--preprocessing_num_workers 20 \
--seed 42 \
--ignore_data_skip true \
--per_device_train_batch_size 128 \
--per_device_eval_batch_size 128 \
--num_train_epochs 2 \
--learning_rate 1e-4 \
--lr_scheduler_type cosine \
--warmup_steps 2000 \
--gradient_accumulation_steps 1 \
--evaluation_strategy steps \
--eval_steps 1000 \
--max_eval_samples 5000 \
--save_strategy epoch \
--save_total_limit 2 \
--logging_dir ${output_model}/logs \
--logging_strategy steps \
--logging_steps 20 \
--disable_tqdm false \
--ddp_find_unused_parameters true \
--overwrite_output_dir \
--report_to wandb \
--do_train \
--do_eval \
--bf16 True \
--torch_dtype float32 \
| tee -a ${output_model}/train.log
# --resume_from_checkpoint ${output_model}/checkpoint-20400 \
# --save_steps 1243 \
|