Add Weevil Transformers model, tokenizer and weights

README.md

@@ -0,0 +1,29 @@
+---
+library_name: transformers
+pipeline_tag: text-generation
+tags:
+- weevil
+- gdn2
+- custom_code
+---
+
+# Weevil-10M
+
+Pretrained causal language model. 9,894,464 parameters, eight parallel GDN2/SwiGLU blocks,
+and a shared 4,096-token BPE tokenizer. Uses the shared Weevil Transformers implementation.
+
+```python
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+repo = "ELBennett/Weevil-10M"
+tokenizer = AutoTokenizer.from_pretrained(repo, trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(repo, trust_remote_code=True)
+inputs = tokenizer("The purpose of this article is", return_tensors="pt")
+output = model.generate(**inputs, max_new_tokens=64, do_sample=True,
+                        temperature=0.8, top_k=40)
+print(tokenizer.decode(output[0], skip_special_tokens=True))
+```
+
+Requires PyTorch and Transformers (`pip install torch 'transformers>=4.57,<6'`).
+Weights are safetensors; generation uses fixed-size recurrent/convolution caches.
+Source checkpoint: `step-00038147`. This is a continuation model, with no chat template.

config.json

@@ -0,0 +1,26 @@
+{
+  "architectures": [
+    "WeevilForCausalLM"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_weevil.WeevilConfig",
+    "AutoModelForCausalLM": "modeling_weevil.WeevilForCausalLM"
+  },
+  "bos_token_id": 0,
+  "conv_kernel": 4,
+  "dtype": "float32",
+  "embedding_dim": 128,
+  "eos_token_id": 0,
+  "hidden_size": 256,
+  "intermediate_size": 960,
+  "model_type": "weevil",
+  "norm_eps": 1e-06,
+  "num_heads": 8,
+  "num_hidden_layers": 8,
+  "pad_token_id": 1,
+  "qkv_rank": null,
+  "tie_word_embeddings": false,
+  "transformers_version": "5.20.2",
+  "use_cache": true,
+  "vocab_size": 4096
+}

configuration_weevil.py

@@ -0,0 +1,37 @@
+"""Hugging Face configuration shared by every Weevil checkpoint."""
+
+from transformers import PretrainedConfig
+
+
+class WeevilConfig(PretrainedConfig):
+    model_type = "weevil"
+
+    def __init__(
+        self, vocab_size=4096, hidden_size=256, embedding_dim=128,
+        intermediate_size=960, num_hidden_layers=8, num_heads=8,
+        qkv_rank=None, conv_kernel=4, norm_eps=1e-6, use_cache=True,
+        **kwargs,
+    ):
+        kwargs.setdefault("bos_token_id", 0)
+        kwargs.setdefault("eos_token_id", 0)
+        kwargs.setdefault("pad_token_id", 1)
+        kwargs.setdefault("tie_word_embeddings", False)
+        super().__init__(**kwargs)
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.embedding_dim = embedding_dim
+        self.intermediate_size = intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_heads = num_heads
+        self.qkv_rank = qkv_rank
+        self.conv_kernel = conv_kernel
+        self.norm_eps = norm_eps
+        self.use_cache = use_cache
+        if min(vocab_size, hidden_size, embedding_dim, intermediate_size,
+               num_hidden_layers, num_heads, conv_kernel) < 1 or hidden_size < num_heads:
+            raise ValueError("Weevil dimensions must be positive; heads cannot exceed hidden size")
+        if qkv_rank is not None and not 1 <= qkv_rank <= hidden_size:
+            raise ValueError("qkv_rank must be between 1 and hidden_size")
+
+
+WeevilConfig.register_for_auto_class()

export_metadata.json

@@ -0,0 +1,9 @@
+{
+  "variant": "Weevil-10M",
+  "source_checkpoint": "reports/tst-s6-8gpu-10b/step-00038147",
+  "source_model_sha256": "e30b2a0277d008bab462bd7ae9aa5b0a8022b7ce1cb5b3b6687a50dde6d4820f",
+  "source_step": 38147,
+  "parameters": 9894464,
+  "model_spec": null,
+  "weights_dtype": "float32"
+}

generation_config.json

@@ -0,0 +1,10 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 0,
+  "eos_token_id": 0,
+  "output_attentions": false,
+  "output_hidden_states": false,
+  "pad_token_id": 1,
+  "transformers_version": "5.20.2",
+  "use_cache": true
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7c0147378710e100fa483cacb0a4a2a1b1a55773b09493881259cfa399d44e95
+size 39595600

modeling_weevil.py

@@ -0,0 +1,234 @@
+"""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")

tokenizer_config.json

@@ -0,0 +1,15 @@
+{
+  "backend": "tokenizers",
+  "bos_token": "<|endoftext|>",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|endoftext|>",
+  "mask_token": "<|mask|>",
+  "model_input_names": [
+    "input_ids",
+    "attention_mask"
+  ],
+  "model_max_length": 32768,
+  "pad_token": "<|padding|>",
+  "padding_side": "left",
+  "tokenizer_class": "PreTrainedTokenizerFast"
+}