#!/bin/bash
# Grid 10 - is depth failing for OPTIMISATION reasons? This must be settled
# before the depth-scaling test means anything.
#
# At 8 digits a 6-layer/384 model reached 99.9%, while 10-layer/512 reached
# 66.7% with more training. Deeper transformers usually need a lower learning
# rate and longer warmup; if that is the cause, the depth sweep at a single LR
# would fake a "depth does not help" result.
#SBATCH -J g10
#SBATCH -p gpu-h100
#SBATCH --gpus=h100:1
#SBATCH -c 8
#SBATCH --mem=64g
#SBATCH -t 10:00:00
#SBATCH -o logs/g10-%A_%a.out
#SBATCH --array=0-35%12
source /etc/profile
SCR=/scratch/zt1/project/msml612/user/$USER
module load python/3.10.10
export PYTHONPATH=
PY=$SCR/hf_env/bin/python
cd $SCR/pcmd/src

I=$SLURM_ARRAY_TASK_ID
LAYERS=(6 10); LR=(3e-4 1e-4 5e-5); WARM=(500 4000); SEEDS=(0 1 2)
L=${LAYERS[$((I / 18))]}
R=${LR[$(((I % 18) / 6))]}
W=${WARM[$(((I % 6) / 3))]}
S=${SEEDS[$((I % 3))]}

NAME="g10-L$L-lr$R-w$W-s$S"
echo "=== $NAME ==="
$PY train_add.py --mode full --digits 10 --seq_len 34 --seed $S \
  --n_train 500000 --n_eval 400 --eval_bs 100 \
  --d 384 --layers $L --heads 6 --bs 128 --steps 40000 --lr $R --warmup $W \
  --K 4 --lam 1.0 --out $SCR/pcmd/runs/$NAME