Automated push from Zaratan

results/g5-d4-distill-s0.json

@@ -0,0 +1,38 @@
+{
+  "args": {
+    "mode": "distill",
+    "K": 4,
+    "lam": 1.0,
+    "digits": 4,
+    "seq_len": 16,
+    "n_train": 200000,
+    "n_eval": 500,
+    "eval_bs": 100,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 128,
+    "steps": 25000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 0,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g5-d4-distill-s0"
+  },
+  "accuracy_by_passes": {
+    "1": 1.0,
+    "2": 1.0,
+    "3": 1.0,
+    "4": 1.0,
+    "6": 1.0,
+    "8": 1.0,
+    "12": 1.0,
+    "16": 1.0
+  },
+  "minutes": 7.094603908061981,
+  "compute": {
+    "params": 10649088,
+    "train_tokens": 51200000,
+    "fwd_per_step": 2
+  }
+}

results/g5-d4-distill-s1.json

@@ -0,0 +1,38 @@
+{
+  "args": {
+    "mode": "distill",
+    "K": 4,
+    "lam": 1.0,
+    "digits": 4,
+    "seq_len": 16,
+    "n_train": 200000,
+    "n_eval": 500,
+    "eval_bs": 100,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 128,
+    "steps": 25000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 1,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g5-d4-distill-s1"
+  },
+  "accuracy_by_passes": {
+    "1": 1.0,
+    "2": 1.0,
+    "3": 1.0,
+    "4": 1.0,
+    "6": 1.0,
+    "8": 1.0,
+    "12": 1.0,
+    "16": 1.0
+  },
+  "minutes": 7.182636868953705,
+  "compute": {
+    "params": 10649088,
+    "train_tokens": 51200000,
+    "fwd_per_step": 2
+  }
+}

results/g5-d4-distill-s2.json

@@ -0,0 +1,38 @@
+{
+  "args": {
+    "mode": "distill",
+    "K": 4,
+    "lam": 1.0,
+    "digits": 4,
+    "seq_len": 16,
+    "n_train": 200000,
+    "n_eval": 500,
+    "eval_bs": 100,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 128,
+    "steps": 25000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 2,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g5-d4-distill-s2"
+  },
+  "accuracy_by_passes": {
+    "1": 1.0,
+    "2": 1.0,
+    "3": 1.0,
+    "4": 1.0,
+    "6": 1.0,
+    "8": 1.0,
+    "12": 1.0,
+    "16": 1.0
+  },
+  "minutes": 7.177619067827861,
+  "compute": {
+    "params": 10649088,
+    "train_tokens": 51200000,
+    "fwd_per_step": 2
+  }
+}

results/g5-d4-full-s0.json

@@ -0,0 +1,38 @@
+{
+  "args": {
+    "mode": "full",
+    "K": 4,
+    "lam": 1.0,
+    "digits": 4,
+    "seq_len": 16,
+    "n_train": 200000,
+    "n_eval": 500,
+    "eval_bs": 100,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 128,
+    "steps": 25000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 0,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g5-d4-full-s0"
+  },
+  "accuracy_by_passes": {
+    "1": 1.0,
+    "2": 1.0,
+    "3": 1.0,
+    "4": 1.0,
+    "6": 1.0,
+    "8": 1.0,
+    "12": 1.0,
+    "16": 1.0
+  },
+  "minutes": 7.199072353045145,
+  "compute": {
+    "params": 10649088,
+    "train_tokens": 51200000,
+    "fwd_per_step": 2
+  }
+}

results/g5-d4-full-s1.json

@@ -0,0 +1,38 @@
+{
+  "args": {
+    "mode": "full",
+    "K": 4,
+    "lam": 1.0,
+    "digits": 4,
+    "seq_len": 16,
+    "n_train": 200000,
+    "n_eval": 500,
+    "eval_bs": 100,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 128,
+    "steps": 25000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 1,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g5-d4-full-s1"
+  },
+  "accuracy_by_passes": {
+    "1": 1.0,
+    "2": 1.0,
+    "3": 1.0,
+    "4": 1.0,
+    "6": 1.0,
+    "8": 1.0,
+    "12": 1.0,
+    "16": 1.0
+  },
+  "minutes": 7.2110337853431705,
+  "compute": {
+    "params": 10649088,
+    "train_tokens": 51200000,
+    "fwd_per_step": 2
+  }
+}

results/g5-d4-full-s2.json

@@ -0,0 +1,38 @@
+{
+  "args": {
+    "mode": "full",
+    "K": 4,
+    "lam": 1.0,
+    "digits": 4,
+    "seq_len": 16,
+    "n_train": 200000,
+    "n_eval": 500,
+    "eval_bs": 100,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 128,
+    "steps": 25000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 2,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g5-d4-full-s2"
+  },
+  "accuracy_by_passes": {
+    "1": 1.0,
+    "2": 1.0,
+    "3": 1.0,
+    "4": 1.0,
+    "6": 1.0,
+    "8": 1.0,
+    "12": 1.0,
+    "16": 1.0
+  },
+  "minutes": 7.120807111263275,
+  "compute": {
+    "params": 10649088,
+    "train_tokens": 51200000,
+    "fwd_per_step": 2
+  }
+}

results/g5-d4-matched-s0.json

