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, }