#!/bin/bash
# Grid 7 - LEARNABILITY GATE for long addition, before any parallelism claim.
#
# At 14-16 digits every method scored 0 even at FULLY SEQUENTIAL decoding, so
# the models had not learned addition at all. Any statement about parallel
# decoding there would be vacuous. This sweeps capacity and training length to
# find where the task becomes learnable; only then does measuring parallelism
# mean anything.
#SBATCH -J g7
#SBATCH -p gpu-h100
#SBATCH --gpus=h100:1
#SBATCH -c 8
#SBATCH --mem=64g
#SBATCH -t 10:00:00
#SBATCH -o logs/g7-%A_%a.out
#SBATCH --array=0-23%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
# capacity x training length, on the two settings that failed outright
CFG=(
  "12 8 512 60000" "12 12 512 60000" "12 8 512 120000"
  "16 8 512 60000" "16 12 512 60000" "16 12 768 120000"
)
MODES=(mdlm full)
C=${CFG[$((I / 4))]}
M=${MODES[$(((I % 4) / 2))]}
S=$((I % 2))
read -r D LY DM ST <<< "$C"
SL=$(( 3 * D + 4 ))

NAME="g7-d$D-L$LY-dm$DM-st$ST-$M-s$S"
echo "=== $NAME ==="
$PY train_add.py --mode $M --digits $D --seq_len $SL --seed $S \
  --n_train 500000 --n_eval 400 --eval_bs 100 \
  --d $DM --layers $LY --heads 8 --bs 128 --steps $ST --lr 3e-4 \
  --K 4 --lam 1.0 --out $SCR/pcmd/runs/$NAME