@@ -0,0 +1,38 @@
+{
+  "args": {
+    "mode": "matched",
+    "K": 4,
+    "lam": 1.0,
+    "digits": 4,
+    "seq_len": 16,
+    "n_train": 200000,
+    "n_eval": 500,
+    "eval_bs": 100,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 128,
+    "steps": 25000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 0,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g5-d4-matched-s0"
+  },
+  "accuracy_by_passes": {
+    "1": 1.0,
+    "2": 1.0,
+    "3": 1.0,
+    "4": 1.0,
+    "6": 1.0,
+    "8": 1.0,
+    "12": 1.0,
+    "16": 1.0
+  },
+  "minutes": 4.066857294241587,
+  "compute": {
+    "params": 10649088,
+    "train_tokens": 51200000,
+    "fwd_per_step": 1
+  }
+}

results/g5-d4-matched-s1.json

@@ -0,0 +1,38 @@
+{
+  "args": {
+    "mode": "matched",
+    "K": 4,
+    "lam": 1.0,
+    "digits": 4,
+    "seq_len": 16,
+    "n_train": 200000,
+    "n_eval": 500,
+    "eval_bs": 100,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 128,
+    "steps": 25000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 1,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g5-d4-matched-s1"
+  },
+  "accuracy_by_passes": {
+    "1": 1.0,
+    "2": 1.0,
+    "3": 1.0,
+    "4": 1.0,
+    "6": 1.0,
+    "8": 1.0,
+    "12": 1.0,
+    "16": 1.0
+  },
+  "minutes": 4.046454687913259,
+  "compute": {
+    "params": 10649088,
+    "train_tokens": 51200000,
+    "fwd_per_step": 1
+  }
+}

results/g5-d4-matched-s2.json

@@ -0,0 +1,38 @@
+{
+  "args": {
+    "mode": "matched",
+    "K": 4,
+    "lam": 1.0,
+    "digits": 4,
+    "seq_len": 16,
+    "n_train": 200000,
+    "n_eval": 500,
+    "eval_bs": 100,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 128,
+    "steps": 25000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 2,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g5-d4-matched-s2"
+  },
+  "accuracy_by_passes": {
+    "1": 1.0,
+    "2": 1.0,
+    "3": 1.0,
+    "4": 1.0,
+    "6": 1.0,
+    "8": 1.0,
+    "12": 1.0,
+    "16": 1.0
+  },
+  "minutes": 4.117052431901296,
+  "compute": {
+    "params": 10649088,
+    "train_tokens": 51200000,
+    "fwd_per_step": 1
+  }
+}

results/g5-d4-mdlm-s0.json

@@ -0,0 +1,38 @@
+{
+  "args": {
+    "mode": "mdlm",
+    "K": 4,
+    "lam": 1.0,
+    "digits": 4,
+    "seq_len": 16,
+    "n_train": 200000,
+    "n_eval": 500,
+    "eval_bs": 100,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 128,
+    "steps": 25000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 0,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g5-d4-mdlm-s0"
+  },
+  "accuracy_by_passes": {
+    "1": 1.0,
+    "2": 1.0,
+    "3": 1.0,
+    "4": 1.0,
+    "6": 1.0,
+    "8": 1.0,
+    "12": 1.0,
+    "16": 1.0
+  },
+  "minutes": 4.046641910076142,
+  "compute": {
+    "params": 10649088,
+    "train_tokens": 51200000,
+    "fwd_per_step": 1
+  }
+}

results/g5-d4-mdlm-s1.json

@@ -0,0 +1,38 @@
+{
+  "args": {
+    "mode": "mdlm",
+    "K": 4,
+    "lam": 1.0,
+    "digits": 4,
+    "seq_len": 16,
+    "n_train": 200000,
+    "n_eval": 500,
+    "eval_bs": 100,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 128,
+    "steps": 25000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 1,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g5-d4-mdlm-s1"
+  },
+  "accuracy_by_passes": {
+    "1": 1.0,
+    "2": 1.0,
+    "3": 1.0,
+    "4": 1.0,
+    "6": 1.0,
+    "8": 1.0,
+    "12": 1.0,
+    "16": 1.0
+  },
+  "minutes": 4.090229777495066,
+  "compute": {
+    "params": 10649088,
+    "train_tokens": 51200000,
+    "fwd_per_step": 1
+  }
+}

results/g5-d4-mdlm-s2.json

@@ -0,0 +1,38 @@
+{
+  "args": {
+    "mode": "mdlm",
+    "K": 4,
+    "lam": 1.0,
+    "digits": 4,
+    "seq_len": 16,
+    "n_train": 200000,
+    "n_eval": 500,
+    "eval_bs": 100,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 128,
+    "steps": 25000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 2,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g5-d4-mdlm-s2"
+  },
+  "accuracy_by_passes": {
+    "1": 1.0,
+    "2": 1.0,
+    "3": 1.0,
+    "4": 1.0,
+    "6": 1.0,
+    "8": 1.0,
+    "12": 1.0,
+    "16": 1.0
+  },
+  "minutes": 4.043796098232269,
+  "compute": {
+    "params": 10649088,
+    "train_tokens": 51200000,
+    "fwd_per_step": 1
+  }
+}

results/g5-d6-distill-s0.json

@@ -0,0 +1,38 @@
+{
+  "args": {
+    "mode": "distill",
+    "K": 4,
+    "lam": 1.0,
+    "digits": 6,
+    "seq_len": 22,
+    "n_train": 200000,
+    "n_eval": 500,
+    "eval_bs": 100,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 128,
+    "steps": 25000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 0,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g5-d6-distill-s0"
+  },
+  "accuracy_by_passes": {
+    "1": 1.0,
+    "2": 1.0,
+    "3": 1.0,
+    "4": 1.0,
+    "6": 1.0,
+    "8": 1.0,
+    "12": 1.0,
+    "16": 1.0
+  },
+  "minutes": 7.2282985806465145,
+  "compute": {
+    "params": 10649088,
+    "train_tokens": 70400000,
+    "fwd_per_step": 2
+  }
+}

