#!/bin/bash
# Grid 11 - THE HEADLINE EXPERIMENT, with hyperparameters that are not broken.
#
# Grid 10 showed every "the method fails at 10+ digits" result was a learning
# rate artifact: at lr 3e-4 with 500 warmup a 10-layer model reached 47%, while
# at lr 1e-4 the same setting reached 100% on 3/3 seeds. Depth was never the
# problem. This reruns the full difficulty sweep at lr 1e-4 / warmup 4000, with
# depth scaled as log2(digits) per THEORY.md.
#SBATCH -J g11
#SBATCH -p gpu-h100
#SBATCH --gpus=h100:1
#SBATCH -c 8
#SBATCH --mem=64g
#SBATCH -t 12:00:00
#SBATCH -o logs/g11-%A_%a.out
#SBATCH --array=0-39%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
DIG=(8 12 16 24 32); MODES=(mdlm full); SEEDS=(0 1 2 3)
D=${DIG[$((I / 8))]}
M=${MODES[$(((I % 8) / 4))]}
S=${SEEDS[$((I % 4))]}
SL=$(( 3 * D + 4 ))
# depth ~ 2*log2(n), generous but still logarithmic
case $D in
  8)  L=8  ;;
  12) L=8  ;;
  16) L=10 ;;
  24) L=10 ;;
  32) L=12 ;;
esac

NAME="g11-d$D-$M-s$S"
echo "=== $NAME (layers=$L) ==="
$PY train_add.py --mode $M --digits $D --seq_len $SL --seed $S \
  --n_train 800000 --n_eval 400 --eval_bs 100 \
  --d 384 --layers $L --heads 6 --bs 128 --steps 60000 --lr 1e-4 --warmup 4000 \
  --K 4 --lam 1.0 --out $SCR/pcmd/runs/$NAME