File size: 1,458 Bytes
2220ce4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 | #!/bin/bash
# Grid 5 - Parallelism-Matched Diffusion on addition (unique answer, chained).
#
# mdlm standard uniform-t training (baseline)
# matched train only at the masking ratios a K-pass decode visits
# distill uniform training + one-pass supervision
# full matched schedule + one-pass supervision (the method)
#
# Prediction: the gap should be largest at FEW passes (high parallelism), where
# the baseline has to fit the whole carry chain into one forward pass having
# been trained mostly on nearly-complete states it never sees there. At many
# passes all methods should converge, since sequential decoding amortises the
# computation anyway.
#SBATCH -J g5
#SBATCH -p gpu-h100
#SBATCH --gpus=h100:1
#SBATCH -c 8
#SBATCH --mem=64g
#SBATCH -t 06:00:00
#SBATCH -o logs/g5-%A_%a.out
#SBATCH --array=0-35%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=(4 6 8); SEEDS=(0 1 2)
M=${MODES[$((I / 9))]}
D=${DIG[$(((I % 9) / 3))]}
S=${SEEDS[$((I % 3))]}
SL=$(( 3 * D + 4 ))
NAME="g5-d$D-$M-s$S"
echo "=== $NAME ==="
$PY train_add.py --mode $M --digits $D --seq_len $SL --seed $S \
--n_train 200000 --n_eval 500 --eval_bs 100 \
--d 384 --layers 6 --heads 6 --bs 128 --steps 25000 --lr 3e-4 \
--K 4 --lam 1.0 --out $SCR/pcmd/runs/$NAME
|