File size: 6,543 Bytes
2220ce4 55003bd c62918e 2220ce4 55003bd 2220ce4 55003bd 2220ce4 55003bd 2220ce4 55003bd 2220ce4 55003bd 2220ce4 c62918e 55003bd c62918e 2220ce4 c62918e 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 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 | """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", "prog",
"anybudget", "anybudget_nc"],
default="mdlm",
help="anybudget = sample K and condition on it (the method); "
"anybudget_nc = same K sampling WITHOUT conditioning, "
"the ablation that isolates what conditioning buys")
p.add_argument("--budget_bins", type=int, default=10,
help="log2 bins for budget conditioning; 0 disables")
p.add_argument("--K_max", type=int, default=16)
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):
"""Accuracy across the whole pass grid.
The grid now runs out to a FULLY SEQUENTIAL decode (one token per pass), not
just 16: a 20-digit answer has 21 slots, so stopping at 16 never measures
sequential decoding at all and cannot tell "failed to parallelise" apart
from "failed to learn".
"""
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)
cond = getattr(model, "budget_bins", 0) > 0
grid = sorted(set(PASS_GRID) | {int(A[0].sum())})
for steps in 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_budget(model, xb, ab, tok, steps, device)
if cond else 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)
cond = a.mode == "anybudget"
model = Transformer(len(tok), a.d, a.layers, a.heads, a.pe,
causal=False, max_len=a.seq_len + 8,
budget_bins=a.budget_bins if cond else 0).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):
if a.mode == "prog":
loss, K_now, lam_now = dfn.pmd_progressive(
model, xb, ab, tok, step / max(1, a.steps), a.K_max, a.lam)
elif a.mode in ("anybudget", "anybudget_nc"):
loss = dfn.cond_any_budget_loss(model, xb, ab, tok,
condition=cond)
else:
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:
extra = f" K={K_now} lam={lam_now:.2f}" if a.mode == "prog" else ""
print(f"step {step:6d} loss {loss.item():.4f}{extra} "
f"({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()
|