| """Top-20 plus complementary-tail forward KL, with chunk-recomputed logits.""" |
| from __future__ import annotations |
|
|
| from typing import Callable, Literal, Protocol |
|
|
| import torch |
| from torch import Tensor |
| from torch.utils.checkpoint import checkpoint |
|
|
| from .model import Drafter |
|
|
| POSITION_WEIGHTS = (1., .8, .65, .5, .4) |
|
|
|
|
| class Vocab(Protocol): |
| @property |
| def vocab(self) -> int: ... |
|
|
|
|
| class LogitModel(Protocol): |
| @property |
| def c(self) -> Vocab: ... |
| def logit_chunk(self, hidden: Tensor, previous: Tensor, start: int, end: int) -> Tensor: ... |
|
|
|
|
| def complementary_logaddexp(a: Tensor, b: Tensor) -> Tensor: |
| """Merge possibly empty chunks without the undefined (-inf, -inf) gradient.""" |
| empty = torch.isneginf(a) & torch.isneginf(b) |
| return torch.logaddexp(a.masked_fill(empty, 0), b.masked_fill(empty, 0)).masked_fill(empty, -torch.inf) |
|
|
|
|
| def bucket_divergence(logq: Tensor, logqt: Tensor, ids: Tensor, logp: Tensor, tail: Tensor, |
| divergence: Literal["kl", "tvd"]) -> Tensor: |
| """Coarse top-k plus tail divergence; zero teacher buckets contribute zero. |
| |
| This is not full-vocabulary KL/TV: within-tail allocation is unobserved. |
| logqt comes from the complementary logsumexp, never 1-sum(q_top). |
| """ |
| p = logp.exp().masked_fill(ids < 0, 0) |
| if divergence == "kl": |
| top_loss = (torch.xlogy(p, p) - p * logq.masked_fill(p == 0, 0)).sum(-1) |
| tail_loss = torch.xlogy(tail, tail) - tail * logqt.masked_fill(tail == 0, 0) |
| return top_loss + tail_loss |
| q = logq.exp().masked_fill(ids < 0, 0) |
| return .5 * ((p-q).abs().sum(-1) + (tail-logqt.exp()).abs()) |
|
|
|
|
| def bucket_statistics(logits: Callable[[int, int], Tensor], ids: Tensor, vocab: int, chunk: int, |
| recompute: bool = True) -> tuple[Tensor, Tensor, Tensor]: |
| """logZ, logZ(complement), selected logits. No subtraction of nearly equal masses. |
| |
| Each chunk checkpoint recomputes its logits in backward rather than retaining |
| the 129K-vocabulary activation/Markov head graph. -1 is missing top-k padding. |
| """ |
| valid = ids >= 0 |
| if chunk < 1 or bool((ids >= vocab).any()): |
| raise ValueError("invalid chunk size or top-k id") |
| sorted_ids = ids.masked_fill(~valid, vocab).sort(-1).values |
| if bool(((sorted_ids[..., 1:] == sorted_ids[..., :-1]) & (sorted_ids[..., 1:] != vocab)).any()): |
| raise ValueError("duplicate top-k ids") |
| lse = torch.full(ids.shape[:-1], -torch.inf, device=ids.device, dtype=torch.float32) |
| tail = lse.clone() |
| selected = torch.zeros_like(ids, dtype=torch.float32) |
| for lo in range(0, vocab, chunk): |
| hi = min(vocab, lo + chunk) |
|
|
| def part(dummy: Tensor, start: int = lo, end: int = hi) -> tuple[Tensor, Tensor, Tensor]: |
| z = logits(start, end).float() |
| inside = valid & (ids >= start) & (ids < end) |
| loc = (ids - start).clamp(0, end-start-1) |
| chosen = z.gather(-1, loc).masked_fill(~inside, 0) |
| |
| |
| counts = torch.zeros_like(z, dtype=torch.int32).scatter_add(-1, loc, inside.int()) |
| remainder = z.masked_fill(counts != 0, -torch.inf) |
| any_tail = (counts == 0).any(-1) |
| safe = torch.where(any_tail.unsqueeze(-1), remainder, torch.zeros_like(remainder)) |
| tail_z = safe.logsumexp(-1).masked_fill(~any_tail, -torch.inf) |
| return z.logsumexp(-1), tail_z, chosen |
|
|
| if recompute and torch.is_grad_enabled(): |
| |
| a, b, t = checkpoint(part, lse, use_reentrant=False) |
| else: |
| a, b, t = part(lse) |
| lse, tail, selected = torch.logaddexp(lse, a), complementary_logaddexp(tail, b), selected + t |
| return lse, tail, selected |
|
|
|
|
| def teacher_argmax(ids: Tensor, logp: Tensor, vocab: int) -> Tensor: |
| """Teacher top-k mode; equal log probabilities choose the lowest token id. |
| |
| Empty rows get a harmless lookup id and are excluded by the ACE mask. |
| """ |
| valid = ids >= 0 |
| best = logp.masked_fill(~valid, -torch.inf).amax(-1, keepdim=True) |
| picks = ids.masked_fill(~valid | (logp != best), vocab).amin(-1) |
| return picks.clamp_max(vocab-1) |
|
|
|
|
| def distillation_loss(model: LogitModel, hidden: Tensor, previous: Tensor, top_ids: Tensor, top_logp: Tensor, |
| tail_mass: Tensor, labels: Tensor, present: Tensor, position_mask: Tensor, |
| chunk: int = 4096, divergence: Literal["kl", "tvd", "ace"] = "kl" |
| ) -> tuple[Tensor, Tensor]: |
| if divergence not in ("kl", "tvd", "ace"): |
| raise ValueError(divergence) |
| valid = top_ids >= 0 |
| p = top_logp.exp().masked_fill(~valid, 0) |
| if bool((~torch.isfinite(p)).any()) or bool((~torch.isfinite(tail_mass)).any()): |
| raise ValueError("nonfinite target distribution") |
| if bool(((tail_mass < 0) | (tail_mass > 1)).any()): |
| raise ValueError("tail mass outside [0,1]") |
| observed = present & valid.any(-1) if divergence == "ace" else present |
| if not torch.allclose((p.sum(-1)+tail_mass)[observed], torch.ones_like(tail_mass[observed]), atol=1e-4, rtol=0): |
| raise ValueError("top-k probabilities plus tail do not sum to one") |
| if divergence != "ace" and bool((present & ~valid.any(-1)).any()): |
| raise ValueError("present top-k distribution is empty") |
| if divergence == "ace": |
| position_mask = position_mask & present & valid.any(-1) |
| if isinstance(model, Drafter) and model.fast_training: |
| from .fast_loss import bucket_loss |
| if chunk < 1 or bool((top_ids >= model.c.vocab).any()): |
| raise ValueError("invalid chunk size or top-k id") |
| ordered = top_ids.masked_fill(~valid, model.c.vocab).sort(-1).values |
| if bool(((ordered[..., 1:] == ordered[..., :-1]) & (ordered[..., 1:] != model.c.vocab)).any()): |
| raise ValueError("duplicate top-k ids") |
| missing = ~present & position_mask |
| if bool(missing.any()) and (bool(((labels < 0) & missing).any()) or bool((labels >= model.c.vocab).any())): |
| raise ValueError("missing-top20 position lacks a sampled label") |
| if not bool(position_mask.any()): |
| raise ValueError("batch has no scored positions") |
| loss, per = bucket_loss(hidden, model.t["head.weight"], model.last("markov_head.embed.weight"), |
| model.last("markov_head.head.weight"), previous, top_ids, top_logp, |
| tail_mass, labels, present, position_mask, chunk, divergence) |
| if not bool(torch.isfinite(loss)): |
| raise FloatingPointError("nonfinite distillation loss") |
| return loss, per |
| |
| |
| fn = lambda lo, hi: model.logit_chunk(hidden, previous, lo, hi) |
| lse, tail_lse, selected = bucket_statistics(fn, top_ids, model.c.vocab, chunk) |
| logq = selected - lse.unsqueeze(-1) |
| logqt = tail_lse - lse |
| if divergence == "ace": |
| picks = teacher_argmax(top_ids, top_logp, model.c.vocab) |
| per = lse - selected.masked_fill(top_ids != picks.unsqueeze(-1), 0).sum(-1) |
| per = per.masked_fill(~position_mask, 0) |
| else: |
| per = bucket_divergence(logq, logqt, top_ids, top_logp, tail_mass, divergence) |
| missing = ~present & position_mask |
| if bool(missing.any()): |
| safe_label = labels.clamp_min(0) |
| if bool(((labels < 0) & missing).any()) or bool((labels >= model.c.vocab).any()): |
| raise ValueError("missing-top20 position lacks a sampled label") |
| |
| |
| _, _, label_logits = bucket_statistics(fn, safe_label.unsqueeze(-1), model.c.vocab, chunk) |
| per = torch.where(present, per, .1 * (lse-label_logits.squeeze(-1))) |
| else: |
| per = torch.where(present, per, 0) |
| weights = hidden.new_tensor(POSITION_WEIGHTS).expand_as(per) * position_mask |
| if not bool(weights.sum() > 0): |
| raise ValueError("batch has no scored positions") |
| loss = (per * weights).sum() / weights.sum() |
| if not bool(torch.isfinite(loss)): |
| raise FloatingPointError("nonfinite distillation loss") |
| return loss, per.detach() |
|
|
|
|
| def confidence_loss(model: Drafter, hidden: Tensor, head_x: Tensor, previous: Tensor, top_ids: Tensor, |
| top_logp: Tensor, position_mask: Tensor, chunk: int = 4096) -> tuple[Tensor | None, int]: |
| |
| |
| with torch.no_grad(): |
| best = torch.full(previous.shape, -torch.inf, device=hidden.device) |
| picks = torch.zeros_like(previous) |
| for lo in range(0, model.c.vocab, chunk): |
| val, idx = model.logit_chunk(hidden, previous, lo, min(lo+chunk, model.c.vocab)).max(-1) |
| picks = torch.where(val > best, idx+lo, picks) |
| best = torch.maximum(best, val) |
| match = top_ids == picks.unsqueeze(-1) |
| known = match.any(-1) & position_mask |
| target = top_logp.exp().masked_fill(~match, 0).sum(-1) |
| if not bool(known.any()): |
| return None, 0 |
| prediction = model.confidence(head_x.detach(), previous).sigmoid() |
| return (prediction[known]-target[known]).square().mean(), int(known.sum()) |
|
|
