| #!/bin/bash | |
| #SBATCH --job-name=ar-ckpt-test | |
| #SBATCH --partition=dev | |
| #SBATCH --nodes=1 | |
| #SBATCH --ntasks-per-node=1 | |
| #SBATCH --gres=gpu:8 | |
| #SBATCH --cpus-per-task=96 | |
| #SBATCH --exclusive | |
| #SBATCH --exclude=ip-10-0-136-14,ip-10-0-136-94,ip-10-0-136-253,ip-10-0-138-178,ip-10-0-139-113,ip-10-0-140-247 | |
| #SBATCH --time=00:20:00 | |
| #SBATCH --output=logs/ckpt_test_%j.log | |
| #SBATCH --error=logs/ckpt_test_%j.err | |
| set -euo pipefail | |
| cd /fsx/ajay-stu/AR | |
| source .venv/bin/activate | |
| export PYTHONPATH=$PWD:${PYTHONPATH:-} | |
| export MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n1) | |
| export MASTER_PORT=$((29500 + SLURM_JOB_ID % 1000)) | |
| export FI_PROVIDER=efa | |
| export NCCL_SOCKET_IFNAME=enp73s0 | |
| export LD_LIBRARY_PATH=/fsx/ajay-stu/AR/.venv/lib/python3.10/site-packages/nvidia/cudnn/lib:$LD_LIBRARY_PATH | |
| export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True | |
| export TOKENIZERS_PARALLELISM=false | |
| : "${MAX_STEPS:=60}" | |
| echo "MASTER_ADDR=$MASTER_ADDR MAX_STEPS=$MAX_STEPS port=$MASTER_PORT" | |
| srun bash -c ' | |
| torchrun \ | |
| --nnodes=1 \ | |
| --nproc_per_node=8 \ | |
| --rdzv_backend=c10d \ | |
| --rdzv_endpoint=$MASTER_ADDR:$MASTER_PORT \ | |
| -m ar.train.train_fsdp \ | |
| --bin /fsx/ajay-stu/ar-data/bin-val/fineweb.bin \ | |
| --model_size small \ | |
| --batch_size 2 \ | |
| --seq_len 1024 \ | |
| --max_steps '"$MAX_STEPS"' \ | |
| --warmup_steps 5 \ | |
| --ckpt_every 20 \ | |
| --ckpt_dir /fsx/ajay-stu/ar-data/checkpoints-test | |
| ' |