Automated push from Zaratan

README.md

@@ -0,0 +1,18 @@
+---
+license: apache-2.0
+tags: [masked-diffusion, discrete-diffusion, length-generalization, positional-encoding, from-scratch]
+---
+
+# Length Generalization in Masked Diffusion Language Models
+
+From-scratch masked diffusion language models (MDLM-style absorbing state) and
+matched autoregressive baselines, trained on arithmetic and algorithmic tasks and
+evaluated far beyond the lengths seen in training.
+
+**No pretrained weights or tokenizers are used anywhere.**
+
+Code: https://github.com/GIND123/MSML612
+
+## Results
+
+_(pending)_

results/g1-agree-mdlm-s0.json

@@ -0,0 +1,41 @@
+{
+  "accuracy_by_passes": {
+    "1": 0.0,
+    "2": 0.0,
+    "4": 0.0,
+    "8": 0.0,
+    "16": 0.0,
+    "32": 0.0,
+    "64": 0.0,
+    "128": 1.0
+  },
+  "args": {
+    "objective": "mdlm",
+    "k_frac": 0.25,
+    "correct_weight": 1.0,
+    "task": "agree",
+    "coupling": 4,
+    "seq_len": 128,
+    "n_train": 200000,
+    "n_eval": 500,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 64,
+    "steps": 20000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 0,
+    "data_root": "data",
+    "eval_bs": 100,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g1-agree-mdlm-s0"
+  },
+  "minutes": 3.730570578575134,
+  "compute": {
+    "params": 10660608,
+    "train_tokens": 163840000,
+    "train_flops_approx": 10479804088320000,
+    "fwd_per_step": 1
+  }
+}

results/g1-agree-mdlm-s1.json

@@ -0,0 +1,41 @@
+{
+  "accuracy_by_passes": {
+    "1": 0.0,
+    "2": 0.0,
+    "4": 0.0,
+    "8": 0.0,
+    "16": 0.0,
+    "32": 0.0,
+    "64": 0.0,
+    "128": 1.0
+  },
+  "args": {
+    "objective": "mdlm",
+    "k_frac": 0.25,
+    "correct_weight": 1.0,
+    "task": "agree",
+    "coupling": 4,
+    "seq_len": 128,
+    "n_train": 200000,
+    "n_eval": 500,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 64,
+    "steps": 20000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 1,
+    "data_root": "data",
+    "eval_bs": 100,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g1-agree-mdlm-s1"
+  },
+  "minutes": 3.6946280360221864,
+  "compute": {
+    "params": 10660608,
+    "train_tokens": 163840000,
+    "train_flops_approx": 10479804088320000,
+    "fwd_per_step": 1
+  }
+}

results/g1-agree-selfcorrupt-s0.json

@@ -0,0 +1,41 @@
+{
+  "accuracy_by_passes": {
+    "1": 0.0,
+    "2": 0.0,
+    "4": 0.0,
+    "8": 0.0,
+    "16": 0.0,
+    "32": 0.0,
+    "64": 0.0,
+    "128": 0.0
+  },
+  "args": {
+    "objective": "selfcorrupt",
+    "k_frac": 0.25,
+    "correct_weight": 1.0,
+    "task": "agree",
+    "coupling": 4,
+    "seq_len": 128,
+    "n_train": 200000,
+    "n_eval": 500,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 64,
+    "steps": 20000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 0,
+    "data_root": "data",
+    "eval_bs": 100,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g1-agree-selfcorrupt-s0"
+  },
+  "minutes": 4.3024773756663,
+  "compute": {
+    "params": 10660608,
+    "train_tokens": 163840000,
+    "train_flops_approx": 10479804088320000,
+    "fwd_per_step": 2
+  }
+}

results/g1-agree-selfcorrupt-s1.json

@@ -0,0 +1,41 @@
+{
+  "accuracy_by_passes": {
+    "1": 0.0,
+    "2": 0.0,
+    "4": 0.0,
+    "8": 0.0,
+    "16": 0.0,
+    "32": 0.0,
+    "64": 0.0,
+    "128": 0.0
+  },
+  "args": {
+    "objective": "selfcorrupt",
+    "k_frac": 0.25,
+    "correct_weight": 1.0,
+    "task": "agree",
+    "coupling": 4,
+    "seq_len": 128,
+    "n_train": 200000,
+    "n_eval": 500,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 64,
+    "steps": 20000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 1,
+    "data_root": "data",
+    "eval_bs": 100,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g1-agree-selfcorrupt-s1"
+  },
+  "minutes": 4.266684722900391,
+  "compute": {
+    "params": 10660608,
+    "train_tokens": 163840000,
+    "train_flops_approx": 10479804088320000,
+    "fwd_per_step": 2
+  }
+}

results/g1-copy-mdlm-s0.json

@@ -0,0 +1,41 @@
+{
+  "accuracy_by_passes": {
+    "1": 1.0,
+    "2": 1.0,
+    "4": 1.0,
+    "8": 1.0,
+    "16": 1.0,
+    "32": 1.0,
+    "64": 1.0,
+    "128": 1.0
+  },
+  "args": {
+    "objective": "mdlm",
+    "k_frac": 0.25,
+    "correct_weight": 1.0,
+    "task": "copy",
+    "coupling": 4,
+    "seq_len": 128,
+    "n_train": 200000,
+    "n_eval": 500,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 64,
+    "steps": 20000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 0,
+    "data_root": "data",
+    "eval_bs": 100,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g1-copy-mdlm-s0"
+  },
+  "minutes": 3.757214069366455,
+  "compute": {
+    "params": 10660608,
+    "train_tokens": 163840000,
+    "train_flops_approx": 10479804088320000,
+    "fwd_per_step": 1
+  }
+}

results/g1-copy-mdlm-s1.json