results/g5-d6-distill-s1.json

@@ -0,0 +1,38 @@
+{
+  "args": {
+    "mode": "distill",
+    "K": 4,
+    "lam": 1.0,
+    "digits": 6,
+    "seq_len": 22,
+    "n_train": 200000,
+    "n_eval": 500,
+    "eval_bs": 100,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 128,
+    "steps": 25000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 1,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g5-d6-distill-s1"
+  },
+  "accuracy_by_passes": {
+    "1": 1.0,
+    "2": 1.0,
+    "3": 1.0,
+    "4": 1.0,
+    "6": 1.0,
+    "8": 1.0,
+    "12": 1.0,
+    "16": 1.0
+  },
+  "minutes": 7.147757164637247,
+  "compute": {
+    "params": 10649088,
+    "train_tokens": 70400000,
+    "fwd_per_step": 2
+  }
+}

results/g5-d6-distill-s2.json

@@ -0,0 +1,38 @@
+{
+  "args": {
+    "mode": "distill",
+    "K": 4,
+    "lam": 1.0,
+    "digits": 6,
+    "seq_len": 22,
+    "n_train": 200000,
+    "n_eval": 500,
+    "eval_bs": 100,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 128,
+    "steps": 25000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 2,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g5-d6-distill-s2"
+  },
+  "accuracy_by_passes": {
+    "1": 1.0,
+    "2": 1.0,
+    "3": 1.0,
+    "4": 1.0,
+    "6": 1.0,
+    "8": 1.0,
+    "12": 1.0,
+    "16": 1.0
+  },
+  "minutes": 7.205686577161154,
+  "compute": {
+    "params": 10649088,
+    "train_tokens": 70400000,
+    "fwd_per_step": 2
+  }
+}

results/g5-d6-full-s0.json

@@ -0,0 +1,38 @@
+{
+  "args": {
+    "mode": "full",
+    "K": 4,
+    "lam": 1.0,
+    "digits": 6,
+    "seq_len": 22,
+    "n_train": 200000,
+    "n_eval": 500,
+    "eval_bs": 100,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 128,
+    "steps": 25000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 0,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g5-d6-full-s0"
+  },
+  "accuracy_by_passes": {
+    "1": 1.0,
+    "2": 1.0,
+    "3": 1.0,
+    "4": 1.0,
+    "6": 1.0,
+    "8": 1.0,
+    "12": 1.0,
+    "16": 1.0
+  },
+  "minutes": 7.124441786607107,
+  "compute": {
+    "params": 10649088,
+    "train_tokens": 70400000,
+    "fwd_per_step": 2
+  }
+}

results/g5-d6-full-s1.json

@@ -0,0 +1,38 @@
+{
+  "args": {
+    "mode": "full",
+    "K": 4,
+    "lam": 1.0,
+    "digits": 6,
+    "seq_len": 22,
+    "n_train": 200000,
+    "n_eval": 500,
+    "eval_bs": 100,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 128,
+    "steps": 25000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 1,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g5-d6-full-s1"
+  },
+  "accuracy_by_passes": {
+    "1": 1.0,
+    "2": 1.0,
+    "3": 1.0,
+    "4": 1.0,
+    "6": 1.0,
+    "8": 1.0,
+    "12": 1.0,
+    "16": 1.0
+  },
+  "minutes": 7.137926022211711,
+  "compute": {
+    "params": 10649088,
+    "train_tokens": 70400000,
+    "fwd_per_step": 2
+  }
+}

results/g5-d6-full-s2.json

@@ -0,0 +1,38 @@
+{
+  "args": {
+    "mode": "full",
+    "K": 4,
+    "lam": 1.0,
+    "digits": 6,
+    "seq_len": 22,
+    "n_train": 200000,
+    "n_eval": 500,
+    "eval_bs": 100,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 128,
+    "steps": 25000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 2,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g5-d6-full-s2"
+  },
+  "accuracy_by_passes": {
+    "1": 1.0,
+    "2": 1.0,
+    "3": 1.0,
+    "4": 1.0,
+    "6": 1.0,
+    "8": 1.0,
+    "12": 1.0,
+    "16": 1.0
+  },
+  "minutes": 7.225931032498678,
+  "compute": {
+    "params": 10649088,
+    "train_tokens": 70400000,
+    "fwd_per_step": 2
+  }
+}

results/g5-d6-matched-s0.json

@@ -0,0 +1,38 @@
+{
+  "args": {
+    "mode": "matched",
+    "K": 4,
+    "lam": 1.0,
+    "digits": 6,
+    "seq_len": 22,
+    "n_train": 200000,
+    "n_eval": 500,
+    "eval_bs": 100,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 128,
+    "steps": 25000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 0,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g5-d6-matched-s0"
+  },
+  "accuracy_by_passes": {
+    "1": 1.0,
+    "2": 1.0,
+    "3": 1.0,
+    "4": 1.0,
+    "6": 1.0,
+    "8": 1.0,
+    "12": 1.0,
+    "16": 1.0
+  },
+  "minutes": 4.056125191847483,
+  "compute": {
+    "params": 10649088,
+    "train_tokens": 70400000,
+    "fwd_per_step": 1
+  }
+}

results/g5-d6-matched-s1.json

