#!/bin/bash
# Grid 6 - push past saturation. 4 and 6 digits are 100% for every method, so
# the effect only shows once the carry chain exceeds what one forward pass can
# compute. This traces the frontier and checks the gap widens with depth.
#SBATCH -J g6
#SBATCH -p gpu-h100
#SBATCH --gpus=h100:1
#SBATCH -c 8
#SBATCH --mem=64g
#SBATCH -t 08:00:00
#SBATCH -o logs/g6-%A_%a.out
#SBATCH --array=0-47%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
MODES=(mdlm matched distill full); DIG=(10 12 14 16); SEEDS=(0 1 2)
M=${MODES[$((I / 12))]}
D=${DIG[$(((I % 12) / 3))]}
S=${SEEDS[$((I % 3))]}
SL=$(( 3 * D + 4 ))

NAME="g6-d$D-$M-s$S"
echo "=== $NAME ==="
$PY train_add.py --mode $M --digits $D --seq_len $SL --seed $S \
  --n_train 300000 --n_eval 500 --eval_bs 100 \
  --d 384 --layers 6 --heads 6 --bs 128 --steps 30000 --lr 3e-4 \
  --K 4 --lam 1.0 --out $SCR/pcmd/runs/$NAME