File size: 7,272 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 | """Chunked GPU bucket loss with explicit head derivatives and bounded logit cache.
Only FP32 logits are retained, rather than the per-chunk autograd graphs and
full-table slice gradients. This saves a repeated frozen-head forward in
backward. No CPU reductions or trainable weight copies leave the GPU.
"""
from __future__ import annotations
from typing import Callable, Literal, Protocol, cast
import torch
from torch import Tensor
from torch.nn import functional as F
from .ops import check_cancelled, head_product
class BucketContext(Protocol):
chunk: int
divergence: str
shape: torch.Size
saved_tensors: tuple[Tensor, ...]
def save_for_backward(self, *tensors: Tensor) -> None: ...
def mark_non_differentiable(self, *tensors: Tensor) -> None: ...
class _BucketLoss(torch.autograd.Function):
@staticmethod
def forward(ctx: BucketContext, hidden: Tensor, head: Tensor, embed: Tensor, markov: Tensor, previous: Tensor,
ids: Tensor, logp: Tensor, tail: Tensor, labels: Tensor, present: Tensor, mask: Tensor,
chunk: int, divergence: Literal["kl", "tvd", "ace"]) -> tuple[Tensor, Tensor]:
from .loss import POSITION_WEIGHTS, bucket_divergence, complementary_logaddexp, teacher_argmax
if divergence == "ace":
mask = mask & present & (ids >= 0).any(-1)
labels = teacher_argmax(ids, logp, head.shape[0])
if not bool(mask.any()):
raise ValueError("batch has no scored positions")
flat = hidden.reshape(-1, hidden.shape[-1])
emb = F.embedding(previous, embed).reshape(-1, embed.shape[-1])
shape = previous.shape
valid = ids >= 0
lse = torch.full(shape, -torch.inf, device=hidden.device)
tail_lse = lse.clone()
selected = torch.zeros_like(logp)
label_z = torch.zeros_like(tail)
saved_logits = []
for lo in range(0, head.shape[0], chunk):
check_cancelled()
hi = min(lo + chunk, head.shape[0])
z = (head_product(flat, head[lo:hi].T) + F.linear(emb, markov[lo:hi])).reshape(*shape, hi-lo)
saved_logits.append(z)
inside = valid & (ids >= lo) & (ids < hi)
loc = (ids-lo).clamp(0, hi-lo-1)
selected += 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))
part_tail = safe.logsumexp(-1).masked_fill(~any_tail, -torch.inf)
lse = torch.logaddexp(lse, z.logsumexp(-1))
tail_lse = complementary_logaddexp(tail_lse, part_tail)
chosen_label = z.gather(-1, (labels-lo).clamp(0, hi-lo-1).unsqueeze(-1)).squeeze(-1)
label_z += chosen_label.masked_fill((labels < lo) | (labels >= hi), 0)
p = logp.exp().masked_fill(~valid, 0)
logq = selected - lse.unsqueeze(-1)
logqt = tail_lse - lse
if divergence == "ace":
per = (lse-label_z).masked_fill(~mask, 0)
else:
per = bucket_divergence(logq, logqt, ids, logp, tail, divergence)
per = torch.where(present, per, .1 * (lse-label_z) if bool((~present & mask).any()) else 0.)
weights = hidden.new_tensor(POSITION_WEIGHTS).expand_as(per) * mask
norm_weights = weights / weights.sum()
loss = (per * weights).sum() / weights.sum()
ctx.chunk, ctx.divergence, ctx.shape = chunk, divergence, hidden.shape
ctx.save_for_backward(head, embed, markov, previous, ids, p, tail, labels, present, norm_weights,
lse, tail_lse, selected, *saved_logits)
detached = per.detach().clone()
ctx.mark_non_differentiable(detached)
return loss, detached
@staticmethod
def backward(ctx: BucketContext, grad: Tensor, grad_per: Tensor | None) -> tuple[Tensor | None, ...]:
(head, embed, markov, previous, ids, p, tail, labels, present, weights,
lse, tail_lse, selected, *logits) = ctx.saved_tensors
emb = F.embedding(previous, embed).reshape(-1, embed.shape[-1])
dh = torch.zeros(previous.numel(), head.shape[-1], device=head.device)
de = torch.zeros_like(emb)
dm = torch.empty_like(markov)
valid = ids >= 0
if ctx.divergence == "tvd":
qtop = (selected-lse.unsqueeze(-1)).exp().masked_fill(~valid, 0)
qt = (tail_lse-lse).exp()
atop = .5 * (qtop-p).sign()
atail = .5 * (qt-tail).sign()
weighted = (atop*qtop).sum(-1) + atail*qt
for part, z in enumerate(logits):
check_cancelled()
lo = part * ctx.chunk
hi = lo + z.shape[-1]
inside = valid & (ids >= lo) & (ids < hi)
loc = (ids-lo).clamp(0, hi-lo-1)
counts = torch.zeros_like(z, dtype=torch.int32).scatter_add(-1, loc, inside.int())
q = (z-lse.unsqueeze(-1)).exp()
if ctx.divergence == "ace":
dz = q.clone()
label_inside = (labels >= lo) & (labels < hi)
dz.scatter_add_(-1, (labels-lo).clamp(0, hi-lo-1).unsqueeze(-1), -label_inside.float().unsqueeze(-1))
elif ctx.divergence == "kl":
# Use the supplied mass, including its permitted rounding error;
# the reference differentiates this exact, unrenormalized target.
dz = q * (p.sum(-1)+tail).unsqueeze(-1)
dz -= tail.unsqueeze(-1) * (z-tail_lse.unsqueeze(-1)).exp().masked_fill(counts != 0, 0)
dz.scatter_add_(-1, loc, -p.masked_fill(~inside, 0))
else:
coeff = atail.unsqueeze(-1).expand_as(z).clone()
coeff.scatter_add_(-1, loc, (atop-atail.unsqueeze(-1)).masked_fill(~inside, 0))
dz = q * (coeff-weighted.unsqueeze(-1))
ce = q.clone()
label_inside = (labels >= lo) & (labels < hi)
ce.scatter_add_(-1, (labels-lo).clamp(0, hi-lo-1).unsqueeze(-1), -label_inside.float().unsqueeze(-1))
if ctx.divergence != "ace":
dz = torch.where(present.unsqueeze(-1), dz, .1 * ce)
dz *= (weights * grad).unsqueeze(-1)
dz2 = dz.reshape(-1, hi-lo)
dh += head_product(dz2, head[lo:hi])
de.addmm_(dz2, markov[lo:hi])
dm[lo:hi] = dz2.T @ emb
dembed = torch.zeros_like(embed)
dembed.index_add_(0, previous.flatten(), de)
return (dh.reshape(ctx.shape), None, dembed, dm, None, None, None, None, None, None, None, None, None)
def bucket_loss(hidden: Tensor, head: Tensor, embed: Tensor, markov: Tensor, previous: Tensor,
ids: Tensor, logp: Tensor, tail: Tensor, labels: Tensor, present: Tensor, mask: Tensor,
chunk: int, divergence: Literal["kl", "tvd", "ace"]) -> tuple[Tensor, Tensor]:
return cast(Callable[..., tuple[Tensor, Tensor]], _BucketLoss.apply)(
hidden, head, embed, markov, previous, ids, logp, tail, labels, present, mask, chunk, divergence)
|