@@ -0,0 +1,41 @@
+{
+  "accuracy_by_passes": {
+    "1": 1.0,
+    "2": 1.0,
+    "4": 1.0,
+    "8": 1.0,
+    "16": 1.0,
+    "32": 1.0,
+    "64": 1.0,
+    "128": 1.0
+  },
+  "args": {
+    "objective": "mdlm",
+    "k_frac": 0.25,
+    "correct_weight": 1.0,
+    "task": "copy",
+    "coupling": 4,
+    "seq_len": 128,
+    "n_train": 200000,
+    "n_eval": 500,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 64,
+    "steps": 20000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 1,
+    "data_root": "data",
+    "eval_bs": 100,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g1-copy-mdlm-s1"
+  },
+  "minutes": 3.7269402066866557,
+  "compute": {
+    "params": 10660608,
+    "train_tokens": 163840000,
+    "train_flops_approx": 10479804088320000,
+    "fwd_per_step": 1
+  }
+}

results/g1-copy-selfcorrupt-s0.json

@@ -0,0 +1,41 @@
+{
+  "accuracy_by_passes": {
+    "1": 0.0,
+    "2": 0.0,
+    "4": 0.0,
+    "8": 0.0,
+    "16": 0.0,
+    "32": 0.0,
+    "64": 0.0,
+    "128": 0.0
+  },
+  "args": {
+    "objective": "selfcorrupt",
+    "k_frac": 0.25,
+    "correct_weight": 1.0,
+    "task": "copy",
+    "coupling": 4,
+    "seq_len": 128,
+    "n_train": 200000,
+    "n_eval": 500,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 64,
+    "steps": 20000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 0,
+    "data_root": "data",
+    "eval_bs": 100,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g1-copy-selfcorrupt-s0"
+  },
+  "minutes": 4.29505068063736,
+  "compute": {
+    "params": 10660608,
+    "train_tokens": 163840000,
+    "train_flops_approx": 10479804088320000,
+    "fwd_per_step": 2
+  }
+}

results/g1-copy-selfcorrupt-s1.json

@@ -0,0 +1,41 @@
+{
+  "accuracy_by_passes": {
+    "1": 0.0,
+    "2": 0.0,
+    "4": 0.0,
+    "8": 0.0,
+    "16": 0.0,
+    "32": 0.0,
+    "64": 0.0,
+    "128": 0.0
+  },
+  "args": {
+    "objective": "selfcorrupt",
+    "k_frac": 0.25,
+    "correct_weight": 1.0,
+    "task": "copy",
+    "coupling": 4,
+    "seq_len": 128,
+    "n_train": 200000,
+    "n_eval": 500,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 64,
+    "steps": 20000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 1,
+    "data_root": "data",
+    "eval_bs": 100,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g1-copy-selfcorrupt-s1"
+  },
+  "minutes": 4.269979015986125,
+  "compute": {
+    "params": 10660608,
+    "train_tokens": 163840000,
+    "train_flops_approx": 10479804088320000,
+    "fwd_per_step": 2
+  }
+}

results/g1-parity-mdlm-s0.json

@@ -0,0 +1,41 @@
+{
+  "accuracy_by_passes": {
+    "1": 0.0,
+    "2": 0.0,
+    "4": 0.0,
+    "8": 0.0,
+    "16": 0.0,
+    "32": 0.0,
+    "64": 0.0,
+    "128": 1.0
+  },
+  "args": {
+    "objective": "mdlm",
+    "k_frac": 0.25,
+    "correct_weight": 1.0,
+    "task": "parity",
+    "coupling": 4,
+    "seq_len": 128,
+    "n_train": 200000,
+    "n_eval": 500,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 64,
+    "steps": 20000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 0,
+    "data_root": "data",
+    "eval_bs": 100,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g1-parity-mdlm-s0"
+  },
+  "minutes": 3.7436817288398743,
+  "compute": {
+    "params": 10660608,
+    "train_tokens": 163840000,
+    "train_flops_approx": 10479804088320000,
+    "fwd_per_step": 1
+  }
+}

results/g1-parity-mdlm-s1.json

@@ -0,0 +1,41 @@
+{
+  "accuracy_by_passes": {
+    "1": 0.0,
+    "2": 0.0,
+    "4": 0.0,
+    "8": 0.0,
+    "16": 0.0,
+    "32": 0.0,
+    "64": 0.0,
+    "128": 1.0
+  },
+  "args": {
+    "objective": "mdlm",
+    "k_frac": 0.25,
+    "correct_weight": 1.0,
+    "task": "parity",
+    "coupling": 4,
+    "seq_len": 128,
+    "n_train": 200000,
+    "n_eval": 500,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 64,
+    "steps": 20000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 1,
+    "data_root": "data",
+    "eval_bs": 100,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g1-parity-mdlm-s1"
+  },
+  "minutes": 3.6899612665176393,
+  "compute": {
+    "params": 10660608,
+    "train_tokens": 163840000,
+    "train_flops_approx": 10479804088320000,
+    "fwd_per_step": 1
+  }
+}

results/g1-parity-selfcorrupt-s0.json

@@ -0,0 +1,41 @@
+{
+  "accuracy_by_passes": {
+    "1": 0.0,
+    "2": 0.0,
+    "4": 0.0,
+    "8": 0.0,
+    "16": 0.0,
+    "32": 0.0,
+    "64": 0.0,
+    "128": 0.0
+  },
+  "args": {
+    "objective": "selfcorrupt",
+    "k_frac": 0.25,
+    "correct_weight": 1.0,
+    "task": "parity",
+    "coupling": 4,
+    "seq_len": 128,
+    "n_train": 200000,
+    "n_eval": 500,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 64,
+    "steps": 20000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 0,
+    "data_root": "data",
+    "eval_bs": 100,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g1-parity-selfcorrupt-s0"
+  },
+  "minutes": 4.269179002443949,
+  "compute": {
+    "params": 10660608,
+    "train_tokens": 163840000,
+    "train_flops_approx": 10479804088320000,
+    "fwd_per_step": 2
+  }
+}

results/g1-parity-selfcorrupt-s1.json

