File size: 1,114 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 | #!/bin/bash
#SBATCH --job-name=pretrain
#SBATCH --output=your_output_directory/output_%j.out
#SBATCH --error=your_error_directory/error_%j.err
#SBATCH --nodes=4
#SBATCH --gpus-per-node=4
#SBATCH --ntasks-per-node=4
#SBATCH --time=23:59:00
#SBATCH -p compute_full_node
# Load required modules
module --ignore_cache load cuda/11.4.4
module --ignore_cache load anaconda3
source activate your_environment
# Change to the working directory
cd your_working_directory
# Set environment variables
export CUBLAS_WORKSPACE_CONFIG=:4096:2
export PYTHONPATH="${PWD}/model:${PYTHONPATH:-}"
export NCCL_DEBUG=INFO
export PYTHONFAULTHANDLER=1
# Run the Python script with arguments
stdbuf -oL -eL srun python scripts/pretrain.py \
--tokenizer_path model/src/CodonTransformerTokenizer.json \
--train_data_path your_data_directory/pretrain_dataset.json \
--checkpoint_dir your_checkpoint_directory \
--batch_size 6 \
--max_epochs 5 \
--num_workers 5 \
--accumulate_grad_batches 1 \
--num_gpus 16 \
--learning_rate 0.00005 \
--warmup_fraction 0.1 \
--save_interval 5 \
--seed 123
|