File size: 9,757 Bytes
e1acc15 | 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 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 | """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)
# scatter_add counts instead of scatter avoids padding overwriting
# a real top-id at the same clamped index.
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():
# Non-reentrant checkpoint tracks closed-over model parameters too.
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
# Include sampled label in statistics only as a separate lookup; it must
# not be removed from the complement bucket unless actually in the top-k.
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")
# This second streamed lookup computes ordinary sampled CE only for
# positions with missing top20. It never substitutes CE on known targets.
_, _, 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]:
# Calibrate each greedy conditional at a committed previous token. If q's
# argmax is outside top20 its exact p is unknown; do not invent a zero label.
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 # an unlabelled block is skipped, never assigned a fabricated zero probability
prediction = model.confidence(head_x.detach(), previous).sigmoid()
return (prediction[known]-target[known]).square().mean(), int(known.sum())
|