| """PyTorch implementation of the Weevil recurrent GDN2 decoder. |
| |
| Prelude runs once. The four-layer core then repeats ``core_passes`` times, |
| adding ``input_scale * prelude`` on every pass. Coda runs once. Next-token |
| logits are tied through the factorized embedding. |
| |
| Each mixer keeps a convolution window and a Gated DeltaNet-2 state. Prefill |
| runs the chunk schedule over the prompt and stores that state. Later tokens |
| advance the same state one step at a time. Depth 16 is the evaluated setting. |
| """ |
|
|
| from __future__ import annotations |
|
|
| from dataclasses import dataclass |
|
|
| import torch |
| import torch.nn.functional as F |
| from torch import nn |
| from transformers import PreTrainedModel |
| from transformers.generation import GenerationMixin |
| from transformers.modeling_outputs import CausalLMOutputWithPast |
|
|
| from .configuration_weevil_recurrent import WeevilRecurrentConfig |
|
|
|
|
| def _dtype(name: str) -> torch.dtype: |
| return torch.bfloat16 if name == "bfloat16" else torch.float32 |
|
|
|
|
| def apply_linear(x: torch.Tensor, layer: nn.Linear) -> torch.Tensor: |
| weight = layer.weight.to(dtype=x.dtype) |
| bias = None if layer.bias is None else layer.bias.to(dtype=x.dtype) |
| return F.linear(x, weight, bias) |
|
|
|
|
| def rmsnorm(x: torch.Tensor, eps: float) -> torch.Tensor: |
| x32 = x.float() |
| inv = torch.rsqrt(x32.square().mean(dim=-1, keepdim=True) + eps) |
| return (x32 * inv).to(dtype=x.dtype) |
|
|
|
|
| class WeevilRecurrentState: |
| """Convolution windows and GDN2 states for one batch of sequences. |
| |
| ``recurrent_s`` has a separate state for every core pass. Passes share |
| weights, but each pass sees a different residual and keeps its own time |
| state. ``seq_length`` is how many tokens have already been folded in. |
| """ |
|
|
| def __init__( |
| self, |
| prelude_s: torch.Tensor, |
| prelude_hist: torch.Tensor, |
| recurrent_s: torch.Tensor, |
| recurrent_hist: torch.Tensor, |
| coda_s: torch.Tensor, |
| coda_hist: torch.Tensor, |
| seq_length: int, |
| core_passes: int, |
| ) -> None: |
| self.prelude_s = prelude_s |
| self.prelude_hist = prelude_hist |
| self.recurrent_s = recurrent_s |
| self.recurrent_hist = recurrent_hist |
| self.coda_s = coda_s |
| self.coda_hist = coda_hist |
| self.seq_length = seq_length |
| self.core_passes = core_passes |
|
|
| def get_seq_length(self, layer_idx: int = 0) -> int: |
| return self.seq_length |
|
|
| def get_max_length(self) -> int: |
| return -1 |
|
|
| @property |
| def batch_size(self) -> int: |
| return self.prelude_s.shape[1] |
|
|
| def clone(self) -> "WeevilRecurrentState": |
| """Independent copy. Token steps write the state in place, so a fork has to own its tensors.""" |
|
|
| return WeevilRecurrentState( |
| self.prelude_s.clone(), |
| self.prelude_hist.clone(), |
| self.recurrent_s.clone(), |
| self.recurrent_hist.clone(), |
| self.coda_s.clone(), |
| self.coda_hist.clone(), |
| self.seq_length, |
| self.core_passes, |
| ) |
|
|
| def repeat(self, count: int) -> "WeevilRecurrentState": |
| """``count`` independent copies of each sequence, in batch order.""" |
|
|
| if count < 1: |
| raise ValueError("count must be positive") |
|
|
| def copies(tensor: torch.Tensor, dim: int) -> torch.Tensor: |
| return tensor.repeat_interleave(count, dim=dim) |
|
|
| return WeevilRecurrentState( |
| copies(self.prelude_s, 1), |
| copies(self.prelude_hist, 2), |
| copies(self.recurrent_s, 2), |
| copies(self.recurrent_hist, 3), |
| copies(self.coda_s, 1), |
| copies(self.coda_hist, 2), |
| self.seq_length, |
| self.core_passes, |
| ) |
|
|
|
|
| @dataclass |
| class WeevilCausalLMOutput(CausalLMOutputWithPast): |
| """Causal LM output. ``recurrent_state`` is the carried GDN2 state.""" |
|
|
| recurrent_state: WeevilRecurrentState | None = None |
|
|
|
|
| class CausalConv(nn.Module): |
| def __init__(self, channels: int, kernel: int) -> None: |
| super().__init__() |
| self.weight = nn.Parameter(torch.empty(channels, kernel)) |
|
|
| def forward( |
| self, x: torch.Tensor, hist: torch.Tensor | None |
| ) -> tuple[torch.Tensor, torch.Tensor]: |
| """Causal depthwise conv. ``weight[:, -1]`` multiplies the current token. |
| |
| ``hist`` is the previous ``kernel - 1`` inputs, or ``None`` for zeros. |
| Returns the activated outputs and the window to carry forward. |
| """ |
|
|
| channels, kernel = self.weight.shape |
| lookback = kernel - 1 |
| if hist is None: |
| hist = x.new_zeros(x.shape[0], lookback, channels) |
| full = x if lookback == 0 else torch.cat((hist, x), dim=1) |
| rhs = self.weight.to(dtype=x.dtype).view(channels, 1, kernel) |
| y = F.conv1d(full.transpose(1, 2), rhs, groups=channels) |
| new_hist = full[:, -lookback:] if lookback else full[:, :0] |
| return F.silu(y.transpose(1, 2)), new_hist |
|
|
|
|
| def _gdn_prepare(q, k, v, g, b, w, valid, chunk_size): |
| """Mask, normalize, and clamp. Shared by the chunk prefill and the token step.""" |
|
|
| if valid is not None: |
| mask = valid.to(dtype=q.dtype)[..., None, None] |
| q, k, v, g, b, w = q * mask, k * mask, v * mask, g * mask, b * mask, w * mask |
| q = q.float() |
| k = k.float() |
| q = q * torch.rsqrt(q.square().sum(dim=-1, keepdim=True) + 1e-6) |
| k = k * torch.rsqrt(k.square().sum(dim=-1, keepdim=True) + 1e-6) |
| q = q * (q.shape[-1] ** -0.5) |
| v = v.float() |
| g = torch.clamp(g.float(), min=-60.0 / chunk_size) |
| return q, k, v, g, b.float(), w.float() |
|
|
|
|
| def _forward_substitute(lower: torch.Tensor) -> torch.Tensor: |
| size = lower.shape[-1] |
| solved = lower |
| index = torch.arange(size, device=lower.device) |
| for row_index in range(1, size): |
| mask = (index < row_index).to(dtype=lower.dtype) |
| row = solved[..., row_index, :] |
| masked_row = row * mask |
| masked_block = solved * mask[..., :, None] * mask[..., None, :] |
| updated = row + torch.matmul(masked_row.unsqueeze(-2), masked_block).squeeze(-2) |
| updated = torch.where(mask.bool(), updated, row) |
| solved = solved.clone() |
| solved[..., row_index, :] = updated |
| eye = torch.eye(size, dtype=solved.dtype, device=solved.device) |
| return solved + eye |
|
|
|
|
| def gdn2_chunk(q, k, v, g, b, w, chunk_size: int) -> tuple[torch.Tensor, torch.Tensor]: |
| """Chunk prefill from a zero state. Inputs are already prepared, ``[B, T, H, D]``.""" |
|
|
| def heads_first(tensor: torch.Tensor) -> torch.Tensor: |
| return tensor.permute(0, 2, 1, 3) |
|
|
| q, k, v, g, b, w = (heads_first(tensor) for tensor in (q, k, v, g, b, w)) |
| batch, heads, length, k_dim = k.shape |
| v_dim = v.shape[-1] |
| pad = (chunk_size - length % chunk_size) % chunk_size |
| if pad: |
| padding = (0, 0, 0, pad) |
| q, k, v, g, b, w = (F.pad(tensor, padding) for tensor in (q, k, v, g, b, w)) |
| total = length + pad |
| n_chunks = total // chunk_size |
|
|
| def as_chunks(tensor: torch.Tensor) -> torch.Tensor: |
| return tensor.reshape(batch, heads, n_chunks, chunk_size, tensor.shape[-1]) |
|
|
| q, k, v, g, b, w = (as_chunks(tensor) for tensor in (q, k, v, g, b, w)) |
| g = torch.cumsum(g, dim=-2) |
| decay = torch.exp(g) |
| center = g[..., -1:, :] * 0.5 |
| decay_centered = torch.exp(g - center) |
| inv_decay_centered = torch.exp(center - g) |
| erase = decay * b * k |
| erase_centered = decay_centered * b * k |
| key_inv = k * inv_decay_centered |
| write = w * v |
| interaction = -torch.matmul(erase_centered, key_inv.transpose(-1, -2)) |
| strict = torch.tril(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=q.device), diagonal=-1) |
| lower = torch.where(strict, interaction, torch.zeros_like(interaction)) |
| transition = _forward_substitute(lower) |
| write = torch.matmul(transition, write) |
| k_cumdecay = torch.matmul(transition, erase) |
| query_decay = q * decay |
| query_decay_centered = q * decay_centered |
| g_last = g[..., -1:, :] |
|
|
| state = q.new_zeros(batch, heads, k_dim, v_dim) |
| outputs = [] |
| causal = torch.tril(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=q.device)) |
| for index in range(n_chunks): |
| q_decay = query_decay[:, :, index] |
| q_centered = query_decay_centered[:, :, index] |
| key_inv_c = key_inv[:, :, index] |
| write_c = write[:, :, index] |
| k_cum_c = k_cumdecay[:, :, index] |
| key_c = k[:, :, index] |
| g_c = g[:, :, index] |
| g_end = g_last[:, :, index] |
| attn = torch.matmul(q_centered, key_inv_c.transpose(-1, -2)) |
| attn = torch.where(causal, attn, torch.zeros_like(attn)) |
| residual = write_c - torch.matmul(k_cum_c, state) |
| out = torch.matmul(q_decay, state) + torch.matmul(attn, residual) |
| key_bar = key_c * torch.exp(g_end - g_c) |
| decay_state = torch.exp(g_end.squeeze(-2)).unsqueeze(-1) |
| state = state * decay_state + torch.matmul(key_bar.transpose(-1, -2), residual) |
| outputs.append(out) |
| stacked = torch.stack(outputs, dim=2).reshape(batch, heads, total, v_dim)[:, :, :length] |
| return stacked.permute(0, 2, 1, 3), state |
|
|
|
|
| def _gdn2_token_pytorch(q, k, v, g, b, w, state: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: |
| """Reference one-step update. ``q``..``w`` are ``[B, H, D]``, state is ``[B, H, D, D]``.""" |
|
|
| alpha = g.exp() |
| decayed = state * alpha.unsqueeze(-1) |
| erase = torch.einsum("bhd,bhde->bhe", b * k, decayed) |
| updated = decayed - torch.einsum("bhd,bhe->bhde", k, erase) |
| updated = updated + torch.einsum("bhd,bhe->bhde", k, w * v) |
| out = torch.einsum("bhde,bhd->bhe", updated, q) |
| return out, updated |
|
|
|
|
| try: |
| import triton |
| import triton.language as tl |
|
|
| @triton.jit |
| def _gdn2_token_kernel( |
| q_ptr, k_ptr, v_ptr, g_ptr, b_ptr, w_ptr, s_ptr, o_ptr, |
| stride_vb, stride_vh, stride_sb, stride_sh, stride_sd, |
| D: tl.constexpr, |
| ): |
| """One head. Load the DxD state, apply the delta rule, store it back.""" |
|
|
| pid_b = tl.program_id(0) |
| pid_h = tl.program_id(1) |
| offs = tl.arange(0, D) |
| base = pid_b * stride_vb + pid_h * stride_vh |
| sbase = pid_b * stride_sb + pid_h * stride_sh |
| q = tl.load(q_ptr + base + offs) |
| k = tl.load(k_ptr + base + offs) |
| v = tl.load(v_ptr + base + offs) |
| g = tl.load(g_ptr + base + offs) |
| b = tl.load(b_ptr + base + offs) |
| w = tl.load(w_ptr + base + offs) |
| state = tl.load(s_ptr + sbase + offs[:, None] * stride_sd + offs[None, :]) |
| alpha = tl.exp(g) |
| decayed = state * alpha[:, None] |
| erase = tl.sum(decayed * (b * k)[:, None], axis=0) |
| updated = decayed + k[:, None] * ((w * v)[None, :] - erase[None, :]) |
| tl.store(s_ptr + sbase + offs[:, None] * stride_sd + offs[None, :], updated) |
| tl.store(o_ptr + base + offs, tl.sum(updated * q[:, None], axis=0)) |
|
|
| @triton.jit |
| def _conv_token_kernel( |
| proj_ptr, hist_ptr, weight_ptr, out_ptr, |
| proj_stride, hist_stride_stream, hist_stride_batch, hist_stride_tap, weight_stride, |
| channels, |
| BLOCK: tl.constexpr, |
| ): |
| """Four-tap depthwise conv for one token, then slide the window down.""" |
|
|
| batch = tl.program_id(0) |
| stream = tl.program_id(1) |
| block = tl.program_id(2) |
| offs = block * BLOCK + tl.arange(0, BLOCK) |
| mask = offs < channels |
| hist_base = stream * hist_stride_stream + batch * hist_stride_batch |
| weight_base = (stream * channels + offs) * weight_stride |
| h0 = tl.load(hist_ptr + hist_base + 0 * hist_stride_tap + offs, mask=mask, other=0.0).to(tl.float32) |
| h1 = tl.load(hist_ptr + hist_base + 1 * hist_stride_tap + offs, mask=mask, other=0.0).to(tl.float32) |
| h2 = tl.load(hist_ptr + hist_base + 2 * hist_stride_tap + offs, mask=mask, other=0.0).to(tl.float32) |
| current = tl.load(proj_ptr + batch * proj_stride + stream * channels + offs, mask=mask, other=0.0) |
| w0 = tl.load(weight_ptr + weight_base + 0, mask=mask, other=0.0).to(tl.float32) |
| w1 = tl.load(weight_ptr + weight_base + 1, mask=mask, other=0.0).to(tl.float32) |
| w2 = tl.load(weight_ptr + weight_base + 2, mask=mask, other=0.0).to(tl.float32) |
| w3 = tl.load(weight_ptr + weight_base + 3, mask=mask, other=0.0).to(tl.float32) |
| acc = h0 * w0 + h1 * w1 + h2 * w2 + current.to(tl.float32) * w3 |
| out = (acc * tl.sigmoid(acc)).to(current.dtype) |
| tl.store(out_ptr + (batch * 3 + stream) * channels + offs, out, mask=mask) |
| tl.store(hist_ptr + hist_base + 0 * hist_stride_tap + offs, h1.to(current.dtype), mask=mask) |
| tl.store(hist_ptr + hist_base + 1 * hist_stride_tap + offs, h2.to(current.dtype), mask=mask) |
| tl.store(hist_ptr + hist_base + 2 * hist_stride_tap + offs, current, mask=mask) |
|
|
| except ImportError: |
| _gdn2_token_kernel = None |
| _conv_token_kernel = None |
|
|
|
|
| def _conv_token_cuda(projected: torch.Tensor, hist: torch.Tensor, weight: torch.Tensor): |
| """Slide the q/k/v windows and return the activated conv outputs, each ``[B, 1, C]``.""" |
|
|
| batch = projected.shape[0] |
| channels = hist.shape[-1] |
| out = torch.empty(batch, 3, channels, device=projected.device, dtype=projected.dtype) |
| block = 128 |
| _conv_token_kernel[(batch, 3, (channels + block - 1) // block)]( |
| projected, hist, weight, out, |
| projected.stride(0), hist.stride(0), hist.stride(1), hist.stride(2), weight.stride(0), |
| channels, BLOCK=block, |
| ) |
| return out[:, 0:1], out[:, 1:2], out[:, 2:3] |
|
|
|
|
| def _gdn2_token_cuda(q, k, v, g, b, w, state: torch.Tensor) -> torch.Tensor: |
| batch, heads, dim = q.shape |
| out = torch.empty(batch, heads, dim, device=q.device, dtype=torch.float32) |
| _gdn2_token_kernel[batch, heads]( |
| q, k, v, g, b, w, state, out, |
| q.stride(0), q.stride(1), state.stride(0), state.stride(1), state.stride(2), |
| D=dim, num_warps=4, |
| ) |
| return out |
|
|
|
|
| def gdn2_token(q, k, v, g, b, w, state: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: |
| """One time step. Inputs are prepared ``[B, 1, H, D]``. ``state`` is ``[B, H, D, D]``. |
| |
| On CUDA the state is updated in place. The returned state is that same tensor. |
| """ |
|
|
| q, k, v, g, b, w = (tensor[:, 0] for tensor in (q, k, v, g, b, w)) |
| vectors = (q, k, v, g, b, w) |
| inplace = ( |
| _gdn2_token_kernel is not None |
| and state.is_cuda |
| and state.is_contiguous() |
| and state.dtype == torch.float32 |
| and all(tensor.is_contiguous() and tensor.dtype == torch.float32 for tensor in vectors) |
| ) |
| if inplace: |
| out = _gdn2_token_cuda(*vectors, state) |
| return out[:, None], state |
| out, updated = _gdn2_token_pytorch(*vectors, state) |
| return out[:, None], updated |
|
|
|
|
| class Mixer(nn.Module): |
| def __init__(self, config: WeevilRecurrentConfig) -> None: |
| super().__init__() |
| hidden = config.hidden_size |
| width = config.num_heads * config.head_dim |
| head = config.head_dim |
| self.config = config |
| self.q_proj = nn.Linear(hidden, width, bias=False) |
| self.k_proj = nn.Linear(hidden, width, bias=False) |
| self.v_proj = nn.Linear(hidden, width, bias=False) |
| self.q_conv = CausalConv(width, config.conv_kernel) |
| self.k_conv = CausalConv(width, config.conv_kernel) |
| self.v_conv = CausalConv(width, config.conv_kernel) |
| self.f_in = nn.Linear(hidden, head, bias=False) |
| self.f_out = nn.Linear(head, width, bias=False) |
| self.b_proj = nn.Linear(hidden, width, bias=False) |
| self.w_proj = nn.Linear(hidden, width, bias=False) |
| self.g_in = nn.Linear(hidden, head, bias=False) |
| self.g_out = nn.Linear(head, width, bias=True) |
| self.o_proj = nn.Linear(width, hidden, bias=False) |
| self.o_norm = nn.Parameter(torch.ones(head)) |
| self.A_log = nn.Parameter(torch.empty(config.num_heads)) |
| self.dt_bias = nn.Parameter(torch.empty(width)) |
| self._packed_input = None |
| self._packed_conv = None |
|
|
| def _pack(self, dtype: torch.dtype) -> tuple[torch.Tensor, torch.Tensor]: |
| """One input GEMM and one depthwise conv cover q, k, v, erase, write, and the two small ins.""" |
|
|
| if ( |
| self._packed_input is not None |
| and self._packed_input.dtype == dtype |
| and self._packed_input.device == self.q_proj.weight.device |
| ): |
| return self._packed_input, self._packed_conv |
| parts = ( |
| self.q_proj.weight, |
| self.k_proj.weight, |
| self.v_proj.weight, |
| self.b_proj.weight, |
| self.w_proj.weight, |
| self.f_in.weight, |
| self.g_in.weight, |
| ) |
| self._packed_input = torch.cat([part.detach().to(dtype=dtype) for part in parts], dim=0) |
| convs = (self.q_conv.weight, self.k_conv.weight, self.v_conv.weight) |
| self._packed_conv = torch.cat([part.detach().to(dtype=dtype) for part in convs], dim=0) |
| return self._packed_input, self._packed_conv |
|
|
| def forward( |
| self, |
| x: torch.Tensor, |
| valid: torch.Tensor | None, |
| s: torch.Tensor | None, |
| hist: torch.Tensor | None, |
| ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: |
| config = self.config |
| compute = _dtype(config.compute_dtype) |
| batch, length, _ = x.shape |
| heads, head = config.num_heads, config.head_dim |
| width = heads * head |
| x_c = x.to(dtype=compute) |
| packed_w, packed_conv = self._pack(compute) |
| projected = F.linear(x_c, packed_w) |
| q, k, v, erase_in, write_in, decay_in, gate_in = projected.split( |
| (width, width, width, width, width, head, head), dim=-1 |
| ) |
| projected_q, projected_k, projected_v = q, k, v |
| channels = packed_conv.shape[0] |
| lookback = packed_conv.shape[1] - 1 |
| fast_conv = ( |
| _conv_token_kernel is not None |
| and length == 1 |
| and hist is not None |
| and lookback == 3 |
| and projected.is_cuda |
| ) |
| if fast_conv: |
| q, k, v = _conv_token_cuda(projected, hist, packed_conv) |
| new_hist = hist |
| else: |
| qkv = torch.cat((projected_q, projected_k, projected_v), dim=-1) |
| if hist is None: |
| conv_hist = qkv.new_zeros(batch, lookback, channels) |
| else: |
| conv_hist = torch.cat((hist[0], hist[1], hist[2]), dim=-1) |
| full = torch.cat((conv_hist, qkv), dim=1) |
| if length == 1: |
| activated = F.silu((full * packed_conv.transpose(0, 1)).sum(dim=1, keepdim=True)) |
| else: |
| y = F.conv1d( |
| full.transpose(1, 2), |
| packed_conv.view(channels, 1, lookback + 1), |
| groups=channels, |
| ) |
| activated = F.silu(y.transpose(1, 2)) |
| q, k, v = activated.split((width, width, width), dim=-1) |
| if length == 1 and hist is not None and lookback: |
| for shift in range(lookback - 1): |
| hist[:, :, shift].copy_(hist[:, :, shift + 1]) |
| hist[0, :, -1].copy_(projected_q[:, 0]) |
| hist[1, :, -1].copy_(projected_k[:, 0]) |
| hist[2, :, -1].copy_(projected_v[:, 0]) |
| new_hist = hist |
| else: |
| new_window = full[:, -lookback:] |
| q_hist, k_hist, v_hist = new_window.split((width, width, width), dim=-1) |
| new_hist = torch.stack((q_hist, k_hist, v_hist), dim=0) |
| decay = F.softplus( |
| apply_linear(decay_in, self.f_out).float() + self.dt_bias.float() |
| ) |
| decay = decay.reshape(batch, length, heads, head) |
| rate = torch.exp(self.A_log.float()) |
| log_decay = -rate[None, None, :, None] * decay |
| erase = torch.sigmoid(erase_in).reshape(batch, length, heads, head) |
| write = torch.sigmoid(write_in).reshape(batch, length, heads, head) |
| q, k, v, g, b, w = _gdn_prepare( |
| q.reshape(batch, length, heads, head), |
| k.reshape(batch, length, heads, head), |
| v.reshape(batch, length, heads, head), |
| log_decay, |
| erase, |
| write, |
| valid, |
| config.chunk_size, |
| ) |
| if length == 1: |
| if s is None: |
| s = q.new_zeros(batch, heads, head, head) |
| out, s_new = gdn2_token(q, k, v, g, b, w, s) |
| else: |
| out, s_new = gdn2_chunk(q, k, v, g, b, w, config.chunk_size) |
| gate = apply_linear(gate_in, self.g_out).reshape(batch, length, heads, head) |
| out32 = out.float() |
| inv = torch.rsqrt(out32.square().mean(dim=-1, keepdim=True) + config.norm_eps) |
| out = out32 * inv * self.o_norm * torch.sigmoid(gate.float()) |
| flat = out.reshape(batch, length, config.num_heads * head).to(dtype=compute) |
| return apply_linear(flat, self.o_proj), s_new, new_hist |
|
|
|
|
| class SwiGLU(nn.Module): |
| def __init__(self, config: WeevilRecurrentConfig) -> None: |
| super().__init__() |
| self.config = config |
| self.in_proj = nn.Linear(config.hidden_size, 2 * config.intermediate_size, bias=False) |
| self.out_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False) |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| compute = _dtype(self.config.compute_dtype) |
| gate, up = apply_linear(x.to(dtype=compute), self.in_proj).chunk(2, dim=-1) |
| return apply_linear(F.silu(gate) * up, self.out_proj) |
|
|
|
|
| class Gate(nn.Module): |
| def __init__(self, hidden: int) -> None: |
| super().__init__() |
| self.pre = nn.Parameter(torch.empty(hidden)) |
| self.post = nn.Parameter(torch.empty(hidden)) |
| self.alpha = nn.Parameter(torch.empty(2)) |
| self.bias = nn.Parameter(torch.empty(2)) |
|
|
| def forward(self, normed: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: |
| flat = normed.float() |
| pre = (flat * self.pre.float()).sum(dim=-1) |
| post = (flat * self.post.float()).sum(dim=-1) |
| pre = self.alpha[0].float() * pre + self.bias[0].float() |
| post = self.alpha[1].float() * post + self.bias[1].float() |
| return torch.sigmoid(pre), 2 * torch.sigmoid(post) |
|
|
|
|
| class Block(nn.Module): |
| def __init__(self, config: WeevilRecurrentConfig) -> None: |
| super().__init__() |
| self.config = config |
| self.mixer = Mixer(config) |
| self.mlp = SwiGLU(config) |
| self.mixer_gate = Gate(config.hidden_size) |
| self.mlp_gate = Gate(config.hidden_size) |
|
|
| def _read(self, residual: torch.Tensor, gate: Gate) -> tuple[torch.Tensor, torch.Tensor]: |
| normed = rmsnorm(residual, self.config.norm_eps) |
| pre, post = gate(normed) |
| hidden = (pre.unsqueeze(-1) * normed).to(dtype=_dtype(self.config.compute_dtype)) |
| return hidden, post |
|
|
| def forward( |
| self, |
| residual: torch.Tensor, |
| valid: torch.Tensor | None, |
| s: torch.Tensor | None, |
| hist: torch.Tensor | None, |
| ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: |
| hidden, post = self._read(residual, self.mixer_gate) |
| mixed, s_new, hist_new = self.mixer(hidden, valid, s, hist) |
| residual = residual + post.unsqueeze(-1) * mixed.float() |
| hidden, post = self._read(residual, self.mlp_gate) |
| residual = residual + post.unsqueeze(-1) * self.mlp(hidden).float() |
| if valid is not None: |
| residual = residual * valid.to(dtype=residual.dtype).unsqueeze(-1) |
| return residual, s_new, hist_new |
|
|
|
|
| class WeevilRecurrentPreTrainedModel(PreTrainedModel): |
| config_class = WeevilRecurrentConfig |
| base_model_prefix = "model" |
| supports_gradient_checkpointing = False |
| _no_split_modules = ["Block"] |
|
|
|
|
| class Stage(nn.Module): |
| def __init__(self, config: WeevilRecurrentConfig, count: int) -> None: |
| super().__init__() |
| self.layers = nn.ModuleList(Block(config) for _ in range(count)) |
|
|
| def forward( |
| self, |
| hidden: torch.Tensor, |
| valid: torch.Tensor | None, |
| s: torch.Tensor | None, |
| hist: torch.Tensor | None, |
| ) -> torch.Tensor: |
| for index, layer in enumerate(self.layers): |
| layer_s = None if s is None else s[index] |
| layer_hist = None if hist is None else hist[index] |
| hidden, s_new, hist_new = layer(hidden, valid, layer_s, layer_hist) |
| if s is not None: |
| if s_new.data_ptr() != s[index].data_ptr(): |
| s[index].copy_(s_new) |
| if hist_new.data_ptr() != hist[index].data_ptr(): |
| hist[index].copy_(hist_new) |
| return hidden |
|
|
|
|
| class WeevilRecurrentModel(WeevilRecurrentPreTrainedModel): |
| def __init__(self, config: WeevilRecurrentConfig) -> None: |
| super().__init__(config) |
| self.embed = nn.Embedding(config.vocab_size, config.embedding_dim) |
| self.proj = nn.Linear(config.embedding_dim, config.hidden_size, bias=False) |
| self.input_scale = nn.Parameter(torch.zeros(config.hidden_size)) |
| self.prelude = Stage(config, config.prelude_layers) |
| self.recurrent = Stage(config, config.recurrent_layers) |
| self.coda = Stage(config, config.coda_layers) |
| self.post_init() |
|
|
| def empty_state(self, batch: int, passes: int, device: torch.device) -> WeevilRecurrentState: |
| config = self.config |
| heads, dim = config.num_heads, config.head_dim |
| width = heads * dim |
| lookback = config.conv_kernel - 1 |
| compute = _dtype(config.compute_dtype) |
|
|
| def states(count: int) -> torch.Tensor: |
| return torch.zeros(count, batch, heads, dim, dim, device=device) |
|
|
| def windows(count: int) -> torch.Tensor: |
| return torch.zeros(count, 3, batch, lookback, width, device=device, dtype=compute) |
|
|
| return WeevilRecurrentState( |
| prelude_s=states(config.prelude_layers), |
| prelude_hist=windows(config.prelude_layers), |
| recurrent_s=states(passes * config.recurrent_layers).view( |
| passes, config.recurrent_layers, batch, heads, dim, dim |
| ), |
| recurrent_hist=windows(passes * config.recurrent_layers).view( |
| passes, config.recurrent_layers, 3, batch, lookback, width |
| ), |
| coda_s=states(config.coda_layers), |
| coda_hist=windows(config.coda_layers), |
| seq_length=0, |
| core_passes=passes, |
| ) |
|
|
| def embed_tokens(self, input_ids: torch.Tensor, valid: torch.Tensor | None) -> torch.Tensor: |
| compute = _dtype(self.config.compute_dtype) |
| tokens = self.embed(input_ids).to(dtype=compute) |
| hidden = apply_linear(tokens, self.proj).float() |
| if valid is not None: |
| hidden = hidden * valid.to(dtype=hidden.dtype).unsqueeze(-1) |
| return hidden |
|
|
| def _run( |
| self, |
| embedded: torch.Tensor, |
| valid: torch.Tensor | None, |
| state: WeevilRecurrentState, |
| passes: int, |
| ) -> torch.Tensor: |
| if embedded.shape[1] > 1 and state.seq_length != 0: |
| hidden = embedded |
| for step in range(embedded.shape[1]): |
| step_valid = None if valid is None else valid[:, step : step + 1] |
| hidden_step = self._run(hidden[:, step : step + 1], step_valid, state, passes) |
| if step == 0: |
| outputs = hidden_step |
| else: |
| outputs = torch.cat((outputs, hidden_step), dim=1) |
| return outputs |
| prelude = self.prelude(embedded, valid, state.prelude_s, state.prelude_hist) |
| injected = self.input_scale.float() * prelude |
| hidden = prelude |
| for index in range(passes): |
| hidden = self.recurrent( |
| hidden + injected, valid, state.recurrent_s[index], state.recurrent_hist[index] |
| ) |
| hidden = self.coda(hidden, valid, state.coda_s, state.coda_hist) |
| hidden = rmsnorm(hidden, self.config.norm_eps).to(dtype=_dtype(self.config.compute_dtype)) |
| if valid is not None: |
| hidden = hidden * valid.to(dtype=hidden.dtype).unsqueeze(-1) |
| state.seq_length += embedded.shape[1] |
| return hidden |
|
|
| def forward( |
| self, |
| input_ids: torch.Tensor, |
| attention_mask: torch.Tensor | None = None, |
| core_passes: int | None = None, |
| past_key_values: WeevilRecurrentState | None = None, |
| use_cache: bool | None = None, |
| **kwargs, |
| ) -> tuple[torch.Tensor, WeevilRecurrentState | None]: |
| passes = self.config.core_passes if core_passes is None else core_passes |
| if passes < 1: |
| raise ValueError("core_passes must be positive") |
| if attention_mask is not None and attention_mask.shape[-1] != input_ids.shape[-1]: |
| attention_mask = attention_mask[:, -input_ids.shape[-1] :] |
| valid = None if attention_mask is None else attention_mask.to(dtype=torch.bool) |
| state = past_key_values |
| if state is None and use_cache: |
| state = self.empty_state(input_ids.shape[0], passes, input_ids.device) |
| if state is not None: |
| if state.core_passes != passes: |
| raise ValueError( |
| f"state was built for {state.core_passes} core passes, not {passes}" |
| ) |
| if state.batch_size != input_ids.shape[0]: |
| raise ValueError("state batch does not match input_ids") |
| embedded = self.embed_tokens(input_ids, valid) |
| if state is None: |
| state = self.empty_state(input_ids.shape[0], passes, input_ids.device) |
| hidden = self._run(embedded, valid, state, passes) |
| state = None |
| else: |
| hidden = self._run(embedded, valid, state, passes) |
| return hidden, state |
|
|
|
|
| class WeevilRecurrentForCausalLM(WeevilRecurrentPreTrainedModel, GenerationMixin): |
| def __init__(self, config: WeevilRecurrentConfig) -> None: |
| super().__init__(config) |
| self.model = WeevilRecurrentModel(config) |
| self.post_init() |
| self.generation_config.use_cache = True |
|
|
| @classmethod |
| def _supports_default_dynamic_cache(cls) -> bool: |
| return False |
|
|
| def forward( |
| self, |
| input_ids: torch.Tensor | None = None, |
| attention_mask: torch.Tensor | None = None, |
| core_passes: int | None = None, |
| past_key_values: WeevilRecurrentState | None = None, |
| use_cache: bool | None = None, |
| labels: torch.Tensor | None = None, |
| **kwargs, |
| ) -> WeevilCausalLMOutput: |
| if input_ids is None: |
| raise ValueError("input_ids is required") |
| if use_cache is None: |
| use_cache = past_key_values is not None or labels is None |
| hidden, state = self.model( |
| input_ids, |
| attention_mask=attention_mask, |
| core_passes=core_passes, |
| past_key_values=past_key_values, |
| use_cache=use_cache, |
| ) |
| compute = _dtype(self.config.compute_dtype) |
| low = hidden.to(dtype=compute) @ self.model.proj.weight.to(dtype=compute) |
| logits = (low @ self.model.embed.weight.to(dtype=compute).T).float() |
| loss = None |
| if labels is not None: |
| shift_logits = logits[:, :-1].contiguous() |
| shift_labels = labels[:, 1:].contiguous() |
| loss = F.cross_entropy(shift_logits.view(-1, shift_logits.shape[-1]), shift_labels.view(-1)) |
| return WeevilCausalLMOutput( |
| loss=loss, |
| logits=logits, |
| past_key_values=state, |
| recurrent_state=state, |
| ) |
|
|
| def score_continuations( |
| self, |
| input_ids: torch.Tensor, |
| continuations: torch.Tensor | list, |
| continuation_mask: torch.Tensor | None = None, |
| core_passes: int | None = None, |
| ) -> torch.Tensor: |
| """Mean log-probability of each continuation after one prefill of ``input_ids``. |
| |
| ``continuations`` is a list of token-id sequences or a right-padded ``[choices, length]`` |
| tensor. A padded tensor uses ``pad_token_id`` as padding unless ``continuation_mask`` |
| marks the real tokens. The prompt state is copied once per choice, so the choices share |
| the prefill and do not interfere. |
| """ |
|
|
| if input_ids.ndim == 1: |
| input_ids = input_ids.unsqueeze(0) |
| if input_ids.shape[0] != 1: |
| raise ValueError("score_continuations prefills one prompt") |
| encoded, mask = _pack_continuations( |
| continuations, continuation_mask, self.config.pad_token_id, input_ids.device |
| ) |
| prefill = self(input_ids, core_passes=core_passes, use_cache=True) |
| first = F.log_softmax(prefill.logits[0, -1].float(), dim=-1) |
| token_scores = first.new_zeros(encoded.shape) |
| token_scores[:, 0] = first[encoded[:, 0]] |
| if encoded.shape[1] > 1 and bool(mask[:, 1:].any()): |
| branch = prefill.recurrent_state.repeat(encoded.shape[0]) |
| scored = self( |
| encoded[:, :-1], |
| attention_mask=mask[:, :-1].to(dtype=torch.long), |
| past_key_values=branch, |
| core_passes=core_passes, |
| use_cache=True, |
| ) |
| taken = F.log_softmax(scored.logits.float(), dim=-1) |
| token_scores[:, 1:] = taken.gather(-1, encoded[:, 1:].unsqueeze(-1)).squeeze(-1) |
| return (token_scores * mask).sum(dim=-1) / mask.sum(dim=-1) |
|
|
|
|
| def _pack_continuations(continuations, mask, pad_id: int, device: torch.device): |
| if isinstance(continuations, torch.Tensor): |
| if continuations.ndim == 1: |
| continuations = continuations.unsqueeze(0) |
| if continuations.ndim != 2: |
| raise ValueError("continuations tensor must have shape [choices, length]") |
| encoded = continuations.to(device=device, dtype=torch.long) |
| if encoded.shape[0] == 0 or encoded.shape[1] == 0: |
| raise ValueError("each continuation must contain at least one token") |
| if mask is None: |
| mask = encoded.ne(pad_id) |
| else: |
| mask = mask.to(device=device, dtype=torch.bool) |
| if mask.shape != encoded.shape: |
| raise ValueError("continuation_mask must match continuations") |
| if not bool(mask[:, 0].all()): |
| raise ValueError("continuations must be left-aligned") |
| return encoded, mask |
| rows = [list(row) for row in continuations] |
| if not rows or any(len(row) == 0 for row in rows): |
| raise ValueError("each continuation must contain at least one token") |
| width = max(len(row) for row in rows) |
| encoded = torch.full((len(rows), width), pad_id, dtype=torch.long, device=device) |
| valid = torch.zeros(len(rows), width, dtype=torch.bool, device=device) |
| for index, row in enumerate(rows): |
| encoded[index, : len(row)] = torch.tensor(row, dtype=torch.long, device=device) |
| valid[index, : len(row)] = True |
| return encoded, valid |
|
|
|
|
| WeevilRecurrentConfig.register_for_auto_class() |
| WeevilRecurrentModel.register_for_auto_class("AutoModel") |
| WeevilRecurrentForCausalLM.register_for_auto_class("AutoModelForCausalLM") |
|
|
