| """Autoregressive evaluator for generative perplexity on text8. |
| |
| Generative perplexity - the likelihood an independent language model assigns to |
| our samples - is the standard quality metric for diffusion language models, |
| because the diffusion model's own ELBO says nothing about the quality of what it |
| actually generates at a low number of function evaluations. |
| |
| The evaluator is trained here from scratch on the same corpus rather than using |
| a pretrained GPT-2, for two reasons. The project's constraint is that everything |
| is trained from scratch, and a character-level evaluator trained on text8 itself |
| is the correct judge for text8 samples - a word-piece model trained on web text |
| would score character-level artifacts through a tokenizer that never sees them. |
| |
| The evaluator is trained ONCE and shared by every diffusion run, so all systems |
| are judged by exactly the same yardstick. |
| """ |
| import argparse, json, math, os, time |
| import numpy as np |
| import torch |
| import torch.nn.functional as F |
|
|
| from model import Transformer |
| from text8 import load_text8, batches |
|
|
| p = argparse.ArgumentParser() |
| p.add_argument("--root", default="data/text8") |
| p.add_argument("--seq_len", type=int, default=256) |
| p.add_argument("--d", type=int, default=512) |
| p.add_argument("--layers", type=int, default=8) |
| p.add_argument("--heads", type=int, default=8) |
| p.add_argument("--bs", type=int, default=128) |
| p.add_argument("--steps", type=int, default=60000) |
| p.add_argument("--lr", type=float, default=3e-4) |
| p.add_argument("--warmup", type=int, default=2000) |
| p.add_argument("--seed", type=int, default=0) |
| p.add_argument("--out", default="runs/evallm") |
| a = p.parse_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) |
|
|
| tr, va, te, tok = load_text8(a.root, a.seq_len) |
| print(f"text8: train {tr.shape}, valid {va.shape}, test {te.shape}, vocab {len(tok)}", |
| flush=True) |
|
|
| model = Transformer(len(tok), a.d, a.layers, a.heads, "rope", |
| causal=True, max_len=a.seq_len + 8).to(dev) |
| print(f"evaluator params: {sum(p.numel() for p in model.parameters())/1e6:.1f}M", flush=True) |
| opt = torch.optim.AdamW(model.parameters(), lr=a.lr, weight_decay=0.01, betas=(0.9, 0.95)) |
| sched = torch.optim.lr_scheduler.LambdaLR( |
| opt, lambda s: min((s + 1) / a.warmup, max(0.0, (a.steps - s) / max(1, a.steps - a.warmup)))) |
|
|
| it = batches(tr, a.bs, seed=a.seed) |
| t0 = time.time() |
| for step in range(a.steps): |
| x = torch.from_numpy(next(it)).to(dev) |
| with torch.autocast("cuda", dtype=torch.bfloat16, enabled=(dev == "cuda")): |
| logits = model(x) |
| |
| loss = F.cross_entropy(logits[:, :-1].reshape(-1, len(tok)), |
| x[:, 1:].reshape(-1)) |
| opt.zero_grad(set_to_none=True) |
| 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} " |
| f"bpc {loss.item()/math.log(2):.3f} {time.time()-t0:.0f}s", flush=True) |
|
|
|
|
| @torch.no_grad() |
| def test_bpc(arr, n=200): |
| model.eval() |
| tot, cnt = 0.0, 0 |
| for i in range(0, min(n * a.bs, len(arr)), a.bs): |
| x = torch.from_numpy(arr[i:i + a.bs]).to(dev) |
| if len(x) < 2: |
| continue |
| logits = model(x) |
| ce = F.cross_entropy(logits[:, :-1].reshape(-1, len(tok)), |
| x[:, 1:].reshape(-1), reduction="sum") |
| tot += float(ce); cnt += x[:, 1:].numel() |
| model.train() |
| return tot / cnt / math.log(2) |
|
|
|
|
| bpc = test_bpc(te) |
| print(f"\nevaluator test BPC: {bpc:.4f}", flush=True) |
| torch.save(model.state_dict(), os.path.join(a.out, "model.pt")) |
| json.dump({"args": vars(a), "test_bpc": bpc, "vocab": len(tok)}, |
| open(os.path.join(a.out, "result.json"), "w"), indent=2) |
| print(f"saved to {a.out}", flush=True) |
|
|
