File size: 6,568 Bytes
5407bdc | 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 | """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
# ---------------------------------------------------------------------------
# Perplexity
# ---------------------------------------------------------------------------
@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)))
# ---------------------------------------------------------------------------
# Greedy generation (batched over equal-length prompts)
# ---------------------------------------------------------------------------
@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
# ---------------------------------------------------------------------------
# Expert usage
# ---------------------------------------------------------------------------
@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() # list of (B, N, k)
lang_flat = lang.reshape(-1).numpy()
for li, a in enumerate(assigns):
a = a.reshape(-1, a.size(-1)).cpu().numpy() # (B*N, k)
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) # (L, E)
share = per_layer / per_layer.sum(axis=1, keepdims=True).clip(min=1)
n_exp = share.shape[1]
# Normalised entropy: 1.0 = perfectly balanced, 0.0 = one expert takes all.
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()),
}
|