Spaces:
Sleeping
Sleeping
| """Early-exit heads: full-vocabulary heads and pruned heads with a reject class. | |
| A pruned head k predicts over its vocabulary V_k plus one extra output, the | |
| reject class (⊥). Its training target is the full-model distribution collapsed | |
| onto V_k ∪ {⊥}: p̄(v) = p(v) for v ∈ V_k and p̄(⊥) = p(V \\ V_k). | |
| """ | |
| from __future__ import annotations | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| class RMSNorm(nn.Module): | |
| def __init__(self, dim: int, eps: float): | |
| super().__init__() | |
| self.weight = nn.Parameter(torch.ones(dim)) | |
| self.eps = eps | |
| def forward(self, x: torch.Tensor) -> torch.Tensor: | |
| dtype = x.dtype | |
| x = x.float() | |
| x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) | |
| return self.weight * x.to(dtype) | |
| class Adapter(nn.Module): | |
| """Residual low-rank MLP before an early head (Kangaroo/EAGLE-style, no attention, | |
| hence no extra KV cache): h + up(silu(down(norm(h)))), identity at initialisation.""" | |
| def __init__(self, hidden_size: int, rank: int, eps: float): | |
| super().__init__() | |
| self.norm = RMSNorm(hidden_size, eps) | |
| self.down = nn.Linear(hidden_size, rank, bias=False) | |
| self.up = nn.Linear(rank, hidden_size, bias=False) | |
| nn.init.zeros_(self.up.weight) | |
| def forward(self, h: torch.Tensor) -> torch.Tensor: | |
| return h + self.up(F.silu(self.down(self.norm(h)))) | |
| class ExitHead(nn.Module): | |
| """Norm + linear projection onto `vocab_ids` (+ optional reject class). | |
| vocab_ids=None means the full vocabulary (output index == token id). | |
| """ | |
| def __init__(self, hidden_size: int, vocab_size: int, eps: float, | |
| vocab_ids: torch.Tensor | None = None, reject: bool = False, adapter_rank: int = 0): | |
| super().__init__() | |
| self.adapter = Adapter(hidden_size, adapter_rank, eps) if adapter_rank > 0 else None | |
| self.norm = RMSNorm(hidden_size, eps) | |
| self.vocab_size = vocab_size | |
| n_out = vocab_size if vocab_ids is None else int(vocab_ids.numel()) | |
| self.reject = reject | |
| self.proj = nn.Linear(hidden_size, n_out + int(reject), bias=False) | |
| # The ⊥ logit gets its own scalar bias (initialised to the log of the | |
| # mean out-of-vocabulary mass when pruning). | |
| self.reject_bias = nn.Parameter(torch.zeros(1)) if reject else None | |
| if vocab_ids is None: | |
| self.register_buffer("vocab_ids", None) | |
| else: | |
| self.register_buffer("vocab_ids", vocab_ids.long().clone()) | |
| def n_vocab(self) -> int: | |
| return self.proj.out_features - int(self.reject) | |
| def forward(self, h: torch.Tensor) -> torch.Tensor: | |
| """Logits over V_k (∪ ⊥ as the last column), float32.""" | |
| if self.adapter is not None: | |
| h = self.adapter(h) | |
| logits = self.proj(self.norm(h)).float() | |
| if self.reject: | |
| logits = torch.cat([logits[..., :-1], logits[..., -1:] + self.reject_bias.float()], dim=-1) | |
| return logits | |
| def to_token_ids(self, idx: torch.Tensor) -> torch.Tensor: | |
| # gather (not advanced indexing) so that it stays capturable in CUDA/HIP graphs | |
| if self.vocab_ids is None: | |
| return idx | |
| return self.vocab_ids.gather(0, idx.reshape(-1)).reshape(idx.shape) | |
| def scatter_full(self, probs_k: torch.Tensor) -> torch.Tensor: | |
| """Map probabilities over V_k (no ⊥ column) to the full vocabulary.""" | |
| if self.vocab_ids is None: | |
| return probs_k | |
| out = probs_k.new_zeros(*probs_k.shape[:-1], self.vocab_size) | |
| out[..., self.vocab_ids] = probs_k | |
| return out | |
| def from_full(cls, full: "ExitHead", vocab_ids: torch.Tensor, reject: bool = True, | |
| reject_bias: float = 0.0, adapter_rank: int = 0) -> "ExitHead": | |
| """Prune a full head to vocab_ids, appending a zero-weight ⊥ row (optional adapter).""" | |
| assert full.vocab_ids is None | |
| d = full.proj.in_features | |
| head = cls(d, full.vocab_size, full.norm.eps, vocab_ids, reject, adapter_rank) | |
| with torch.no_grad(): | |
| head.norm.weight.copy_(full.norm.weight) | |
| head.proj.weight[: vocab_ids.numel()].copy_(full.proj.weight[vocab_ids]) | |
| if reject: | |
| head.proj.weight[-1].zero_() | |
| head.reject_bias.fill_(reject_bias) | |
| return head | |
| def collapse_target(p_full: torch.Tensor, vocab_ids: torch.Tensor) -> torch.Tensor: | |
| """p̄ over V_k ∪ {⊥} from full-vocabulary probabilities (…, V).""" | |
| kept = p_full[..., vocab_ids] | |
| rest = (1.0 - kept.sum(-1, keepdim=True)).clamp_min(0.0) | |
| return torch.cat([kept, rest], dim=-1) | |
| def head_draft_distribution(logits: torch.Tensor, reject: bool, temperature: float): | |
| """From head logits return (q over V_k renormalised, reject prob, confidence). | |
| The confidence is the max probability over V_k ∪ {⊥} (⊥ included in the | |
| softmax), so a head that puts its mass on ⊥ is not confident. At T=0 the | |
| returned q is one-hot on the argmax over V_k. | |
| """ | |
| probs = F.softmax(logits, dim=-1) | |
| if reject: | |
| p_rej = probs[..., -1] | |
| conf, arg = probs.max(-1) | |
| is_rej = arg == probs.shape[-1] - 1 | |
| conf = torch.where(is_rej, torch.zeros_like(conf), conf) | |
| body_logits = logits[..., :-1] | |
| else: | |
| p_rej = torch.zeros(probs.shape[:-1], device=probs.device) | |
| conf = probs.max(-1).values | |
| body_logits = logits | |
| if temperature == 0.0: | |
| q = F.one_hot(body_logits.argmax(-1), body_logits.shape[-1]).float() | |
| else: | |
| q = F.softmax(body_logits / temperature, dim=-1) | |
| return q, p_rej, conf | |
| def entropy(logits: torch.Tensor) -> torch.Tensor: | |
| logp = F.log_softmax(logits, dim=-1) | |
| return -(logp.exp() * logp).sum(-1) | |