#!/bin/bash
# Closes the one acknowledged gap in the ablation.
#
# The headline comparison - recurrent 99.70% against 32 distinct layers 94.50% -
# conflates TWO differences: weight sharing, and re-supplying the input at every
# application. --no_inject holds architecture, depth, parameter count and
# objective fixed and changes only the injection, which isolates them.
#
#   no-injection stays near 98%  -> weight sharing is the mechanism
#   no-injection falls toward 94% -> injection is, and our explanation is wrong
#
# Runs on a 5GB MIG slice: at 212k parameters and 81-token sequences the
# optimiser state is a few MB, and that partition bills at 9/min against the
# H100's 144 - 16x more compute from the same remaining allocation.
#SBATCH -J inj
#SBATCH -p gpu-a100_1g.5gb
#SBATCH --gres=gpu:a100_1g.5gb:1
#SBATCH -c 4
#SBATCH --mem=16g
#SBATCH -t 6:00:00
#SBATCH -o logs/inj-%A_%a.out
#SBATCH --array=0-4%5
source /etc/profile
SCR=/scratch/zt1/project/msml612/user/$USER
module load python/3.10.10
export PYTHONPATH=
cd $SCR/pcmd/src
I=$SLURM_ARRAY_TASK_ID
COMMON="--mode full --pe ape --hard $SCR/pcmd/data/sudoku_hard --augment \
  --heads 4 --bs 256 --steps 110000 --lr 3e-4 --warmup 3000 --K 8 --lam 1.0 \
  --eval_bs 200 --ckpt_every 8000"
case $I in
  0) NAME="inj-noinject-s0"; ARGS="--recurrent --R 32 --R_infer 64 --no_inject --d 128 --layers 1 --seed 0" ;;
  1) NAME="inj-noinject-s1"; ARGS="--recurrent --R 32 --R_infer 64 --no_inject --d 128 --layers 1 --seed 1" ;;
  2) NAME="inj-deep32L-s1";  ARGS="--d 128 --layers 32 --seed 1" ;;
  3) NAME="inj-rec-s2";      ARGS="--recurrent --R 32 --R_infer 64 --d 128 --layers 1 --seed 2" ;;
  4) NAME="inj-rec-s3";      ARGS="--recurrent --R 32 --R_infer 64 --d 128 --layers 1 --seed 3" ;;
esac
echo "=== $NAME ==="
$SCR/hf_env/bin/python -u train_sudoku.py $ARGS $COMMON --out $SCR/pcmd/runs/$NAME