Buckets:
| """Teacher top-K union and bounded-memory text projection for UniAR OPD.""" | |
| from dataclasses import dataclass | |
| import math | |
| import torch | |
| from torch.utils.checkpoint import checkpoint | |
| MASK_LOGIT = -1.0e30 | |
| class TextSupport: | |
| token_ids: torch.Tensor | |
| valid: torch.Tensor | |
| retained_mass: torch.Tensor | |
| top_k: int | |
| def select(self, logits): | |
| return logits.gather(-1, self.token_ids).float().masked_fill(~self.valid, MASK_LOGIT) | |
| def build_text_support(logits, *, top_k: int, minimum_mass: float) -> TextSupport: | |
| """Grow K until every teacher view at every position retains enough mass.""" | |
| if not logits or any(value.ndim != 2 or value.shape != logits[0].shape for value in logits): | |
| raise ValueError("aligned [positions, vocabulary] teacher logits are required") | |
| if top_k <= 0 or not 0 < minimum_mass <= 1: | |
| raise ValueError("invalid text support budget or retained mass") | |
| vocab = logits[0].shape[-1] | |
| normalizers = [torch.logsumexp(value.float(), -1, keepdim=True) for value in logits] | |
| k = min(top_k, vocab) | |
| while True: | |
| if k == vocab: | |
| ids = torch.arange(vocab, device=logits[0].device).expand(logits[0].shape[0], -1) | |
| valid = torch.ones_like(ids, dtype=torch.bool) | |
| else: | |
| ids = torch.cat([value.topk(k, dim=-1).indices for value in logits], -1).sort(-1).values | |
| valid = torch.ones_like(ids, dtype=torch.bool) | |
| valid[:, 1:] = ids[:, 1:] != ids[:, :-1] | |
| mass = torch.stack([ | |
| ((value.gather(-1, ids).float() - log_z).exp() * valid).sum(-1) | |
| for value, log_z in zip(logits, normalizers, strict=True) | |
| ]) | |
| if k == vocab or bool((mass >= minimum_mass).all()): | |
| return TextSupport(ids, valid, mass, k) | |
| k = min(k * 2, vocab) | |
| def project_text_support(ar, hidden, support: TextSupport, *, chunk_size: int = 16): | |
| """Recompute small vocabulary projections in backward, preserving LM-head LoRA.""" | |
| def project(states, ids, valid): | |
| # Use the canonical projector (RMSNorm + PEFT lm_head), including all | |
| # trainable head adapters. Never read only lm_head.base_layer.weight. | |
| return ar.project_text(states).gather(-1, ids).float().masked_fill(~valid, MASK_LOGIT) | |
| outputs = [] | |
| for start in range(0, hidden.shape[0], chunk_size): | |
| args = (hidden[start:start + chunk_size], support.token_ids[start:start + chunk_size], support.valid[start:start + chunk_size]) | |
| outputs.append(checkpoint(project, *args, use_reentrant=False) if torch.is_grad_enabled() else project(*args)) | |
| return torch.cat(outputs) | |
| def clipped_text_jsd(target_logp, student_logits, *, contribution_clip): | |
| teacher = target_logp.detach().float() | |
| student = student_logits.float().log_softmax(-1) | |
| mixture = torch.logaddexp(teacher, student) - math.log(2.0) | |
| contribution = 0.5 * (teacher.exp() * (teacher - mixture) + student.exp() * (student - mixture)) | |
| clipped = (contribution > contribution_clip).float() | |
| loss = contribution.clamp(max=contribution_clip).sum(-1).mean() | |
| return loss, clipped | |
Xet Storage Details
- Size:
- 3.17 kB
- Xet hash:
- 9fb2cf7b0788660c87596690c3d1e84dfc09fd4ae38046ba4a70e7da583e9268
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.