| # 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 | |