@@ -0,0 +1,38 @@
+{
+  "args": {
+    "mode": "matched",
+    "K": 4,
+    "lam": 1.0,
+    "digits": 6,
+    "seq_len": 22,
+    "n_train": 200000,
+    "n_eval": 500,
+    "eval_bs": 100,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 128,
+    "steps": 25000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 1,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g5-d6-matched-s1"
+  },
+  "accuracy_by_passes": {
+    "1": 1.0,
+    "2": 1.0,
+    "3": 1.0,
+    "4": 1.0,
+    "6": 1.0,
+    "8": 1.0,
+    "12": 1.0,
+    "16": 1.0
+  },
+  "minutes": 4.098791329065959,
+  "compute": {
+    "params": 10649088,
+    "train_tokens": 70400000,
+    "fwd_per_step": 1
+  }
+}

results/g5-d6-matched-s2.json

@@ -0,0 +1,38 @@
+{
+  "args": {
+    "mode": "matched",
+    "K": 4,
+    "lam": 1.0,
+    "digits": 6,
+    "seq_len": 22,
+    "n_train": 200000,
+    "n_eval": 500,
+    "eval_bs": 100,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 128,
+    "steps": 25000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 2,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g5-d6-matched-s2"
+  },
+  "accuracy_by_passes": {
+    "1": 1.0,
+    "2": 1.0,
+    "3": 1.0,
+    "4": 1.0,
+    "6": 1.0,
+    "8": 1.0,
+    "12": 1.0,
+    "16": 1.0
+  },
+  "minutes": 4.063978902498881,
+  "compute": {
+    "params": 10649088,
+    "train_tokens": 70400000,
+    "fwd_per_step": 1
+  }
+}

results/g5-d6-mdlm-s0.json

@@ -0,0 +1,38 @@
+{
+  "args": {
+    "mode": "mdlm",
+    "K": 4,
+    "lam": 1.0,
+    "digits": 6,
+    "seq_len": 22,
+    "n_train": 200000,
+    "n_eval": 500,
+    "eval_bs": 100,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 128,
+    "steps": 25000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 0,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g5-d6-mdlm-s0"
+  },
+  "accuracy_by_passes": {
+    "1": 1.0,
+    "2": 1.0,
+    "3": 1.0,
+    "4": 1.0,
+    "6": 1.0,
+    "8": 1.0,
+    "12": 1.0,
+    "16": 1.0
+  },
+  "minutes": 4.098314166069031,
+  "compute": {
+    "params": 10649088,
+    "train_tokens": 70400000,
+    "fwd_per_step": 1
+  }
+}

results/g5-d6-mdlm-s1.json

@@ -0,0 +1,38 @@
+{
+  "args": {
+    "mode": "mdlm",
+    "K": 4,
+    "lam": 1.0,
+    "digits": 6,
+    "seq_len": 22,
+    "n_train": 200000,
+    "n_eval": 500,
+    "eval_bs": 100,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 128,
+    "steps": 25000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 1,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g5-d6-mdlm-s1"
+  },
+  "accuracy_by_passes": {
+    "1": 0.9960000000000001,
+    "2": 1.0,
+    "3": 1.0,
+    "4": 1.0,
+    "6": 1.0,
+    "8": 1.0,
+    "12": 1.0,
+    "16": 1.0
+  },
+  "minutes": 4.0784336686134335,
+  "compute": {
+    "params": 10649088,
+    "train_tokens": 70400000,
+    "fwd_per_step": 1
+  }
+}

results/g5-d6-mdlm-s2.json

@@ -0,0 +1,38 @@
+{
+  "args": {
+    "mode": "mdlm",
+    "K": 4,
+    "lam": 1.0,
+    "digits": 6,
+    "seq_len": 22,
+    "n_train": 200000,
+    "n_eval": 500,
+    "eval_bs": 100,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 128,
+    "steps": 25000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 2,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g5-d6-mdlm-s2"
+  },
+  "accuracy_by_passes": {
+    "1": 1.0,
+    "2": 1.0,
+    "3": 1.0,
+    "4": 1.0,
+    "6": 1.0,
+    "8": 1.0,
+    "12": 1.0,
+    "16": 1.0
+  },
+  "minutes": 4.086201139291128,
+  "compute": {
+    "params": 10649088,
+    "train_tokens": 70400000,
+    "fwd_per_step": 1
+  }
+}

results/g5-d8-distill-s0.json

@@ -0,0 +1,38 @@
+{
+  "args": {
+    "mode": "distill",
+    "K": 4,
+    "lam": 1.0,
+    "digits": 8,
+    "seq_len": 28,
+    "n_train": 200000,
+    "n_eval": 500,
+    "eval_bs": 100,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 128,
+    "steps": 25000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 0,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g5-d8-distill-s0"
+  },
+  "accuracy_by_passes": {
+    "1": 1.0,
+    "2": 0.9940000000000001,
+    "3": 0.9739999999999999,
+    "4": 0.968,
+    "6": 0.9620000000000001,
+    "8": 0.96,
+    "12": 0.96,
+    "16": 0.96
+  },
+  "minutes": 7.277184391021729,
+  "compute": {
+    "params": 10649088,
+    "train_tokens": 89600000,
+    "fwd_per_step": 2
+  }
+}

results/g5-d8-distill-s1.json

@@ -0,0 +1,38 @@
+{
+  "args": {
+    "mode": "distill",
+    "K": 4,
+    "lam": 1.0,
+    "digits": 8,
+    "seq_len": 28,
+    "n_train": 200000,
+    "n_eval": 500,
+    "eval_bs": 100,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 128,
+    "steps": 25000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 1,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g5-d8-distill-s1"
+  },
+  "accuracy_by_passes": {
+    "1": 0.998,
+    "2": 1.0,
+    "3": 1.0,
+    "4": 1.0,
+    "6": 1.0,
+    "8": 1.0,
+    "12": 1.0,
+    "16": 1.0
+  },
+  "minutes": 7.428394786516825,
+  "compute": {
+    "params": 10649088,
+    "train_tokens": 89600000,
+    "fwd_per_step": 2
+  }
+}

