| #!/bin/bash | |
| #SBATCH --job-name=ar-lrsweep | |
| #SBATCH --partition=dev | |
| #SBATCH --nodes=2 | |
| #SBATCH --ntasks-per-node=1 | |
| #SBATCH --gpus-per-node=8 | |
| #SBATCH --cpus-per-task=96 | |
| #SBATCH --output=logs/lrsweep_%x_%j.out | |
| #SBATCH --error=logs/lrsweep_%x_%j.err | |
| # --nodelist is set from the command line via --export; see the launch commands below | |
| set -euo pipefail | |
| # --- env --- | |
| cd /fsx/ajay-stu/AR | |
| source .venv/bin/activate | |
| export PYTHONPATH=$PWD:${PYTHONPATH:-} | |
| export MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -1) | |
| export MASTER_PORT=$((29500 + SLURM_JOB_ID % 1000)) | |
| export NCCL_SOCKET_IFNAME=enp73s0 | |
| export FI_PROVIDER=efa | |
| export LD_LIBRARY_PATH=$PWD/.venv/lib/python*/site-packages/nvidia/cudnn/lib:$LD_LIBRARY_PATH | |
| export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True | |
| export TOKENIZERS_PARALLELISM=false | |
| # --- knobs the sweep varies --- | |
| # PEAK_LR must be passed in from the command line via --export | |
| : "${PEAK_LR:?PEAK_LR must be set, e.g. --export=ALL,PEAK_LR=5e-4}" | |
| # --- knobs held constant across the sweep --- | |
| BIN_DIR=/fsx/ajay-stu/ar-data/bin-val | |
| BATCH=4 | |
| SEQ=4096 | |
| STEPS=2000 | |
| WARMUP=200 | |
| echo "=== sweep run: PEAK_LR=$PEAK_LR batch=$BATCH seq=$SEQ steps=$STEPS ===" | |
| srun --cpu-bind=none bash -c ' | |
| export RANK=$SLURM_PROCID | |
| export LOCAL_RANK=$SLURM_LOCALID | |
| export WORLD_SIZE=$SLURM_NTASKS | |
| torchrun \ | |
| --nnodes=$SLURM_NNODES \ | |
| --nproc_per_node=8 \ | |
| --rdzv_backend=c10d \ | |
| --rdzv_endpoint=$MASTER_ADDR:$MASTER_PORT \ | |
| ar/train/train_fsdp.py \ | |
| --bin '"$BIN_DIR"'/fineweb.bin '"$BIN_DIR"'/stack.bin '"$BIN_DIR"'/finemath.bin \ | |
| --weights 0.65 0.20 0.15 \ | |
| --batch_size '"$BATCH"' \ | |
| --seq_len '"$SEQ"' \ | |
| --max_steps '"$STEPS"' \ | |
| --warmup_steps '"$WARMUP"' \ | |
| --peak_lr '"$PEAK_LR"' \ | |
| --model_size full | |
| ' |