File size: 9,797 Bytes
98c193b 7fa1b76 5135496 7fa1b76 98c193b 7fa1b76 98c193b 7fa1b76 5135496 7fa1b76 5135496 7fa1b76 98c193b | 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 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 | """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", "coord", "cmd"],
default="mdlm",
help="mdlm=ELBO baseline; selfcorrupt=+self-conditioning; "
"coord=+coordination; cmd=full method (all three)")
p.add_argument("--alpha", type=float, default=0.5, help="self-conditioning weight")
p.add_argument("--beta", type=float, default=1.0, help="coordination weight")
p.add_argument("--gamma", type=float, default=0.0,
help="marginal-entropy penalty: the explicit price paid in "
"likelihood to buy parallel decodability")
p.add_argument("--dep_weighted", type=int, default=1,
help="1 = sharpen only where positions are measurably dependent; "
"0 = sharpen everywhere (ablation)")
p.add_argument("--elbo_w", type=float, default=1.0,
help="down-weight the ELBO so coordination is not cancelled by it")
p.add_argument("--group", type=int, default=4)
p.add_argument("--coord_every", type=int, default=4,
help="apply the coordination term every N steps; the teacher\n costs `group` sequential passes, so this amortises it")
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("--ramp_frac", type=float, default=0.3,
help="fraction of training over which self-corruption ramps "
"from 0 to --k_frac; corrupting from step 0 trains on a "
"random model's noise and prevents learning entirely")
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):
ramp = min(1.0, step / max(1, a.ramp_frac * a.steps))
uc = (step % a.coord_every == 0)
if a.objective == "cmd":
loss = dfn.cmd_loss(model, x0, tok, a.alpha, a.beta, a.k_frac,
ramp, a.group, use_coord=uc,
gamma=a.gamma, elbo_w=a.elbo_w,
dep_weighted=bool(a.dep_weighted))
elif a.objective == "coord":
loss = dfn.cmd_loss(model, x0, tok, 0.0, a.beta, a.k_frac,
ramp, a.group, use_coord=uc,
gamma=a.gamma, elbo_w=a.elbo_w,
dep_weighted=bool(a.dep_weighted))
elif a.objective == "selfcorrupt":
loss = dfn.cmd_loss(model, x0, tok, a.alpha, 0.0, a.k_frac, ramp)
else:
loss = 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()
|