results/g5-d8-distill-s2.json

@@ -0,0 +1,38 @@
+{
+  "args": {
+    "mode": "distill",
+    "K": 4,
+    "lam": 1.0,
+    "digits": 8,
+    "seq_len": 28,
+    "n_train": 200000,
+    "n_eval": 500,
+    "eval_bs": 100,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 128,
+    "steps": 25000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 2,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g5-d8-distill-s2"
+  },
+  "accuracy_by_passes": {
+    "1": 1.0,
+    "2": 1.0,
+    "3": 1.0,
+    "4": 1.0,
+    "6": 1.0,
+    "8": 1.0,
+    "12": 1.0,
+    "16": 1.0
+  },
+  "minutes": 7.321874841054281,
+  "compute": {
+    "params": 10649088,
+    "train_tokens": 89600000,
+    "fwd_per_step": 2
+  }
+}

results/g5-d8-full-s0.json

@@ -0,0 +1,38 @@
+{
+  "args": {
+    "mode": "full",
+    "K": 4,
+    "lam": 1.0,
+    "digits": 8,
+    "seq_len": 28,
+    "n_train": 200000,
+    "n_eval": 500,
+    "eval_bs": 100,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 128,
+    "steps": 25000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 0,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g5-d8-full-s0"
+  },
+  "accuracy_by_passes": {
+    "1": 0.998,
+    "2": 0.994,
+    "3": 0.998,
+    "4": 0.9899999999999999,
+    "6": 0.9880000000000001,
+    "8": 0.982,
+    "12": 0.982,
+    "16": 0.982
+  },
+  "minutes": 7.317351579666138,
+  "compute": {
+    "params": 10649088,
+    "train_tokens": 89600000,
+    "fwd_per_step": 2
+  }
+}

results/g5-d8-full-s1.json

@@ -0,0 +1,38 @@
+{
+  "args": {
+    "mode": "full",
+    "K": 4,
+    "lam": 1.0,
+    "digits": 8,
+    "seq_len": 28,
+    "n_train": 200000,
+    "n_eval": 500,
+    "eval_bs": 100,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 128,
+    "steps": 25000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 1,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g5-d8-full-s1"
+  },
+  "accuracy_by_passes": {
+    "1": 1.0,
+    "2": 1.0,
+    "3": 1.0,
+    "4": 1.0,
+    "6": 1.0,
+    "8": 1.0,
+    "12": 1.0,
+    "16": 1.0
+  },
+  "minutes": 7.377502389748892,
+  "compute": {
+    "params": 10649088,
+    "train_tokens": 89600000,
+    "fwd_per_step": 2
+  }
+}

results/g5-d8-full-s2.json

@@ -0,0 +1,38 @@
+{
+  "args": {
+    "mode": "full",
+    "K": 4,
+    "lam": 1.0,
+    "digits": 8,
+    "seq_len": 28,
+    "n_train": 200000,
+    "n_eval": 500,
+    "eval_bs": 100,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 128,
+    "steps": 25000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 2,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g5-d8-full-s2"
+  },
+  "accuracy_by_passes": {
+    "1": 1.0,
+    "2": 1.0,
+    "3": 1.0,
+    "4": 1.0,
+    "6": 1.0,
+    "8": 1.0,
+    "12": 1.0,
+    "16": 1.0
+  },
+  "minutes": 7.370485206445058,
+  "compute": {
+    "params": 10649088,
+    "train_tokens": 89600000,
+    "fwd_per_step": 2
+  }
+}

results/g5-d8-matched-s0.json

@@ -0,0 +1,38 @@
+{
+  "args": {
+    "mode": "matched",
+    "K": 4,
+    "lam": 1.0,
+    "digits": 8,
+    "seq_len": 28,
+    "n_train": 200000,
+    "n_eval": 500,
+    "eval_bs": 100,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 128,
+    "steps": 25000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 0,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g5-d8-matched-s0"
+  },
+  "accuracy_by_passes": {
+    "1": 0.916,
+    "2": 0.858,
+    "3": 0.8460000000000001,
+    "4": 0.8400000000000001,
+    "6": 0.866,
+    "8": 0.858,
+    "12": 0.858,
+    "16": 0.858
+  },
+  "minutes": 4.13342049519221,
+  "compute": {
+    "params": 10649088,
+    "train_tokens": 89600000,
+    "fwd_per_step": 1
+  }
+}

results/g5-d8-matched-s1.json

@@ -0,0 +1,38 @@
+{
+  "args": {
+    "mode": "matched",
+    "K": 4,
+    "lam": 1.0,
+    "digits": 8,
+    "seq_len": 28,
+    "n_train": 200000,
+    "n_eval": 500,
+    "eval_bs": 100,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 128,
+    "steps": 25000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 1,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g5-d8-matched-s1"
+  },
+  "accuracy_by_passes": {
+    "1": 0.006,
+    "2": 0.006,
+    "3": 0.004,
+    "4": 0.004,
+    "6": 0.008,
+    "8": 0.008,
+    "12": 0.008,
+    "16": 0.008
+  },
+  "minutes": 4.116305422782898,
+  "compute": {
+    "params": 10649088,
+    "train_tokens": 89600000,
+    "fwd_per_step": 1
+  }
+}

results/g5-d8-matched-s2.json

