modeling_spin.py

import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import PreTrainedModel, GenerationMixin
from transformers.modeling_outputs import CausalLMOutputWithPast
from .configuration_spin import SpinConfig


class RMSNorm(nn.Module):
    def __init__(self, dim: int, eps: float = 1e-5):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        variance = x.pow(2).mean(-1, keepdim=True)
        return x * torch.rsqrt(variance + self.eps) * self.weight


def precompute_freqs_cis(dim: int, max_seq_len: int, theta: float = 10000.0):
    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
    t = torch.arange(max_seq_len, dtype=torch.float32)
    freqs = torch.outer(t, freqs)
    return torch.cos(freqs), torch.sin(freqs)


def apply_rotary_emb(xq, xk, freqs_cos, freqs_sin):
    xq_r, xq_i = xq.float().reshape(*xq.shape[:-1], -1, 2).unbind(-1)
    xk_r, xk_i = xk.float().reshape(*xk.shape[:-1], -1, 2).unbind(-1)

    freqs_cos = freqs_cos.unsqueeze(0).unsqueeze(2)
    freqs_sin = freqs_sin.unsqueeze(0).unsqueeze(2)

    xq_out_r = xq_r * freqs_cos - xq_i * freqs_sin
    xq_out_i = xq_r * freqs_sin + xq_i * freqs_cos
    xk_out_r = xk_r * freqs_cos - xk_i * freqs_sin
    xk_out_i = xk_r * freqs_sin + xk_i * freqs_cos

    xq_out = torch.stack([xq_out_r, xq_out_i], dim=-1).flatten(3)
    xk_out = torch.stack([xk_out_r, xk_out_i], dim=-1).flatten(3)
    return xq_out.type_as(xq), xk_out.type_as(xk)


class SwiGLU(nn.Module):
    def __init__(self, d_model: int, d_ff: int):
        super().__init__()
        self.w_gate = nn.Linear(d_model, d_ff, bias=False)
        self.w_up = nn.Linear(d_model, d_ff, bias=False)
        self.w_down = nn.Linear(d_ff, d_model, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.w_down(F.silu(self.w_gate(x)) * self.w_up(x))


class CausalSelfAttention(nn.Module):
    def __init__(self, config: SpinConfig, layer_idx: int = 0):
        super().__init__()
        self.layer_idx = layer_idx
        self.n_heads = config.n_heads
        self.head_dim = config.d_model // config.n_heads

        self.q_proj = nn.Linear(config.d_model, config.d_model, bias=False)
        self.k_proj = nn.Linear(config.d_model, config.d_model, bias=False)
        self.v_proj = nn.Linear(config.d_model, config.d_model, bias=False)
        self.out_proj = nn.Linear(config.d_model, config.d_model, bias=False)

        mask = torch.full((config.max_seq_len, config.max_seq_len), float("-inf"))
        self.register_buffer("causal_mask", torch.triu(mask, diagonal=1), persistent=False)

    def forward(self, x, freqs_cos, freqs_sin, past_key_value=None):
        B, T, C = x.shape
        q = self.q_proj(x).view(B, T, self.n_heads, self.head_dim)
        k = self.k_proj(x).view(B, T, self.n_heads, self.head_dim)
        v = self.v_proj(x).view(B, T, self.n_heads, self.head_dim)

        q, k = apply_rotary_emb(q, k, freqs_cos, freqs_sin)
        q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)

        # Standard Cache update (handles both DynamicCache and classic tuple)
        if past_key_value is not None:
            if hasattr(past_key_value, "update"):
                k, v = past_key_value.update(k, v, self.layer_idx)
                new_kv_cache = past_key_value
            elif isinstance(past_key_value, tuple):
                prev_k, prev_v = past_key_value
                k = torch.cat([prev_k, k], dim=2)
                v = torch.cat([prev_v, v], dim=2)
                new_kv_cache = (k, v)
            else:
                new_kv_cache = (k, v)
        else:
            new_kv_cache = (k, v)

        scores = (q @ k.transpose(-2, -1)) / math.sqrt(self.head_dim)
        if T > 1:
            scores = scores + self.causal_mask[:T, : k.size(2)]

        attn_weights = F.softmax(scores, dim=-1)
        out = (attn_weights @ v).transpose(1, 2).contiguous().view(B, T, C)
        return self.out_proj(out), new_kv_cache


