#!/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
'