| """DiffuRefill-1B as an lm-eval-harness model, so it is scored on exactly the |
| tasks, prompts and metrics the autoregressive baselines are. |
| |
| A masked-diffusion model has no left-to-right likelihood of its own; it scores a |
| continuation by the chain rule in left-to-right order. For a continuation |
| x_1..x_n after context c, row i is |
| |
| <bos> c x_1 .. x_{i-1} [MASK] * (n - i + 1) |
| |
| and log p(x_i | c, x_<i) is read at the first masked position. The model knows |
| the continuation's length (the masks are there) but sees none of its later |
| tokens, so the sum over i is a proper likelihood of x under the decoding order |
| the model is later sampled in. The same chain gives `is_greedy`: every x_i is |
| the argmax at its own step. |
| |
| Rows are grouped by exact length, so no padding is ever attended to (attention |
| is bidirectional), and only the scored position goes through the 73k-way head. |
| """ |
| import os, sys, time, collections |
| import torch |
| import torch.nn.functional as F |
|
|
| os.environ.setdefault("BG_DOCMASK", "0") |
| sys.path.insert(0, "/tmp") |
| from big_common import Net, MASK_ID, VOCAB, SEQ |
| from lm_eval.api.model import LM |
|
|
| REPO = "Asilarkness/DiffuRefill-1B" |
|
|
|
|
| class DLM(LM): |
| def __init__(self, ckpt="final/DiffuRefill-1B-200k-swa.pt", budget=65536, device="cuda"): |
| super().__init__() |
| from huggingface_hub import hf_hub_download |
| from transformers import AutoTokenizer |
| self.tk = AutoTokenizer.from_pretrained("openbmb/MiniCPM4-0.5B", trust_remote_code=True) |
| self.bos = self.tk.bos_token_id if self.tk.bos_token_id is not None else 1 |
| p = ckpt if ckpt.startswith("/") else hf_hub_download(REPO, ckpt, local_dir="/tmp/dlm") |
| g = torch.load(p, map_location="cpu", weights_only=False) |
| net = Net() |
| net.load_state_dict({k: v.float() for k, v in g["model"].items()}) |
| self.net = net.to(device, torch.bfloat16).eval() |
| self.dev, self.budget = device, budget |
|
|
| def _enc(self, s): |
| return [i for i in self.tk(s, add_special_tokens=False).input_ids if i < VOCAB - 2] |
|
|
| def _pair(self, ctx, cont): |
| |
| n = len(ctx) - len(ctx.rstrip()) |
| if n: |
| cont, ctx = ctx[-n:] + cont, ctx[:-n] |
| whole = self._enc(ctx + cont) |
| c = self._enc(ctx) |
| x = whole[len(c):] |
| if not x or whole[:len(c)] != c: |
| x = self._enc(cont) |
| c = [self.bos] + c |
| c = c[-(SEQ - len(x)):] |
| return c, x |
|
|
| @torch.no_grad() |
| def _score_rows(self, rows): |
| """rows: list of (ids, pos, target); all ids the same length. -> logp, argmax""" |
| ids = torch.tensor([r[0] for r in rows], device=self.dev) |
| pos = torch.tensor([r[1] for r in rows], device=self.dev) |
| tgt = torch.tensor([r[2] for r in rows], device=self.dev) |
| h = self.net.body(ids)[torch.arange(len(rows), device=self.dev), pos] |
| lp = F.log_softmax(self.net.head(h).float()[:, :VOCAB - 2], -1) |
| return lp.gather(1, tgt[:, None])[:, 0].tolist(), (lp.argmax(-1) == tgt).tolist() |
|
|
| def loglikelihood(self, requests, disable_tqdm=False): |
| if os.environ.get("EV_SCORER", "chain") == "elbo": |
| return self._elbo(requests) |
| jobs = collections.defaultdict(list) |
| for k, r in enumerate(requests): |
| c, x = self._pair(*r.args) |
| n = len(x) |
| for i in range(n): |
| ids = c + x[:i] + [MASK_ID] * (n - i) |
| jobs[len(ids)].append((k, ids, len(c) + i, x[i])) |
| lp = [0.0] * len(requests) |
| greedy = [True] * len(requests) |
| t0, done, total = time.time(), 0, sum(len(v) for v in jobs.values()) |
| for L in sorted(jobs): |
| rows = jobs[L] |
| B = max(1, self.budget // L) |
| for s in range(0, len(rows), B): |
| part = rows[s:s + B] |
| with torch.autocast("cuda", torch.bfloat16): |
| l, g = self._score_rows([(r[1], r[2], r[3]) for r in part]) |
| for r, a, b in zip(part, l, g): |
| lp[r[0]] += a |
| greedy[r[0]] &= b |
| done += len(part) |
| print(" ll rows %d/%d (%.0fs)" % (done, total, time.time() - t0), flush=True) |
| return list(zip(lp, greedy)) |
|
|
| @torch.no_grad() |
| def _elbo(self, requests, N=int(os.environ.get("EV_ELBO_N", "32"))): |
| """LLaDA's conditional estimator: mask l of the n continuation tokens |
| (l stratified over 1..n, positions random), score (n/l) * sum of the |
| masked log-probs, average over N draws. A lower bound on log p(x|c); |
| used only as a cross-check of the chain-rule numbers.""" |
| import random |
| rng = random.Random(0) |
| lp = [0.0] * len(requests) |
| jobs = collections.defaultdict(list) |
| for k, r in enumerate(requests): |
| c, x = self._pair(*r.args) |
| n = len(x) |
| for j in range(N): |
| l = 1 + min(n - 1, int((j + rng.random()) / N * n)) |
| pos = rng.sample(range(n), l) |
| ids = c + [MASK_ID if i in pos else t for i, t in enumerate(x)] |
| jobs[len(ids)].append((k, ids, [len(c) + i for i in pos], [x[i] for i in pos], n / l / N)) |
| for L in sorted(jobs): |
| rows = jobs[L] |
| B = max(1, self.budget // L) |
| for s in range(0, len(rows), B): |
| part = rows[s:s + B] |
| ids = torch.tensor([r[1] for r in part], device=self.dev) |
| with torch.autocast("cuda", torch.bfloat16): |
| h = self.net.body(ids) |
| for b, r in enumerate(part): |
| q = F.log_softmax(self.net.head(h[b, r[2]]).float()[:, :VOCAB - 2], -1) |
| lp[r[0]] += r[4] * float(q.gather(1, torch.tensor(r[3], device=self.dev)[:, None]).sum()) |
| return [(v, False) for v in lp] |
|
|
| def loglikelihood_rolling(self, requests, disable_tqdm=False): |
| raise NotImplementedError("perplexity tasks are not used") |
|
|
| def generate_until(self, requests, disable_tqdm=False): |
| raise NotImplementedError("generation tasks are not used") |
|
|
|
|
| if __name__ == "__main__": |
| import json, lm_eval |
| tasks = os.environ.get("EV_TASKS", "hellaswag,arc_easy,arc_challenge,piqa,winogrande,openbookqa,sciq,lambada_openai,mmlu").split(",") |
| out = os.environ.get("EV_OUT", "/tmp/bench_dlm.json") |
| m = DLM(os.environ.get("EV_CKPT", "final/DiffuRefill-1B-200k-swa.pt")) |
| res = {} |
| for t in tasks: |
| try: |
| lim = int(os.environ["EV_LIMIT"]) if os.environ.get("EV_LIMIT") else None |
| r = lm_eval.simple_evaluate(model=m, tasks=[t], num_fewshot=0, bootstrap_iters=0, limit=lim) |
| res.update({k: v for k, v in r["results"].items()}) |
| print("RESULT", t, json.dumps(r["results"].get(t, {})), flush=True) |
| except Exception as e: |
| print("TASKFAIL", t, str(e)[:200], flush=True) |
| json.dump(res, open(out, "w")) |
| print("BENCH DONE", flush=True) |
|
|
