| """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": |
| |
| |
| 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) |
|
|
