| """Evaluation for Part 1: target perplexity, BLEU, and expert-usage maps.""" |
|
|
| from __future__ import annotations |
|
|
| import numpy as np |
| import torch |
|
|
| from .data import EOS, LANG_NAMES, build_prompt, eval_pairs |
| from .ffn import MoE |
|
|
|
|
| |
| |
| |
| @torch.no_grad() |
| def evaluate_loss(model, loader, device, max_batches=None, amp_dtype=torch.bfloat16): |
| """Cross-entropy per *target* token, summed properly across batches. |
| |
| Averaging per-batch means would weight a block with three target tokens |
| the same as one with two hundred, so the totals are accumulated instead. |
| """ |
| model.eval() |
| tot_loss, tot_tok = 0.0, 0 |
| for i, (x, y, m, _) in enumerate(loader): |
| if max_batches is not None and i >= max_batches: |
| break |
| x, y, m = x.to(device), y.to(device), m.to(device) |
| with torch.autocast("cuda", dtype=amp_dtype, enabled=device.type == "cuda"): |
| logits, _, _ = model(x) |
| ce = torch.nn.functional.cross_entropy( |
| logits.view(-1, logits.size(-1)).float(), y.reshape(-1), |
| reduction="none") |
| mask = m.reshape(-1) |
| tot_loss += (ce * mask).sum().item() |
| tot_tok += mask.sum().item() |
| model.train() |
| loss = tot_loss / max(tot_tok, 1) |
| return loss, float(np.exp(min(loss, 20))) |
|
|
|
|
| |
| |
| |
| @torch.no_grad() |
| def greedy_generate(model, prompts, device, max_new=64, n_ctx=256, |
| amp_dtype=torch.bfloat16): |
| """Greedy-decode a list of variable-length prompts. |
| |
| Prompts are grouped by *exact* length so every batch is a clean rectangle |
| with no padding at all. Left-padding would have been the usual trick, but |
| this model has no padding mask and uses RoPE, so pad tokens would both be |
| attended to and shift every real token's position; bucketing sidesteps |
| that entirely at the cost of slightly smaller batches. |
| """ |
| model.eval() |
| by_len = {} |
| for i, p in enumerate(prompts): |
| by_len.setdefault(len(p), []).append(i) |
|
|
| out = [None] * len(prompts) |
| for length, idxs in by_len.items(): |
| for s in range(0, len(idxs), 64): |
| chunk = idxs[s:s + 64] |
| seq = torch.tensor([prompts[i] for i in chunk], device=device) |
| done = torch.zeros(len(chunk), dtype=torch.bool, device=device) |
| gen = [[] for _ in chunk] |
| for _ in range(max_new): |
| window = seq[:, -n_ctx:] |
| with torch.autocast("cuda", dtype=amp_dtype, |
| enabled=device.type == "cuda"): |
| logits, _, _ = model(window) |
| nxt = logits[:, -1].float().argmax(-1) |
| nxt = torch.where(done, torch.full_like(nxt, EOS), nxt) |
| for j, t in enumerate(nxt.tolist()): |
| if not done[j] and t != EOS: |
| gen[j].append(t) |
| done = done | (nxt == EOS) |
| if bool(done.all()): |
| break |
| seq = torch.cat([seq, nxt.unsqueeze(1)], dim=1) |
| for j, i in enumerate(chunk): |
| out[i] = gen[j] |
| model.train() |
| return out |
|
|
|
|
| def compute_bleu(model, tok, device, split="test", n_pairs=400, max_src=96, |
| max_new=64, n_ctx=256): |
| """Corpus BLEU overall and per source language, via sacrebleu.""" |
| import sacrebleu |
|
|
| pairs = eval_pairs(split, n=n_pairs) |
| prompts, refs, langs = [], [], [] |
| for lang, src, ref in pairs: |
| p = build_prompt(tok, lang, src, max_src=max_src) |
| if len(p) > n_ctx - max_new: |
| continue |
| prompts.append(p) |
| refs.append(ref) |
| langs.append(lang) |
|
|
| gen_ids = greedy_generate(model, prompts, device, max_new=max_new, n_ctx=n_ctx) |
| hyps = [tok.decode(g) for g in gen_ids] |
|
|
| res = {"bleu_all": sacrebleu.corpus_bleu(hyps, [refs]).score} |
| for lang in ("vi", "ja"): |
| h = [h for h, l in zip(hyps, langs) if l == lang] |
| r = [x for x, l in zip(refs, langs) if l == lang] |
| if h: |
| res[f"bleu_{lang}"] = sacrebleu.corpus_bleu(h, [r]).score |
| samples = [{"lang": l, "hyp": h, "ref": r} |
| for l, h, r in list(zip(langs, hyps, refs))[:12]] |
| return res, samples |
|
|
|
|
| |
| |
| |
| @torch.no_grad() |
| def expert_usage(model, loader, device, max_batches=40, amp_dtype=torch.bfloat16): |
| """Counts of (layer, expert, language) routing decisions. |
| |
| Returns ``None`` for the dense baseline, which has nothing to route. |
| """ |
| moe_layers = [i for i, b in enumerate(model.blocks) if isinstance(b.ffn, MoE)] |
| if not moe_layers: |
| return None |
| n_exp = model.blocks[moe_layers[0]].ffn.n_experts |
| n_lang = len(LANG_NAMES) |
| counts = np.zeros((len(moe_layers), n_exp, n_lang), dtype=np.int64) |
|
|
| model.eval() |
| for i, (x, _, _, lang) in enumerate(loader): |
| if i >= max_batches: |
| break |
| x = x.to(device) |
| with torch.autocast("cuda", dtype=amp_dtype, enabled=device.type == "cuda"): |
| model(x) |
| assigns = model.expert_assignments() |
| lang_flat = lang.reshape(-1).numpy() |
| for li, a in enumerate(assigns): |
| a = a.reshape(-1, a.size(-1)).cpu().numpy() |
| for slot in range(a.shape[1]): |
| np.add.at(counts[li], (a[:, slot], lang_flat), 1) |
| model.train() |
| return counts |
|
|
|
|
| def usage_summary(counts): |
| """Per-layer load-balance statistics from a usage tensor.""" |
| if counts is None: |
| return None |
| per_layer = counts.sum(axis=2).astype(float) |
| share = per_layer / per_layer.sum(axis=1, keepdims=True).clip(min=1) |
| n_exp = share.shape[1] |
| |
| ent = -(share * np.log(share.clip(1e-12))).sum(axis=1) / np.log(n_exp) |
| return { |
| "share_per_layer": share.tolist(), |
| "entropy_per_layer": ent.tolist(), |
| "mean_entropy": float(ent.mean()), |
| "max_share": float(share.max()), |
| "min_share": float(share.min()), |
| } |
|
|
