Slayer149-ARC: ARC-E 66.58, BLiMP 78.38, wiki ppl 1.907 — overall 81.51

modeling_slayer_balanced.py

@@ -0,0 +1,190 @@
+"""HF-compatible inference wrapper for Slayer's trained BalancedLM weights."""
+import math
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+from transformers import PretrainedConfig, PreTrainedModel
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+
+class SlayerBalancedConfig(PretrainedConfig):
+    model_type = "slayer_balanced"
+
+    def __init__(self, vocab_size=24576, hidden_size=576, intermediate_size=2496,
+                 num_hidden_layers=26, num_attention_heads=9, num_key_value_heads=3,
+                 head_dim=64, max_position_embeddings=2048, rms_norm_eps=1e-6,
+                 rope_parameters=None, tie_word_embeddings=True, **kwargs):
+        super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.num_key_value_heads = num_key_value_heads
+        self.head_dim = head_dim
+        self.max_position_embeddings = max_position_embeddings
+        self.rms_norm_eps = rms_norm_eps
+        self.rope_parameters = rope_parameters or {"rope_theta": 10000.0, "rope_type": "default"}
+        self.use_cache = False
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, width, eps):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(width))
+        self.eps = eps
+
+    def forward(self, x):
+        dtype = x.dtype
+        x = x.float()
+        return self.weight * (x * torch.rsqrt(x.square().mean(-1, keepdim=True) + self.eps)).to(dtype)
+
+
+class ValueProjection(nn.Linear):
+    def __init__(self, width, out, first):
+        super().__init__(width, out, bias=False)
+        if not first:
+            self.vr_lambda = nn.Parameter(torch.zeros(1))
+
+
+class SlayerAttention(nn.Module):
+    def __init__(self, c, index):
+        super().__init__()
+        h = c.hidden_size
+        self.heads, self.kv, self.d = c.num_attention_heads, c.num_key_value_heads, c.head_dim
+        self.first = index == 0
+        self.q_proj = nn.Linear(h, self.heads * self.d, bias=False)
+        self.k_proj = nn.Linear(h, self.kv * self.d, bias=False)
+        self.v_proj = ValueProjection(h, self.kv * self.d, self.first)
+        self.o_proj = nn.Linear(self.heads * self.d, h, bias=False)
+        self.q_norm, self.k_norm = RMSNorm(self.d, c.rms_norm_eps), RMSNorm(self.d, c.rms_norm_eps)
+
+    def forward(self, x, cos, sin, v0, attention_mask=None, output_attentions=False):
+        b, t, _ = x.shape
+        q = self.q_norm(self.q_proj(x).view(b, t, self.heads, self.d)).transpose(1, 2)
+        k = self.k_norm(self.k_proj(x).view(b, t, self.kv, self.d)).transpose(1, 2)
+        v = self.v_proj(x)
+        if self.first:
+            v0 = v
+        else:
+            v = v + self.v_proj.vr_lambda * v0
+        v = v.view(b, t, self.kv, self.d).transpose(1, 2)
+        if self.heads != self.kv:
+            rep = self.heads // self.kv
+            k = k.repeat_interleave(rep, dim=1)
+            v = v.repeat_interleave(rep, dim=1)
+
+        def rotate(z):
+            return torch.cat((-z[..., self.d // 2:], z[..., :self.d // 2]), dim=-1)
+
+        q, k = q * cos + rotate(q) * sin, k * cos + rotate(k) * sin
+        probs = None
+        if output_attentions:
+            scores = (q @ k.transpose(-2, -1)) / math.sqrt(self.d)
+            causal = torch.ones((t, t), dtype=torch.bool, device=x.device).tril()
+            scores = scores.masked_fill(~causal, torch.finfo(scores.dtype).min)
+            if attention_mask is not None:
+                scores = scores.masked_fill(~attention_mask[:, None, None, :].bool(), torch.finfo(scores.dtype).min)
+            probs = scores.softmax(dim=-1)
+            out = probs @ v
+        else:
+            if attention_mask is None:
+                out = F.scaled_dot_product_attention(q, k, v, is_causal=True)
+            else:
+                causal = torch.ones((t, t), dtype=torch.bool, device=x.device).tril()
+                allowed = causal[None, None] & attention_mask[:, None, None, :].bool()
+                out = F.scaled_dot_product_attention(q, k, v, attn_mask=allowed, is_causal=False)
+        out = self.o_proj(out.transpose(1, 2).contiguous().view(b, t, -1))
+        return out, v0, probs
+
+
+class SlayerMLP(nn.Module):
+    def __init__(self, c):
+        super().__init__()
+        h, f = c.hidden_size, c.intermediate_size
+        self.gate_proj = nn.Linear(h, f, bias=False)
+        self.up_proj = nn.Linear(h, f, bias=False)
+        self.down_proj = nn.Linear(f, h, bias=False)
+
+    def forward(self, x):
+        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+class SlayerBlock(nn.Module):
+    def __init__(self, c, index):
+        super().__init__()
+        self.self_attn = SlayerAttention(c, index)
+        self.mlp = SlayerMLP(c)
+        self.input_layernorm = RMSNorm(c.hidden_size, c.rms_norm_eps)
+        self.post_attention_layernorm = RMSNorm(c.hidden_size, c.rms_norm_eps)
+
+    def forward(self, x, cos, sin, v0, mask, output_attentions):
+        attn, v0, probs = self.self_attn(self.input_layernorm(x), cos, sin, v0, mask, output_attentions)
+        x = x + attn
+        return x + self.mlp(self.post_attention_layernorm(x)), v0, probs
+
+
+class SlayerBackbone(nn.Module):
+    def __init__(self, c):
+        super().__init__()
+        self.embed_tokens = nn.Embedding(c.vocab_size, c.hidden_size)
+        self.layers = nn.ModuleList(SlayerBlock(c, i) for i in range(c.num_hidden_layers))
+        self.norm = RMSNorm(c.hidden_size, c.rms_norm_eps)
+
+
+class SlayerForCausalLM(PreTrainedModel):
+    config_class = SlayerBalancedConfig
+    base_model_prefix = "model"
+    _no_split_modules = ["SlayerBlock"]
+    _supports_sdpa = True
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = SlayerBackbone(config)
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.model.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.model.embed_tokens = value
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, value):
+        self.lm_head = value
+
+    def forward(self, input_ids=None, attention_mask=None, position_ids=None, inputs_embeds=None, labels=None,
+                output_attentions=None, use_cache=None, return_dict=None, **kwargs):
+        if inputs_embeds is None:
+            inputs_embeds = self.model.embed_tokens(input_ids)
+        _, t, _ = inputs_embeds.shape
+        d = self.config.head_dim
+        theta = self.config.rope_parameters.get("rope_theta", 10000.0)
+        inv = 1.0 / theta ** (torch.arange(0, d, 2, device=inputs_embeds.device, dtype=torch.float32) / d)
+        if position_ids is None:
+            if attention_mask is None:
+                position_ids = torch.arange(t, device=inputs_embeds.device)[None, :].expand(inputs_embeds.shape[0], -1)
+            else:
+                position_ids = (attention_mask.long().cumsum(-1) - 1).clamp_min(0)
+        phase = position_ids.to(torch.float32).unsqueeze(-1) * inv
+        phase = torch.cat((phase, phase), dim=-1)
+        cos = phase.cos()[:, None].to(inputs_embeds.dtype)
+        sin = phase.sin()[:, None].to(inputs_embeds.dtype)
+        x, v0 = inputs_embeds, None
+        attentions = [] if output_attentions else None
+        for layer in self.model.layers:
+            x, v0, attn = layer(x, cos, sin, v0, attention_mask, bool(output_attentions))
+            if attentions is not None:
+                attentions.append(attn)
+        logits = self.lm_head(self.model.norm(x))
+        loss = None
+        if labels is not None:
+            loss = F.cross_entropy(logits[:, :-1].float().reshape(-1, logits.size(-1)), labels[:, 1:].reshape(-1))
+        return CausalLMOutputWithPast(
+            loss=loss, logits=logits, past_key_values=None,
+            attentions=tuple(attentions) if attentions is not None else None,
+        )