#!/bin/bash
# Grid 8 - does the parallelism curriculum hold ACROSS difficulties?
#
# The previous method was excellent at 8 digits (3/3 seeds ~100%) but unstable
# at 10-12 (1/3 seeds) because it demanded a full carry chain in one pass from
# random initialisation. `prog` anneals the number of parallel commits 16->1 and
# ramps the one-pass term in, so each stage is a small increment. Capacity is
# also raised, since 14-16 digits previously failed even at sequential decoding
# and that is a learnability failure, not a parallelism one.
#SBATCH -J g8
#SBATCH -p gpu-h100
#SBATCH --gpus=h100:1
#SBATCH -c 8
#SBATCH --mem=64g
#SBATCH -t 10:00:00
#SBATCH -o logs/g8-%A_%a.out
#SBATCH --array=0-29%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 10 12 14 16); MODES=(mdlm prog); SEEDS=(0 1 2)
D=${DIG[$((I / 6))]}
M=${MODES[$(((I % 6) / 3))]}
S=${SEEDS[$((I % 3))]}
SL=$(( 3 * D + 4 ))

NAME="g8-d$D-$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 512 --layers 10 --heads 8 --bs 128 --steps 60000 --lr 3e-4 \
  --K_max 16 --lam 1.0 --out $SCR/pcmd/runs/$NAME