@@ -0,0 +1,38 @@
+{
+  "args": {
+    "mode": "matched",
+    "K": 4,
+    "lam": 1.0,
+    "digits": 8,
+    "seq_len": 28,
+    "n_train": 200000,
+    "n_eval": 500,
+    "eval_bs": 100,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 128,
+    "steps": 25000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 2,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g5-d8-matched-s2"
+  },
+  "accuracy_by_passes": {
+    "1": 0.994,
+    "2": 0.9960000000000001,
+    "3": 0.9960000000000001,
+    "4": 0.998,
+    "6": 0.994,
+    "8": 0.994,
+    "12": 0.994,
+    "16": 0.994
+  },
+  "minutes": 4.115453656514486,
+  "compute": {
+    "params": 10649088,
+    "train_tokens": 89600000,
+    "fwd_per_step": 1
+  }
+}

results/g5-d8-mdlm-s0.json

@@ -0,0 +1,38 @@
+{
+  "args": {
+    "mode": "mdlm",
+    "K": 4,
+    "lam": 1.0,
+    "digits": 8,
+    "seq_len": 28,
+    "n_train": 200000,
+    "n_eval": 500,
+    "eval_bs": 100,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 128,
+    "steps": 25000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 0,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g5-d8-mdlm-s0"
+  },
+  "accuracy_by_passes": {
+    "1": 0.064,
+    "2": 0.268,
+    "3": 0.20400000000000001,
+    "4": 0.202,
+    "6": 0.10800000000000001,
+    "8": 0.09999999999999999,
+    "12": 0.04,
+    "16": 0.04
+  },
+  "minutes": 4.100277105967204,
+  "compute": {
+    "params": 10649088,
+    "train_tokens": 89600000,
+    "fwd_per_step": 1
+  }
+}

results/g5-d8-mdlm-s1.json

@@ -0,0 +1,38 @@
+{
+  "args": {
+    "mode": "mdlm",
+    "K": 4,
+    "lam": 1.0,
+    "digits": 8,
+    "seq_len": 28,
+    "n_train": 200000,
+    "n_eval": 500,
+    "eval_bs": 100,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 128,
+    "steps": 25000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 1,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g5-d8-mdlm-s1"
+  },
+  "accuracy_by_passes": {
+    "1": 0.0,
+    "2": 0.0,
+    "3": 0.0,
+    "4": 0.0,
+    "6": 0.0,
+    "8": 0.0,
+    "12": 0.0,
+    "16": 0.0
+  },
+  "minutes": 4.152010726928711,
+  "compute": {
+    "params": 10649088,
+    "train_tokens": 89600000,
+    "fwd_per_step": 1
+  }
+}

results/g5-d8-mdlm-s2.json

@@ -0,0 +1,38 @@
+{
+  "args": {
+    "mode": "mdlm",
+    "K": 4,
+    "lam": 1.0,
+    "digits": 8,
+    "seq_len": 28,
+    "n_train": 200000,
+    "n_eval": 500,
+    "eval_bs": 100,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 128,
+    "steps": 25000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 2,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g5-d8-mdlm-s2"
+  },
+  "accuracy_by_passes": {
+    "1": 0.992,
+    "2": 0.992,
+    "3": 0.998,
+    "4": 0.998,
+    "6": 1.0,
+    "8": 1.0,
+    "12": 1.0,
+    "16": 1.0
+  },
+  "minutes": 4.113563962777456,
+  "compute": {
+    "params": 10649088,
+    "train_tokens": 89600000,
+    "fwd_per_step": 1
+  }
+}

slurm/g5.sh

@@ -0,0 +1,41 @@
+#!/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

slurm/smk.sh

@@ -0,0 +1,32 @@
+#!/bin/bash
+#SBATCH -J smk
+#SBATCH -p gpu-h100
+#SBATCH --gpus=h100:1
+#SBATCH -c 8
+#SBATCH --mem=32g
+#SBATCH -t 00:30:00
+#SBATCH -o logs/smk-%j.out
+source /etc/profile
+SCR=/scratch/zt1/project/msml612/user/$USER
+module load python/3.10.10
+export PYTHONPATH=
+cd $SCR/pcmd/src
+$SCR/hf_env/bin/python - <<'PY'
+import sys; sys.path.insert(0,'.')
+import numpy as np
+from data import build_addition
+X, A, tok = build_addition(300, 6, 32, seed=3, exact=True)
+bad = 0
+for i in range(len(X)):
+    s = [tok.itos[t] for t in X[i] if t != tok.pad]
+    a = int("".join(s[:s.index("+")])); b = int("".join(s[s.index("+")+1:s.index("=")]))
+    ans = int("".join(reversed(s[s.index("=")+1:])))
+    bad += (a + b != ans)
+print(f"  data check: {bad}/300 wrong; answer positions/example = {A.sum(1)[:3]}")
+PY
+for M in mdlm matched distill full; do
+  echo "--- $M ---"
+  $SCR/hf_env/bin/python train_add.py --mode $M --digits 3 --seq_len 16 \
+    --n_train 4000 --n_eval 60 --eval_bs 30 --d 64 --layers 2 --heads 2 \
+    --bs 32 --steps 200 --K 4 --out /tmp/a$M 2>&1 | tail -2
+done

src/data.py

@@ -130,3 +130,60 @@ def build_probe(n, seq_len, coupling, seed=0, task="copy"):
         else:
             raise ValueError(task)
     return X, tok
