| """Portable PyTorch Weevil, with fp32 recurrence and fixed-size generation caches. |
| |
| This reference implementation needs only PyTorch and Transformers, including for |
| the latent-QKV variant. Both embedding factors are tied to the vocabulary readout. |
| """ |
|
|
| from typing import ClassVar |
|
|
| import torch |
| from torch import nn |
| from torch.nn import functional as F |
| from transformers import GenerationMixin, PreTrainedModel |
| from transformers.cache_utils import Cache |
| from transformers.modeling_outputs import CausalLMOutputWithPast |
|
|
| from .configuration_weevil import WeevilConfig |
|
|
|
|
| class WeevilCache(Cache): |
| def __init__(self, states=None, seen_tokens=0): |
| super().__init__(layers=[]) |
| self.states = states or [] |
| self.seen_tokens = seen_tokens |
|
|
| def get_seq_length(self, layer_idx=0): |
| return self.seen_tokens |
|
|
| def reorder_cache(self, beam_idx): |
| self.states = [tuple(x.index_select(0, beam_idx.to(x.device)) for x in state) |
| for state in self.states] |
|
|
| def batch_repeat_interleave(self, repeats): |
| self.states = [tuple(x.repeat_interleave(repeats, dim=0) for x in state) |
| for state in self.states] |
|
|
|
|
| class ZCRMSNorm(nn.Module): |
| def __init__(self, width, eps): |
| super().__init__() |
| self.scale = nn.Parameter(torch.zeros(width)) |
| self.eps = eps |
|
|
| def forward(self, x): |
| x = x.float() |
| return (1 + self.scale.float()) * x * torch.rsqrt(x.square().mean(-1, keepdim=True) + self.eps) |
|
|
|
|
| class TiedEmbedding(nn.Module): |
| def __init__(self, c): |
| super().__init__() |
| self.embed = nn.Embedding(c.vocab_size, c.embedding_dim) |
| self.proj = nn.Linear(c.embedding_dim, c.hidden_size, bias=False) |
|
|
| def forward(self, ids): |
| return self.proj(self.embed(ids)) |
|
|
| def token_logits(self, x): |
| x = F.linear(x.to(self.proj.weight.dtype), self.proj.weight.T) |
| return F.linear(x, self.embed.weight).float() |
|
|
|
|
| class CausalConv(nn.Module): |
| def __init__(self, channels, kernel): |
| super().__init__() |
| self.weight = nn.Parameter(torch.empty(channels, kernel)) |
| nn.init.normal_(self.weight, std=0.02) |
|
|
| def step(self, x, history, valid): |
| taps = torch.cat((history, x.float().unsqueeze(1)), dim=1) |
| output = (taps * self.weight.float().T).sum(1) |
| history = torch.where(valid[:, None, None], taps[:, 1:], history) |
| return F.silu(output.to(x.dtype)), history |
|
|
|
|
| class GDN2Mixer(nn.Module): |
| def __init__(self, c): |
| super().__init__() |
| self.config = c |
| hidden, head = c.hidden_size, c.hidden_size // c.num_heads |
| width = c.num_heads * head |
|
|
| def linear(i, o, bias=False): |
| return nn.Linear(i, o, bias=bias) |
|
|
| if c.qkv_rank is None: |
| self.q_proj, self.k_proj, self.v_proj = [linear(hidden, width) for _ in range(3)] |
| else: |
| self.qkv_a = linear(hidden, c.qkv_rank) |
| self.q_b, self.k_b, self.v_b = [linear(c.qkv_rank, width) for _ in range(3)] |
| self.q_conv, self.k_conv, self.v_conv = [CausalConv(width, c.conv_kernel) for _ in range(3)] |
| self.f_in, self.f_out = linear(hidden, head), linear(head, width) |
| self.b_proj, self.w_proj = linear(hidden, width), linear(hidden, width) |
| self.g_in, self.g_out = linear(hidden, head), linear(head, width, bias=True) |
| self.o_proj = linear(width, hidden) |
| self.o_norm = nn.Parameter(torch.ones(head)) |
| self.A_log = nn.Parameter(torch.zeros(c.num_heads)) |
| self.dt_bias = nn.Parameter(torch.zeros(width)) |
|
|
| def forward(self, x, valid, state=None): |
| c = self.config |
| batch, length, _ = x.shape |
| head = c.hidden_size // c.num_heads |
| width = c.num_heads * head |
| if state is None: |
| recurrent = torch.zeros(batch, c.num_heads, head, head, device=x.device) |
| histories = [torch.zeros(batch, c.conv_kernel - 1, width, device=x.device) for _ in range(3)] |
| else: |
| recurrent, *histories = state |
| if c.qkv_rank is None: |
| projections = [self.q_proj(x), self.k_proj(x), self.v_proj(x)] |
| else: |
| latent = self.qkv_a(x) |
| projections = [self.q_b(latent), self.k_b(latent), self.v_b(latent)] |
| shape = (batch, length, c.num_heads, head) |
| decay = -self.A_log.float().exp()[None, None, :, None] * F.softplus( |
| self.f_out(self.f_in(x)).float() + self.dt_bias.float()).reshape(shape) |
| erase = self.b_proj(x).sigmoid().float().reshape(shape) |
| write = self.w_proj(x).sigmoid().float().reshape(shape) |
| outputs = [] |
| for t in range(length): |
| active = valid[:, t] |
| qkv = [] |
| for i, conv in enumerate((self.q_conv, self.k_conv, self.v_conv)): |
| value, histories[i] = conv.step(projections[i][:, t], histories[i], active) |
| qkv.append(value.float().reshape(batch, c.num_heads, head)) |
| q, k, v = qkv |
| q = q * torch.rsqrt(q.square().sum(-1, keepdim=True) + 1e-6) * head ** -0.5 |
| k = k * torch.rsqrt(k.square().sum(-1, keepdim=True) + 1e-6) |
| decayed = recurrent * decay[:, t].exp().unsqueeze(-1) |
| erased = ((erase[:, t] * k).unsqueeze(-1) * decayed).sum(-2) |
| updated = decayed + k.unsqueeze(-1) * (write[:, t] * v - erased).unsqueeze(-2) |
| recurrent = torch.where(active[:, None, None, None], updated, recurrent) |
| outputs.append((q.unsqueeze(-1) * recurrent).sum(-2)) |
| output = torch.stack(outputs, dim=1) |
| output = output * torch.rsqrt(output.square().mean(-1, keepdim=True) + c.norm_eps) |
| gate = self.g_out(self.g_in(x)).float().reshape(shape).sigmoid() |
| output = output * self.o_norm.float() * gate |
| output = self.o_proj(output.flatten(-2).to(x.dtype)) |
| return output.masked_fill(~valid[..., None], 0), (recurrent, *histories) |
|
|
|
|
| class SwiGLU(nn.Module): |
| def __init__(self, c): |
| super().__init__() |
| self.gate_proj = nn.Linear(c.hidden_size, c.intermediate_size, bias=False) |
| self.up_proj = nn.Linear(c.hidden_size, c.intermediate_size, bias=False) |
| self.down_proj = nn.Linear(c.intermediate_size, c.hidden_size, bias=False) |
|
|
| def forward(self, x): |
| return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x)) |
|
|
|
|
| class WeevilLayer(nn.Module): |
| def __init__(self, c): |
| super().__init__() |
| self.norm = ZCRMSNorm(c.hidden_size, c.norm_eps) |
| self.mixer = GDN2Mixer(c) |
| self.mlp = SwiGLU(c) |
|
|
| def forward(self, x, valid, state): |
| z = self.norm(x).to(self.mixer.o_proj.weight.dtype) |
| mixed, state = self.mixer(z, valid, state) |
| x = x.float() + mixed.float() + self.mlp(z).float() |
| return x.masked_fill(~valid[..., None], 0), state |
|
|
|
|
| class WeevilForCausalLM(PreTrainedModel, GenerationMixin): |
| config_class = WeevilConfig |
| base_model_prefix = "weevil" |
| main_input_name = "input_ids" |
| _no_split_modules: ClassVar[list[str]] = ["WeevilLayer"] |
|
|
| def __init__(self, config): |
| super().__init__(config) |
| self.embedding = TiedEmbedding(config) |
| self.decoder = nn.ModuleDict({"layers": nn.ModuleList([ |
| WeevilLayer(config) for _ in range(config.num_hidden_layers) |
| ])}) |
| self.final_norm = ZCRMSNorm(config.hidden_size, config.norm_eps) |
| self.post_init() |
|
|
| @classmethod |
| def _supports_default_dynamic_cache(cls): |
| return False |
|
|
| def get_input_embeddings(self): |
| return self.embedding.embed |
|
|
| def set_input_embeddings(self, value): |
| self.embedding.embed = value |
|
|
| def prepare_inputs_for_generation(self, input_ids, past_key_values=None, |
| attention_mask=None, use_cache=True, **kwargs): |
| if past_key_values is not None and past_key_values.get_seq_length(): |
| input_ids = input_ids[:, past_key_values.get_seq_length():] |
| return {"input_ids": input_ids, "attention_mask": attention_mask, |
| "past_key_values": past_key_values, "use_cache": use_cache} |
|
|
| def forward(self, input_ids=None, attention_mask=None, past_key_values=None, |
| inputs_embeds=None, labels=None, use_cache=None, return_dict=None, |
| token_type_ids=None, **kwargs): |
| if (input_ids is None) == (inputs_embeds is None): |
| raise ValueError("Provide exactly one of input_ids or inputs_embeds") |
| x = self.embedding(input_ids) if inputs_embeds is None else inputs_embeds |
| batch, length = x.shape[:2] |
| valid = (torch.ones(batch, length, dtype=torch.bool, device=x.device) |
| if attention_mask is None else attention_mask[:, -length:].to(device=x.device, dtype=torch.bool)) |
| x = x.float().masked_fill(~valid[..., None], 0) |
| use_cache = self.config.use_cache if use_cache is None else use_cache |
| if past_key_values is not None and not isinstance(past_key_values, WeevilCache): |
| raise ValueError("Weevil requires a WeevilCache, not an attention KV cache") |
| states = [] |
| for i, layer in enumerate(self.decoder["layers"]): |
| state = past_key_values.states[i] if past_key_values is not None and past_key_values.states else None |
| x, state = layer(x, valid, state) |
| states.append(state) |
| logits = self.embedding.token_logits(self.final_norm(x)) |
| logits[..., self.config.pad_token_id] = -1e9 |
| logits = logits.masked_fill(~valid[..., None], 0) |
| loss = None |
| if labels is not None: |
| targets = labels[:, 1:].to(logits.device).clone() |
| targets.masked_fill_(~valid[:, 1:], -100) |
| loss = F.cross_entropy(logits[:, :-1].reshape(-1, self.config.vocab_size), |
| targets.reshape(-1), ignore_index=-100) |
| seen = 0 if past_key_values is None else past_key_values.get_seq_length() |
| cache = WeevilCache(states, seen + length) if use_cache else None |
| if return_dict is False: |
| result = (logits, cache) if use_cache else (logits,) |
| return (loss, *result) if loss is not None else result |
| return CausalLMOutputWithPast(loss=loss, logits=logits, past_key_values=cache) |
|
|
|
|
| WeevilForCausalLM.register_for_auto_class("AutoModelForCausalLM") |
|
|