@@ -0,0 +1,41 @@
+{
+  "accuracy_by_passes": {
+    "1": 0.0,
+    "2": 0.0,
+    "4": 0.0,
+    "8": 0.0,
+    "16": 0.0,
+    "32": 0.0,
+    "64": 0.0,
+    "128": 0.0
+  },
+  "args": {
+    "objective": "selfcorrupt",
+    "k_frac": 0.25,
+    "correct_weight": 1.0,
+    "task": "parity",
+    "coupling": 4,
+    "seq_len": 128,
+    "n_train": 200000,
+    "n_eval": 500,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 64,
+    "steps": 20000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 1,
+    "data_root": "data",
+    "eval_bs": 100,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g1-parity-selfcorrupt-s1"
+  },
+  "minutes": 4.223804748058319,
+  "compute": {
+    "params": 10660608,
+    "train_tokens": 163840000,
+    "train_flops_approx": 10479804088320000,
+    "fwd_per_step": 2
+  }
+}

slurm/finalize.sh

@@ -0,0 +1,103 @@
+#!/bin/bash
+# Runs on the cluster after a grid finishes (SLURM dependency), so results are
+# collected, backed up and published even if the laptop is closed, the network
+# drops, or the session that launched the work is long gone.
+#
+#   1. regenerate figures + tables from every run
+#   2. back up to HOME (the only nightly-backed-up tier)
+#   3. back up weights to SHELL (1 TB, no 90-day purge unlike scratch)
+#   4. push to the Hugging Face Hub, with retries
+#
+#SBATCH -J finalize
+#SBATCH -t 03:00:00
+#SBATCH -c 4
+#SBATCH --mem=16g
+#SBATCH -p standard
+#SBATCH -o logs/finalize-%j.out
+#SBATCH --open-mode=append
+
+source /etc/profile
+module load pytorch/2.0.1
+
+USER_SCR=/scratch/zt1/project/msml612/user/$USER
+STAR=$USER_SCR/pcmd
+export PYTHONPATH=$USER_SCR/pylibs:$PYTHONPATH
+# Compute nodes cannot reach HF's Xet CAS endpoint; force the plain LFS path.
+export HF_HUB_DISABLE_XET=1
+
+HOME_BK=$HOME/pcmd-backup
+SHELL_BK=/afs/shell.umd.edu/project/msml612/user-$USER/msml612-weights
+
+echo "=== 1. collect + plot ==="
+cd $STAR/src
+RUNS=$STAR/runs OUT=$STAR/figures python collect.py || echo "collect failed (continuing)"
+
+echo "=== 2. back up to HOME (backed up nightly) ==="
+mkdir -p $HOME_BK/{results,figures,src,slurm}
+cp -f $STAR/figures/* $HOME_BK/figures/ 2>/dev/null
+cp -f $STAR/src/*.py $HOME_BK/src/ 2>/dev/null
+cp -f $STAR/*.sh $HOME_BK/slurm/ 2>/dev/null
+for d in $STAR/runs/*/; do
+  n=$(basename $d)
+  [ -f "$d/result.json" ] && cp -f "$d/result.json" "$HOME_BK/results/$n.json"
+done
+echo "  results backed up: $(ls $HOME_BK/results | wc -l)"
+
+echo "=== 3. back up weights to SHELL (not purged) ==="
+if mkdir -p $SHELL_BK 2>/dev/null; then
+  for d in $STAR/runs/*/; do
+    n=$(basename $d)
+    [ -f "$d/model.pt" ] && cp -f "$d/model.pt" "$SHELL_BK/$n.pt" 2>/dev/null
+  done
+  echo "  weights backed up: $(ls $SHELL_BK 2>/dev/null | wc -l)"
+else
+  echo "  SHELL not writable from this node - weights remain on scratch"
+fi
+
+echo "=== 4. push to Hugging Face (with retries) ==="
+# A transient network failure must not mean the results never reach the Hub, so
+# retry with backoff rather than giving up on the first error.
+for attempt in 1 2 3 4 5; do
+  python - <<'PY' && break
+import os, pathlib, sys
+from huggingface_hub import HfApi
+
+tok = pathlib.Path(os.path.expanduser("~/.hf_token")).read_text().strip()
+star = os.path.expanduser(f"/scratch/zt1/project/msml612/user/{os.environ['USER']}/pcmd")
+stage = pathlib.Path(os.path.expanduser("~/pcmd-backup"))
+repo = "GOVINDFROM/parallel-consistent-masked-diffusion"
+
+summary = pathlib.Path(star, "figures", "summary.md")
+table = summary.read_text() if summary.exists() else "_(pending)_"
+
+(stage / "README.md").write_text(f"""---
+license: apache-2.0
+tags: [masked-diffusion, discrete-diffusion, length-generalization, positional-encoding, from-scratch]
+---
+
+# Length Generalization in Masked Diffusion Language Models
+
+From-scratch masked diffusion language models (MDLM-style absorbing state) and
+matched autoregressive baselines, trained on arithmetic and algorithmic tasks and
+evaluated far beyond the lengths seen in training.
+
+**No pretrained weights or tokenizers are used anywhere.**
+
+Code: https://github.com/GIND123/MSML612
+
+## Results
+
+{table}
+""")
+
+api = HfApi(token=tok)
+api.create_repo(repo, repo_type="model", private=True, exist_ok=True)
+api.upload_folder(folder_path=str(stage), repo_id=repo, repo_type="model",
+                  commit_message="Automated push from Zaratan")
+print("pushed:", f"https://huggingface.co/{repo}")
+PY
+  echo "  push attempt $attempt failed; sleeping $((attempt * 60))s"
+  sleep $((attempt * 60))
+done
+
+echo "=== finalize complete ==="

slurm/grid1.sh

@@ -0,0 +1,36 @@
+#!/bin/bash
+# Grid 1 — learnability gate and the first quality-vs-parallelism curves.
+#
+# Doubles as a gate: if a model cannot solve a probe even at FULLY SEQUENTIAL
+# decoding (128 passes for a 128-token canvas), the task is not learnable at
+# this scale and no parallel-decoding claim built on it would mean anything.
+# Launching a large grid before establishing this is the mistake that cost a
+# day on the previous project.
+#SBATCH -J g1
+#SBATCH -p gpu-h100
+#SBATCH --gpus=h100:1
+#SBATCH -c 8
+#SBATCH --mem=64g
+#SBATCH -t 04:00:00
+#SBATCH -o logs/g1-%A_%a.out
+#SBATCH --array=0-11%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
+TASKS=(copy agree parity); OBJ=(mdlm selfcorrupt); SEEDS=(0 1)
+T=${TASKS[$((I / 4))]}
+O=${OBJ[$(((I % 4) / 2))]}
+S=${SEEDS[$((I % 2))]}
+
+NAME="g1-$T-$O-s$S"
+echo "=== $NAME ==="
+$PY train.py --task $T --objective $O --coupling 4 --seed $S \
+  --seq_len 128 --n_train 200000 --n_eval 500 --eval_bs 100 \
+  --d 384 --layers 6 --heads 6 --bs 64 --steps 20000 --lr 3e-4 \
+  --out $SCR/pcmd/runs/$NAME

slurm/watchdog.sh

@@ -0,0 +1,42 @@
+#!/bin/bash
+# Safety net: re-runs finalize periodically, independent of any grid.
+#
+# finalize.sh is chained to each grid, but that only fires when a grid ends
+# cleanly. If a grid is cancelled, times out, or the chained job itself fails,
+# results would sit on scratch un-published. This resubmits itself every few
+# hours so collection and the Hugging Face push keep happening regardless of
+# what the laptop, the network, or the login session are doing.
+#
+# Start once with:   sbatch watchdog.sh
+# Stop with:         scancel -n watchdog
+#
+#SBATCH -J watchdog
+#SBATCH -t 00:45:00
+#SBATCH -c 2
+#SBATCH --mem=8g
+#SBATCH -p standard
+#SBATCH -o logs/watchdog-%j.out
+
+SCR=/scratch/zt1/project/msml612/user/$USER
+cd $SCR/pcmd
+
+RUNS_NOW=$(ls $SCR/pcmd/runs 2>/dev/null | wc -l)
+LAST=$(cat $SCR/pcmd/.watchdog_last 2>/dev/null || echo 0)
+echo "$(date): runs=$RUNS_NOW last_published=$LAST"
+
+# Only spend effort when something actually changed.
+if [ "$RUNS_NOW" -ne "$LAST" ]; then
+  echo "new results detected - running finalize"
+  bash $SCR/pcmd/finalize.sh
+  echo "$RUNS_NOW" > $SCR/pcmd/.watchdog_last
+else
+  echo "nothing new since last publish"
+fi
+
+# Re-arm unless someone asked it to stop.
+if [ ! -f $SCR/pcmd/.watchdog_stop ]; then
+  sbatch --begin=now+3hours $SCR/pcmd/watchdog.sh
+  echo "re-armed for +3h"
+else
+  echo "stop file present - not re-arming"
+fi

src/data.py

@@ -0,0 +1,132 @@
+"""Datasets: text8 for language modelling, plus dependency-controlled probes.
+
+text8 is the standard small-scale benchmark for discrete diffusion language
+models (D3PM, SEDD, MDLM all report bits-per-character on it), which is what
+makes a from-scratch comparison meaningful.
+
+The probes exist because text8 alone cannot isolate the phenomenon we study.
+Parallel decoding fails when simultaneously-revealed tokens depend on each
+other, so we need tasks where that dependence is a knob rather than a property
+of English.
+"""
+import os
+import numpy as np
+
+TEXT8_URL = "http://mattmahoney.net/dc/text8.zip"
+# text8 is lowercase letters plus space: a 27-symbol alphabet, no tokenizer.
+ALPHABET = " abcdefghijklmnopqrstuvwxyz"
+SPECIALS = ["[MASK]", "[PAD]"]
+
+
+class CharTokenizer:
+    def __init__(self):
+        self.itos = SPECIALS + list(ALPHABET)
+        self.stoi = {c: i for i, c in enumerate(self.itos)}
+        self.mask = self.stoi["[MASK]"]
+        self.pad = self.stoi["[PAD]"]
+
+    def __len__(self):
+        return len(self.itos)
+
+    def encode(self, s):
+        return [self.stoi[c] for c in s if c in self.stoi]
+
+    def decode(self, ids):
+        return "".join(self.itos[i] for i in ids)
+
+
+def load_text8(root, seq_len=256):
+    """Standard 90M/5M/5M character split used throughout the literature."""
+    raw = os.path.join(root, "text8")
+    if not os.path.exists(raw):
+        os.makedirs(root, exist_ok=True)
+        import zipfile, io, requests   # only needed on first download
+        r = requests.get(TEXT8_URL, timeout=300)
+        r.raise_for_status()
+        with zipfile.ZipFile(io.BytesIO(r.content)) as z:
+            z.extract("text8", root)
+    text = open(raw).read()
+    assert len(text) == 100_000_000, f"unexpected text8 length {len(text)}"
+
+    tok = CharTokenizer()
+    cache = os.path.join(root, "text8.npy")
+    if os.path.exists(cache):
+        ids = np.load(cache)
+    else:
+        ids = np.array(tok.encode(text), dtype=np.uint8)
+        np.save(cache, ids)
+
+    tr, va = ids[:90_000_000], ids[90_000_000:95_000_000]
+    te = ids[95_000_000:]
+
+    def chunk(a):
+        n = len(a) // seq_len
+        return a[: n * seq_len].reshape(n, seq_len).astype(np.int64)
+
+    return chunk(tr), chunk(va), chunk(te), tok
+
+
+# ---------------------------------------------------------------------------
+# Dependency-controlled probes.
+#
+# Each item is a sequence whose tokens must agree with each other. `coupling`
+# sets how many positions are bound together, which is the exact quantity that
+# breaks parallel decoding: revealing two mutually-dependent positions in one
+# step forces the model to guess them independently.
+# ---------------------------------------------------------------------------
+
+def build_probe(n, seq_len, coupling, seed=0, task="copy"):
+    """
+    copy    : second half must equal the first half. Each answer position is
+              fixed by one context position, so answer positions are mutually
+              INDEPENDENT given the context - parallel decoding should be safe.
+    agree   : within each group of `coupling` positions all tokens must be equal
+              AND adjacent groups must differ. The value of a group is free, so
+              its members are mutually DEPENDENT: revealing two of them in one
+              step is a 1/V coin flip. This is the failure mode, isolated.
+    parity  : groups are a run of bits whose last element is the XOR of the
+              rest, with adjacent groups forced to differ, so the group is
+              jointly constrained without any member being determined alone.
+
+    The adjacent-groups-differ rule matters: without it a model that emits one
+    constant token everywhere satisfies "all members agree" perfectly, and an
+    untrained model scores 100% at maximum parallelism. That degenerate solution
+    made the first version of this metric meaningless.
+    """
+    tok = CharTokenizer()
+    rng = np.random.default_rng(seed)
+    V = 2 + 26                       # specials + letters
+    lo, hi = 2, V                    # sample letters only
+
+    X = np.zeros((n, seq_len), dtype=np.int64)
+    for i in range(n):
+        if task == "copy":
+            half = seq_len // 2
+            src = rng.integers(lo, hi, size=half)
+            X[i, :half] = src
+            X[i, half:half * 2] = src
+        elif task == "agree":
+            g, prev = coupling, -1
+            for s in range(0, seq_len, g):
+                v = int(rng.integers(lo, hi))
+                while v == prev:                      # adjacent groups must differ
+                    v = int(rng.integers(lo, hi))
+                X[i, s:s + g] = v
+                prev = v
+        elif task == "parity":
+            g, prev = coupling, None
+            for s in range(0, seq_len, g):
+                end = min(s + g, seq_len)
+                if end - s < 2:
+                    X[i, s:end] = lo
+                    continue
+                while True:
+                    bits = list(rng.integers(0, 2, size=end - s - 1))
+                    vals = bits + [int(sum(bits) % 2)]
+                    if vals != prev:                  # adjacent groups must differ
+                        break
+                X[i, s:end] = [lo + v for v in vals]
+                prev = vals
+        else:
+            raise ValueError(task)
+    return X, tok

src/diffusion.py

@@ -0,0 +1,159 @@
+"""Masked diffusion: training objectives and decoding strategies.
+
+The central problem this file exists to study:
+
+  MDLM training teaches each masked position its MARGINAL p(x_i | observed).
+  Parallel decoding reveals k positions at once by sampling the PRODUCT of
+  those marginals. The truth is the joint, and the gap is exactly the mutual
+  information among the revealed positions. Quality therefore falls as the
+  parallelism rises - which is the only reason to use diffusion in the first
+  place.
+
+Every published fix reorders or gates commits at INFERENCE. The model has still
+never seen a partially-decoded state containing its own correlated mistakes.
+`self_corrupted_loss` changes that at training time.
+"""
+import torch
+import torch.nn.functional as F
+
+
+# --------------------------------------------------------------- training ---
+def mdlm_loss(model, x0, tok, t_min=None, fixed_t=None):
+    """Standard MDLM objective: mask each token independently with prob t,
+    predict the originals, weight by 1/t."""
+    B, L = x0.shape
+    dev = x0.device
+    t_min = t_min if t_min is not None else 1.0 / L
+    t = (torch.full((B,), fixed_t, device=dev) if fixed_t
+         else t_min + (1.0 - t_min) * torch.rand(B, device=dev))
+    m = torch.rand(B, L, device=dev) < t[:, None]
+    m[torch.arange(B, device=dev), torch.randint(0, L, (B,), device=dev)] = True
+
+    xt = torch.where(m, torch.full_like(x0, tok.mask), x0)
+    logits = model(xt)
+    ce = F.cross_entropy(logits.reshape(-1, logits.size(-1)), x0.reshape(-1),
+                         reduction="none").view(B, L)
+    return ((1.0 / t) * (ce * m).sum(1) / L).mean()
+
+
+def self_corrupted_loss(model, x0, tok, k_frac=0.25, t_min=None,
+                        correct_weight=1.0):
+    """OUR METHOD. Train on states that parallel decoding actually produces.
+
+    1. mask as usual and take a NO-GRAD pass to get the current marginals
+    2. reveal a random subset of the masked positions *in parallel*, sampling
+       each independently - i.e. commit exactly the kind of jointly-inconsistent
+       tokens that parallel decoding commits
+    3. train on the remaining masked positions given that self-generated
+       context, so the model learns p(x_i | observed, its own parallel samples)
+    4. additionally train it to REPAIR the committed tokens, which turns a
+       wrong commit from an absorbing error into a recoverable one
+
+    The model is therefore optimised for the distribution it meets at inference
+    rather than for ground-truth context it will never see.
+    """
+    B, L = x0.shape
+    dev = x0.device
+    t_min = t_min if t_min is not None else 1.0 / L
+    t = t_min + (1.0 - t_min) * torch.rand(B, device=dev)
+    m = torch.rand(B, L, device=dev) < t[:, None]
+    m[torch.arange(B, device=dev), torch.randint(0, L, (B,), device=dev)] = True
+    xt = torch.where(m, torch.full_like(x0, tok.mask), x0)
+
+    with torch.no_grad():
+        probs = model(xt).softmax(-1)
+        samp = torch.distributions.Categorical(probs=probs).sample()
+
+    # commit a random subset of masked positions from the model's own samples
+    commit = m & (torch.rand(B, L, device=dev) < k_frac)
+    x_self = torch.where(commit, samp, xt)
+    still = m & ~commit
+
+    logits = model(x_self)
+    ce = F.cross_entropy(logits.reshape(-1, logits.size(-1)), x0.reshape(-1),
+                         reduction="none").view(B, L)
+
+    denom = still.sum(1).clamp_min(1)
+    main = ((1.0 / t) * (ce * still).sum(1) / denom).mean()
+    # repair term: committed positions are supervised against ground truth too
+    rep = (ce * commit).sum(1) / commit.sum(1).clamp_min(1)
+    return main + correct_weight * rep.mean()
+
+
+# --------------------------------------------------------------- decoding ---
+@torch.no_grad()
+def decode(model, shape, tok, steps, device, strategy="confidence",
+           temperature=0.0, entropy_tau=None, x_init=None, keep_trace=False):
+    """Reveal the canvas over `steps` passes. Fewer passes = more tokens
+    committed per pass = more parallelism = more conditional-independence error.
+
+    strategies
+      confidence : reveal the most confident masked positions (the standard)
+      random     : reveal a random subset (ablation: is confidence doing work?)
+      entropy    : reveal only positions whose entropy is below `entropy_tau`,
+                   an inference-time baseline for dependency-aware revealing
+    """
+    B, L = shape
+    x = torch.full((B, L), tok.mask, device=device, dtype=torch.long)
+    if x_init is not None:
+        x = x_init.clone()
+    trace = []
+
+    for s in range(steps, 0, -1):
+        masked = x == tok.mask
+        if not masked.any():
+            break
+        logits = model(x)
+        probs = logits.softmax(-1)
+        if temperature > 0:
+            pred = torch.distributions.Categorical(
+                logits=logits / temperature).sample()
+            conf = probs.gather(-1, pred[..., None]).squeeze(-1)
+        else:
+            conf, pred = probs.max(-1)
+
+        if strategy == "random":
+            score = torch.rand_like(conf)
+        elif strategy == "entropy":
+            ent = -(probs * probs.clamp_min(1e-9).log()).sum(-1)
+            score = -ent
+        else:
+            score = conf
+        score = score.masked_fill(~masked, -1e9)
+
+        target_left = int(L * (s - 1) / steps)
+        for b in range(B):
+            nm = int(masked[b].sum())
+            k = max(0, nm - target_left)
+            if strategy == "entropy" and entropy_tau is not None and k > 0:
+                ent_b = -(probs[b] * probs[b].clamp_min(1e-9).log()).sum(-1)
+                allowed = int(((ent_b < entropy_tau) & masked[b]).sum())
+                k = max(1, min(k, allowed)) if allowed else 1
+            if k:
+                idx = score[b].topk(k).indices
+                x[b, idx] = pred[b, idx]
+        if keep_trace:
+            trace.append(x.clone())
+    still = x == tok.mask
+    if still.any():
+        x[still] = pred[still]
+    return (x, trace) if keep_trace else x
+
+
+@torch.no_grad()
+def nll_bpc(model, batches, tok, device, mc=8):
+    """Monte-Carlo estimate of the MDLM bound, reported as bits per character -
+    the standard text8 metric."""
+    tot, n = 0.0, 0
+    for x0 in batches:
+        x0 = x0.to(device)
+        B, L = x0.shape
+        for _ in range(mc):
+            t = (1.0 / L) + (1 - 1.0 / L) * torch.rand(B, device=device)
+            m = torch.rand(B, L, device=device) < t[:, None]
+            xt = torch.where(m, torch.full_like(x0, tok.mask), x0)
+            ce = F.cross_entropy(model(xt).reshape(-1, len(tok)),
+                                 x0.reshape(-1), reduction="none").view(B, L)
+            tot += float(((1.0 / t) * (ce * m).sum(1) / L).sum())
+            n += B
+    return tot / max(n, 1) / torch.log(torch.tensor(2.0)).item()

src/model.py

@@ -0,0 +1,135 @@
+"""Transformer denoiser for masked diffusion.
+
+Bidirectional by default (a masked diffusion denoiser must see both sides);
+`causal=True` gives the matched autoregressive baseline. Rotary is the default
+positional encoding - it was the strongest option in our previous study.
+"""
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def sinusoidal(max_len, d):
+    pe = torch.zeros(max_len, d)
+    pos = torch.arange(max_len).unsqueeze(1).float()
+    div = torch.exp(torch.arange(0, d, 2).float() * (-math.log(10000.0) / d))
+    pe[:, 0::2] = torch.sin(pos * div)
+    pe[:, 1::2] = torch.cos(pos * div)
+    return pe
+
+
+def alibi_slopes(n_heads):
+    start = 2 ** (-(2 ** -(math.log2(n_heads) - 3))) if math.log2(n_heads).is_integer() else 0.5
+    return torch.tensor([start * (start ** i) for i in range(n_heads)])
+
+
+def apply_rope(x, cos, sin):
+    # x: (B, H, L, Dh)
+    x1, x2 = x[..., 0::2], x[..., 1::2]
+    rx = torch.stack([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1)
+    return rx.flatten(-2)
+
+
+class Attention(nn.Module):
+    def __init__(self, d, n_heads, pe, causal, max_len):
+        super().__init__()
+        self.h, self.dh, self.pe, self.causal = n_heads, d // n_heads, pe, causal
+        self.qkv = nn.Linear(d, 3 * d, bias=False)
+        self.proj = nn.Linear(d, d, bias=False)
+        if pe == "alibi":
+            self.register_buffer("slopes", alibi_slopes(n_heads), persistent=False)
+        if pe == "rope":
+            inv = 1.0 / (10000 ** (torch.arange(0, self.dh, 2).float() / self.dh))
+            ang = torch.arange(max_len).float()[:, None] * inv[None, :]
+            self.register_buffer("cos", ang.cos(), persistent=False)
+            self.register_buffer("sin", ang.sin(), persistent=False)
+
+    def forward(self, x, pad_mask, pos_ids=None):
+        B, L, D = x.shape
+        q, k, v = self.qkv(x).chunk(3, dim=-1)
+        q = q.view(B, L, self.h, self.dh).transpose(1, 2)
+        k = k.view(B, L, self.h, self.dh).transpose(1, 2)
+        v = v.view(B, L, self.h, self.dh).transpose(1, 2)
+
+        if self.pe == "rope":
+            if pos_ids is None:
+                cos, sin = self.cos[:L][None, None], self.sin[:L][None, None]
+            else:
+                cos, sin = self.cos[pos_ids][:, None], self.sin[pos_ids][:, None]
+            q, k = apply_rope(q, cos, sin), apply_rope(k, cos, sin)
+
+        att = (q @ k.transpose(-2, -1)) / math.sqrt(self.dh)
+
+        if self.pe == "alibi":
+            if pos_ids is None:
+                pos = torch.arange(L, device=x.device)
+                rel = (pos[None, :] - pos[:, None]).abs().float()[None]
+            else:
+                rel = (pos_ids[:, None, :] - pos_ids[:, :, None]).abs().float()
+            att = att - self.slopes.view(1, -1, 1, 1) * rel[:, None]
+
+        if self.causal:
+            cm = torch.ones(L, L, device=x.device, dtype=torch.bool).tril()
+            att = att.masked_fill(~cm[None, None], float("-inf"))
+        if pad_mask is not None:
+            att = att.masked_fill(~pad_mask[:, None, None, :], float("-inf"))
+
+        att = att.softmax(-1)
+        out = (att @ v).transpose(1, 2).reshape(B, L, D)
+        return self.proj(out)
+
+
+class Block(nn.Module):
+    def __init__(self, d, n_heads, pe, causal, max_len):
+        super().__init__()
+        self.n1, self.n2 = nn.LayerNorm(d), nn.LayerNorm(d)
+        self.attn = Attention(d, n_heads, pe, causal, max_len)
+        self.mlp = nn.Sequential(nn.Linear(d, 4 * d), nn.GELU(), nn.Linear(4 * d, d))
+
+    def forward(self, x, pad_mask, pos_ids=None):
+        x = x + self.attn(self.n1(x), pad_mask, pos_ids)
+        return x + self.mlp(self.n2(x))
+
+
+class Transformer(nn.Module):
+    def __init__(self, vocab, d=384, n_layers=6, n_heads=6, pe="ape", causal=False, max_len=256):
+        super().__init__()
+        self.pe_kind, self.causal = pe, causal
+        self.emb = nn.Embedding(vocab, d)
+        # Segment embedding: which part of the equation a token belongs to.
+        # Place-value ids deliberately collide across operands and answer, so a
+        # bidirectional model needs this to tell those tokens apart; a causal
+        # model gets the same information free from the attention mask.
+
+        if pe == "ape":
+            self.pos = nn.Embedding(max_len, d)
+        elif pe == "sin":
+            self.register_buffer("pos_sin", sinusoidal(max_len, d), persistent=False)
+        self.blocks = nn.ModuleList([Block(d, n_heads, pe, causal, max_len) for _ in range(n_layers)])
+        self.norm = nn.LayerNorm(d)
+        self.head = nn.Linear(d, vocab, bias=False)
+        self.apply(self._init)
+
+    @staticmethod
+    def _init(m):
+        if isinstance(m, (nn.Linear, nn.Embedding)):
+            nn.init.normal_(m.weight, std=0.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.zeros_(m.bias)
+
+    def forward(self, idx, pad_mask=None, pos_ids=None):
+        B, L = idx.shape
+        x = self.emb(idx)
+        if self.pe_kind == "ape":
+            p = torch.arange(L, device=idx.device)[None] if pos_ids is None else pos_ids
+            x = x + self.pos(p)
+        elif self.pe_kind == "sin":
+            p = torch.arange(L, device=idx.device)[None] if pos_ids is None else pos_ids
+            x = x + self.pos_sin[p]
+        for b in self.blocks:
+            x = b(x, pad_mask, pos_ids)
+        return self.head(self.norm(x))
+
+    def n_params(self):
+        return sum(p.numel() for p in self.parameters() if p.requires_grad)

src/train.py

@@ -0,0 +1,171 @@
+"""Train a masked diffusion LM from scratch and measure quality vs parallelism.
+
+The headline axis is accuracy (or bits-per-character) as a function of the
+number of denoising passes. Fewer passes means more tokens committed per pass,
+which is the regime where the conditional-independence assumption bites and
+where every acceleration method must eventually be judged.
+"""
+import argparse, json, math, os, time
+import numpy as np
+import torch
+
+from data import load_text8, build_probe, CharTokenizer
+from model import Transformer
+import diffusion as dfn
+
+
+def get_args():
+    p = argparse.ArgumentParser()
+    p.add_argument("--objective", choices=["mdlm", "selfcorrupt"], default="mdlm",
+                   help="mdlm = standard marginal training; selfcorrupt = ours")
+    p.add_argument("--k_frac", type=float, default=0.25,
+                   help="fraction of masked positions committed from the model's "
+                        "own parallel samples during self-corrupted training")
+    p.add_argument("--correct_weight", type=float, default=1.0)
+    p.add_argument("--task", choices=["text8", "copy", "agree", "parity"], default="agree")
+    p.add_argument("--coupling", type=int, default=4,
+                   help="probe tasks: how many positions are bound together")
+    p.add_argument("--seq_len", type=int, default=128)
+    p.add_argument("--n_train", type=int, default=200000)
+    p.add_argument("--n_eval", type=int, default=1000)
+    p.add_argument("--d", type=int, default=384)
+    p.add_argument("--layers", type=int, default=6)
+    p.add_argument("--heads", type=int, default=6)
+    p.add_argument("--pe", default="rope")
+    p.add_argument("--bs", type=int, default=64)
+    p.add_argument("--steps", type=int, default=20000)
+    p.add_argument("--lr", type=float, default=3e-4)
+    p.add_argument("--warmup", type=int, default=500)
+    p.add_argument("--seed", type=int, default=0)
+    p.add_argument("--data_root", default="data")
+    p.add_argument("--eval_bs", type=int, default=100)
+    p.add_argument("--out", default="runs/dev")
+    return p.parse_args()
+
+
+# The whole point: sweep the number of passes from fully sequential (L) down to
+# maximally parallel (1) and watch where each method falls over.
+PASS_GRID = [1, 2, 4, 8, 16, 32, 64, 128]
+
+
+def probe_accuracy(x, task, coupling, tok):
+    """Fraction of sequences satisfying the task constraint exactly."""
+    ok = 0
+    for row in x.tolist():
+        if task == "copy":
+            h = len(row) // 2
+            ok += row[:h] == row[h:h * 2]
+        elif task == "agree":
+            groups = [row[s:s + coupling] for s in range(0, len(row), coupling)
+                      if len(row[s:s + coupling]) == coupling]
+            good = all(len(set(g)) == 1 for g in groups) and \
+                   all(groups[j][0] != groups[j + 1][0] for j in range(len(groups) - 1))
+            ok += good
+        elif task == "parity":
+            groups = [row[s:s + coupling] for s in range(0, len(row), coupling)
+                      if len(row[s:s + coupling]) >= 2]
+            good = True
+            for g in groups:
+                bits = [v - 2 for v in g]
+                if any(b not in (0, 1) for b in bits) or sum(bits[:-1]) % 2 != bits[-1]:
+                    good = False
+                    break
+            if good:
+                good = all(groups[j] != groups[j + 1] for j in range(len(groups) - 1))
+            ok += good
+    return ok / len(x)
+
+
+@torch.no_grad()
+def evaluate(model, args, tok, device, eval_data):
+    model.eval()
+    out = {}
+    for steps in PASS_GRID:
+        if steps > args.seq_len:
+            continue
+        accs = []
+        for i in range(0, args.n_eval, args.eval_bs):
+            b = min(args.eval_bs, args.n_eval - i)
+            x_init = None
+            if args.task == "copy":
+                # Conditional: the first half is GIVEN, so there is one correct
+                # completion. Unconditionally a constant sequence copies itself,
+                # which any degenerate model satisfies.
+                from data import build_probe as _bp
+                src, _ = _bp(b, args.seq_len, args.coupling, seed=9000 + i, task="copy")
+                x_init = torch.full((b, args.seq_len), tok.mask, device=device, dtype=torch.long)
+                h = args.seq_len // 2
+                x_init[:, :h] = torch.from_numpy(src[:, :h]).to(device)
+            x = dfn.decode(model, (b, args.seq_len), tok, steps, device, x_init=x_init)
+            if args.task == "text8":
+                accs.append(0.0)      # quality for text8 is reported via bpc
+            else:
+                accs.append(probe_accuracy(x.cpu(), args.task, args.coupling, tok))
+        out[str(steps)] = float(np.mean(accs))
+    res = {"accuracy_by_passes": out}
+    if args.task == "text8" and eval_data is not None:
+        res["bpc"] = dfn.nll_bpc(model, eval_data[:8], tok, device)
+    model.train()
+    return res
+
+
+def main():
+    a = get_args()
+    torch.manual_seed(a.seed); np.random.seed(a.seed)
+    dev = "cuda" if torch.cuda.is_available() else "cpu"
+    os.makedirs(a.out, exist_ok=True)
+
+    if a.task == "text8":
+        tr, va, te, tok = load_text8(a.data_root, a.seq_len)
+        train = torch.from_numpy(tr)
+        eval_data = [torch.from_numpy(va[i:i + a.eval_bs]) for i in range(0, 800, a.eval_bs)]
+    else:
+        X, tok = build_probe(a.n_train, a.seq_len, a.coupling, seed=a.seed, task=a.task)
+        train = torch.from_numpy(X)
+        eval_data = None
+
+    model = Transformer(len(tok), a.d, a.layers, a.heads, a.pe,
+                        causal=False, max_len=a.seq_len + 8).to(dev)
+    print(f"[{a.objective}/{a.task}] params={model.n_params()/1e6:.2f}M "
+          f"vocab={len(tok)} seq={a.seq_len} device={dev}", flush=True)
+
+    opt = torch.optim.AdamW(model.parameters(), lr=a.lr, weight_decay=0.01)
+    sched = torch.optim.lr_scheduler.LambdaLR(
+        opt, lambda s: min(1.0, (s + 1) / a.warmup) *
+                       0.5 * (1 + math.cos(math.pi * min(1.0, s / a.steps))))
+    amp = dict(device_type="cuda", dtype=torch.bfloat16) if dev == "cuda" \
+        else dict(device_type="cpu", enabled=False)
+
+    t0 = time.time()
+    for step in range(a.steps):
+        idx = torch.randint(0, len(train), (a.bs,))
+        x0 = train[idx].to(dev)
+        opt.zero_grad(set_to_none=True)
+        with torch.autocast(**amp):
+            loss = (dfn.self_corrupted_loss(model, x0, tok, a.k_frac,
+                                            correct_weight=a.correct_weight)
+                    if a.objective == "selfcorrupt"
+                    else dfn.mdlm_loss(model, x0, tok))
+        loss.backward()
+        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+        opt.step(); sched.step()
+        if step % 1000 == 0:
+            print(f"step {step:6d} loss {loss.item():.4f} ({time.time()-t0:.0f}s)", flush=True)
+
+    res = evaluate(model, a, tok, dev, eval_data)
+    n_par = model.n_params()
+    res.update({"args": vars(a), "minutes": (time.time() - t0) / 60,
+                "compute": {"params": n_par,
+                            "train_tokens": a.steps * a.bs * a.seq_len,
+                            "train_flops_approx": 6 * n_par * a.steps * a.bs * a.seq_len,
+                            # self-corruption costs one extra no-grad pass
+                            "fwd_per_step": 2 if a.objective == "selfcorrupt" else 1}})
+    json.dump(res, open(os.path.join(a.out, "result.json"), "w"), indent=2)
+    torch.save(model.state_dict(), os.path.join(a.out, "model.pt"))
+    print("FINAL", json.dumps(res["accuracy_by_passes"]), flush=True)
+    if "bpc" in res:
+        print(f"BPC {res['bpc']:.4f}", flush=True)
+
+
+if __name__ == "__main__":
+    main()