+
+
+# ---------------------------------------------------------------------------
+# Addition: the canonical UNIQUE-answer task with chained inter-token
+# dependence. Given the operands the answer is fully determined, so every
+# position's true marginal is a point mass and parallel decoding faces no
+# information-theoretic barrier - unlike `agree`/`parity`, where the value is
+# genuinely free and the bound is unbeatable. What remains is a computation
+# problem: digit i needs the carry out of digit i-1, and sequential decoding
+# amortises that chain across passes while parallel decoding must fit it inside
+# one fixed-depth forward pass.
+# ---------------------------------------------------------------------------
+
+DIGITS10 = "0123456789"
+
+
+class AddTokenizer:
+    def __init__(self):
+        self.itos = ["[MASK]", "[PAD]", "+", "="] + list(DIGITS10)
+        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 build_addition(n, digits, seq_len, seed=0, exact=False):
+    """Returns (X, answer_mask, tok).
+
+    Layout: operand a, '+', operand b, '=', then the answer written
+    units-first. `answer_mask` marks the positions the model must produce; the
+    prompt is always visible, so the task is conditional and the answer unique.
+    """
+    tok = AddTokenizer()
+    rng = np.random.default_rng(seed)
+    X = np.full((n, seq_len), tok.pad, dtype=np.int64)
+    A = np.zeros((n, seq_len), dtype=bool)
+
+    for i in range(n):
+        d = digits if exact else int(rng.integers(1, digits + 1))
+
+        def draw(nd):
+            ds = rng.integers(0, 10, size=nd)
+            if nd > 1 and ds[0] == 0:
+                ds[0] = rng.integers(1, 10)
+            return "".join(str(int(v)) for v in ds)
+
+        pa, pb = draw(d), draw(d)
+        s = str(int(pa) + int(pb))[::-1]          # units first
+        seq = list(pa) + ["+"] + list(pb) + ["="] + list(s)
+        if len(seq) > seq_len:
+            raise ValueError(f"seq_len {seq_len} too small for {d} digits")
+        ids = [tok.stoi[c] for c in seq]
+        X[i, : len(ids)] = ids
+        A[i, len(pa) + len(pb) + 2: len(ids)] = True   # answer positions only
+    return X, A, tok

src/diffusion.py

@@ -317,3 +317,118 @@ def dependence_weighted_entropy(model, x0, tok, t_min=None):
     p = model(xt).softmax(-1)
     h2 = -(p.pow(2).sum(-1).clamp_min(1e-9)).log()           # Renyi-2 entropy
     return (d * h2 * m).sum() / m.sum().clamp_min(1)
+
+
+# ===========================================================================
+# Conditional objectives for UNIQUE-answer tasks (prompt visible, answer
+# masked). Here the target is deterministic, so unlike the free-choice probes
+# there is a correct answer to aim at and coordination is learnable.
+# ===========================================================================
+
+def _cond_mask(x0, amask, tok, t):
+    """Mask a fraction t of the ANSWER positions only; the prompt stays visible."""
+    dev = x0.device
+    r = torch.rand_like(amask, dtype=torch.float)
+    m = (r < t[:, None]) & amask
+    # never leave an example with nothing to predict
+    empty = ~m.any(1)
+    if empty.any():
+        first = amask.float().argmax(1)
+        m[empty, first[empty]] = True
+    return m
+
+
+def cond_mdlm_loss(model, x0, amask, tok, t_dist="uniform", K=None):
+    """Baseline, plus the schedule-matched variant.
+
+    t_dist="uniform" is standard MDLM: every masking ratio is trained equally.
+    t_dist="matched" instead draws t from the ratios a K-pass decode actually
+    visits, namely {1, (K-1)/K, ..., 1/K}. Standard training spends most of its
+    capacity on nearly-complete states that high-parallelism decoding never
+    sees, which is a plain train/inference mismatch.
+    """
+    B, L = x0.shape
+    dev = x0.device
+    if t_dist == "matched" and K:
+        idx = torch.randint(1, K + 1, (B,), device=dev).float()
+        t = idx / K
+    else:
+        t = torch.rand(B, device=dev).clamp_min(1.0 / L)
+
+    m = _cond_mask(x0, amask, tok, t)
+    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 ((ce * m).sum(1) / m.sum(1).clamp_min(1)).mean()
+
+
+def sequential_distill_loss(model, x0, amask, tok, chunk=8):
+    """THE METHOD: distil the sequential computation into a single pass.
+
+    Sequential decoding is strong because each step can read the digits already
+    written, letting the carry chain unroll across passes. A one-pass parallel
+    prediction has to fit that whole chain inside a fixed depth, which is why
+    accuracy collapses as passes are reduced.
+
+    So we train the ONE-PASS prediction, from the fully masked answer, to match
+    the answer the model produces when it is allowed to decode sequentially.
+    Because the answer is unique, the teacher's target is the same every time it
+    sees the same prompt - which is exactly what was missing on the free-choice
+    probes, where the teacher's arbitrary pick averaged back to uniform and the
+    signal cancelled.
+
+    Ground truth is available here, so the teacher is only a convenience: the
+    loss below supervises the fully-masked one-pass prediction directly against
+    the true answer, which is the strongest possible version of the same idea.
+    """
+    B, L = x0.shape
+    xt = torch.where(amask, torch.full_like(x0, tok.mask), x0)   # whole answer hidden
+    logits = model(xt)
+    ce = F.cross_entropy(logits.reshape(-1, logits.size(-1)), x0.reshape(-1),
+                         reduction="none").view(B, L)
+    return ((ce * amask).sum(1) / amask.sum(1).clamp_min(1)).mean()
+
+
+def pmd_loss(model, x0, amask, tok, mode="mdlm", K=8, lam=1.0):
+    """Parallelism-Matched Diffusion.
+
+      mdlm     standard uniform-t training (baseline)
+      matched  train only at the masking ratios a K-pass decode visits
+      distill  uniform training + one-pass supervision (the method)
+      full     matched schedule + one-pass supervision
+    """
+    if mode == "mdlm":
+        return cond_mdlm_loss(model, x0, amask, tok, "uniform")
+    if mode == "matched":
+        return cond_mdlm_loss(model, x0, amask, tok, "matched", K)
+    base = cond_mdlm_loss(model, x0, amask, tok,
+                          "matched" if mode == "full" else "uniform", K)
+    return base + lam * sequential_distill_loss(model, x0, amask, tok)
+
+
+@torch.no_grad()
+def cond_decode(model, x0, amask, tok, steps, device):
+    """Reveal the answer over `steps` passes, prompt fixed throughout."""
+    x = torch.where(amask, torch.full_like(x0, tok.mask), x0).to(device)
+    am = amask.to(device)
+    B, L = x.shape
+    n_ans = int(am[0].sum())
+    pred = None
+    for s in range(steps, 0, -1):
+        masked = (x == tok.mask) & am
+        if not masked.any():
+            break
+        logits = model(x)
+        conf, pred = logits.softmax(-1).max(-1)
+        conf = conf.masked_fill(~masked, -1e9)
+        left = int(n_ans * (s - 1) / steps)
+        for b in range(B):
+            k = max(0, int(masked[b].sum()) - left)
+            if k:
+                idx = conf[b].topk(k).indices
+                x[b, idx] = pred[b, idx]
+    rem = (x == tok.mask) & am
+    if rem.any() and pred is not None:
+        x[rem] = pred[rem]
+    return x

