modeling_slayer_balanced.py

"""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,
        )