class TransformerBlock(nn.Module):
    def __init__(self, config: SpinConfig, layer_idx: int = 0):
        super().__init__()
        self.attn_norm = RMSNorm(config.d_model, eps=config.norm_eps)
        self.attn = CausalSelfAttention(config, layer_idx=layer_idx)
        self.ffn_norm = RMSNorm(config.d_model, eps=config.norm_eps)
        self.ffn = SwiGLU(config.d_model, config.d_ff)

    def forward(self, x, freqs_cos, freqs_sin, past_key_value=None):
        attn_out, next_kv = self.attn(self.attn_norm(x), freqs_cos, freqs_sin, past_key_value=past_key_value)
        x = x + attn_out
        x = x + self.ffn(self.ffn_norm(x))
        return x, next_kv


class SpinForCausalLM(PreTrainedModel, GenerationMixin):
    config_class = SpinConfig
    _tied_weights_keys = {"lm_head.weight": "tok_embeddings.weight"}
    _supports_cache_class = True

    def __init__(self, config: SpinConfig):
        super().__init__(config)
        self.config = config
        self.tok_embeddings = nn.Embedding(config.vocab_size, config.d_model)
        self.layers = nn.ModuleList([TransformerBlock(config, layer_idx=i) for i in range(config.n_layers)])
        self.norm = RMSNorm(config.d_model, eps=config.norm_eps)
        self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)

        # Tie weights
        self.lm_head.weight = self.tok_embeddings.weight

        head_dim = config.d_model // config.n_heads
        freqs_cos, freqs_sin = precompute_freqs_cis(head_dim, config.max_seq_len)
        self.register_buffer("freqs_cos", freqs_cos, persistent=False)
        self.register_buffer("freqs_sin", freqs_sin, persistent=False)

        self.post_init()

    def get_input_embeddings(self):
        return self.tok_embeddings

    def set_input_embeddings(self, value):
        self.tok_embeddings = value

    def get_output_embeddings(self):
        return self.lm_head

    def set_output_embeddings(self, new_embeddings):
        self.lm_head = new_embeddings

    def forward(
        self,
        input_ids: torch.Tensor = None,
        attention_mask: torch.Tensor = None,
        labels: torch.Tensor = None,
        past_key_values=None,
        use_cache: bool = False,
        return_dict: bool = True,
        **kwargs,
    ):
        B, T = input_ids.shape
        x = self.tok_embeddings(input_ids)

        # Calculate start position for RoPE
        start_pos = 0
        if past_key_values is not None:
            if hasattr(past_key_values, "get_seq_length"):
                start_pos = past_key_values.get_seq_length()
            elif isinstance(past_key_values, (tuple, list)) and len(past_key_values) > 0 and past_key_values[0] is not None:
                start_pos = past_key_values[0][0].shape[2]

        freqs_cos = self.freqs_cos[start_pos : start_pos + T]
        freqs_sin = self.freqs_sin[start_pos : start_pos + T]

        legacy_kv_caches = []
        for i, layer in enumerate(self.layers):
            if hasattr(past_key_values, "update"):
                layer_cache = past_key_values
            elif isinstance(past_key_values, (tuple, list)) and len(past_key_values) > i:
                layer_cache = past_key_values[i]
            else:
                layer_cache = None

            x, new_cache = layer(x, freqs_cos, freqs_sin, past_key_value=layer_cache)
            if not hasattr(past_key_values, "update"):
                legacy_kv_caches.append(new_cache)

        x = self.norm(x)
        logits = self.lm_head(x)

        loss = None
        if labels is not None:
            loss = F.cross_entropy(logits.view(-1, self.config.vocab_size), labels.view(-1), ignore_index=-100)

        if use_cache:
            output_cache = past_key_values if hasattr(past_key_values, "update") else tuple(legacy_kv_caches)
        else:
            output_cache = None

        if not return_dict:
            return (logits, loss, output_cache)

        return CausalLMOutputWithPast(
            loss=loss,
            logits=logits,
            past_key_values=output_cache,
        )

    def prepare_inputs_for_generation(self, input_ids, past_key_values=None, attention_mask=None, **kwargs):
        past_length = 0
        if past_key_values is not None:
            if hasattr(past_key_values, "get_seq_length"):
                past_length = past_key_values.get_seq_length()
            elif isinstance(past_key_values, (tuple, list)) and len(past_key_values) > 0 and past_key_values[0] is not None:
                past_length = past_key_values[0][0].shape[2]

        if past_length > 0:
            input_ids = input_ids[:, -1:]

        return {
            "input_ids": input_ids,
            "past_key_values": past_key_values,
            "attention_mask": attention_mask,
            "use_cache": True,
        }