src/train_add.py

@@ -0,0 +1,120 @@
+"""Parallelism-Matched Diffusion on addition: a unique-answer, chained task.
+
+Addition is the right testbed because the answer is fully determined by the
+prompt. Every position's true marginal is a point mass, so there is no
+information-theoretic obstacle to parallel decoding - unlike free-choice tasks,
+where two independent draws simply cannot agree and no training can help. What
+is left is a computation problem, and computation problems are trainable.
+
+Headline axis: exact-match accuracy against the number of denoising passes,
+from fully parallel (1) to fully sequential.
+"""
+import argparse, json, math, os, time
+import numpy as np
+import torch
+
+from data import build_addition
+from model import Transformer
+import diffusion as dfn
+
+
+def get_args():
+    p = argparse.ArgumentParser()
+    p.add_argument("--mode", choices=["mdlm", "matched", "distill", "full"], default="mdlm")
+    p.add_argument("--K", type=int, default=8, help="parallelism the schedule is matched to")
+    p.add_argument("--lam", type=float, default=1.0, help="one-pass supervision weight")
+    p.add_argument("--digits", type=int, default=6)
+    p.add_argument("--seq_len", type=int, default=32)
+    p.add_argument("--n_train", type=int, default=200000)
+    p.add_argument("--n_eval", type=int, default=500)
+    p.add_argument("--eval_bs", type=int, default=100)
+    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=128)
+    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("--out", default="runs/dev")
+    return p.parse_args()
+
+
+PASS_GRID = [1, 2, 3, 4, 6, 8, 12, 16]
+
+
+def exact_match(pred, gold, amask, tok):
+    ok = 0
+    for p, g, a in zip(pred.tolist(), gold.tolist(), amask.tolist()):
+        pa = [p[i] for i in range(len(a)) if a[i]]
+        ga = [g[i] for i in range(len(a)) if a[i]]
+        ok += pa == ga
+    return ok / len(gold)
+
+
+@torch.no_grad()
+def evaluate(model, a, tok, device):
+    model.eval()
+    out = {}
+    X, A, _ = build_addition(a.n_eval, a.digits, a.seq_len, seed=90_000, exact=True)
+    X, A = torch.from_numpy(X), torch.from_numpy(A)
+    for steps in PASS_GRID:
+        accs = []
+        for i in range(0, len(X), a.eval_bs):
+            xb, ab = X[i:i + a.eval_bs].to(device), A[i:i + a.eval_bs].to(device)
+            pred = dfn.cond_decode(model, xb, ab, tok, steps, device)
+            accs.append(exact_match(pred.cpu(), xb.cpu(), ab.cpu(), tok))
+        out[str(steps)] = float(np.mean(accs))
+    model.train()
+    return out
+
+
+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)
+
+    X, A, tok = build_addition(a.n_train, a.digits, a.seq_len, seed=a.seed)
+    X, A = torch.from_numpy(X), torch.from_numpy(A)
+
+    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.mode}] params={model.n_params()/1e6:.2f}M digits<={a.digits} "
+          f"K={a.K} 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(X), (a.bs,))
+        xb, ab = X[idx].to(dev), A[idx].to(dev)
+        opt.zero_grad(set_to_none=True)
+        with torch.autocast(**amp):
+            loss = dfn.pmd_loss(model, xb, ab, tok, a.mode, a.K, a.lam)
+        loss.backward()
+        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+        opt.step(); sched.step()
+        if step % 2000 == 0:
+            print(f"step {step:6d} loss {loss.item():.4f} ({time.time()-t0:.0f}s)", flush=True)
+
+    acc = evaluate(model, a, tok, dev)
+    n_par = model.n_params()
+    json.dump({"args": vars(a), "accuracy_by_passes": acc,
+               "minutes": (time.time() - t0) / 60,
+               "compute": {"params": n_par,
+                           "train_tokens": a.steps * a.bs * a.seq_len,
+                           "fwd_per_step": 2 if a.mode in ("distill", "full") else 1}},
+              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(acc), flush=True)
+
+
+if __name__ == "__main__":
+    main()