from configuration_spin import SpinConfig import math import torch import torch.nn as nn import torch.nn.functional as F from transformers import PreTrainedModel from transformers.modeling_outputs import CausalLMOutputWithPast 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) freqs_cos = torch.cos(freqs) freqs_sin = torch.sin(freqs) return freqs_cos, freqs_sin def apply_rotary_emb( xq: torch.Tensor, xk: torch.Tensor, freqs_cos: torch.Tensor, freqs_sin: torch.Tensor, ): 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): super().__init__() 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") ) mask = torch.triu(mask, diagonal=1) self.register_buffer("causal_mask", mask) def forward( self, x: torch.Tensor, freqs_cos: torch.Tensor, freqs_sin: torch.Tensor, kv_cache: tuple[torch.Tensor, torch.Tensor] | None = None, ) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]: 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 = q.transpose(1, 2) k = k.transpose(1, 2) v = v.transpose(1, 2) if kv_cache is not None: prev_k, prev_v = kv_cache k = torch.cat([prev_k, k], dim=2) v = torch.cat([prev_v, v], dim=2) new_kv_cache = (k, v) total_k_len = k.size(2) scores = (q @ k.transpose(-2, -1)) / math.sqrt(self.head_dim) if T > 1: scores = scores + self.causal_mask[:T, :total_k_len] attn_weights = F.softmax(scores, dim=-1) out = attn_weights @ v out = out.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): super().__init__() self.attn_norm = RMSNorm(config.d_model, eps=config.norm_eps) self.attn = CausalSelfAttention(config) self.ffn_norm = RMSNorm(config.d_model, eps=config.norm_eps) self.ffn = SwiGLU(config.d_model, config.d_ff) def forward( self, x: torch.Tensor, freqs_cos: torch.Tensor, freqs_sin: torch.Tensor, kv_cache: tuple[torch.Tensor, torch.Tensor] | None = None, ) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]: norm_x = self.attn_norm(x) attn_out, next_kv = self.attn( norm_x, freqs_cos, freqs_sin, kv_cache=kv_cache ) x = x + attn_out x = x + self.ffn(self.ffn_norm(x)) return x, next_kv class SpinForCausalLM(PreTrainedModel): config_class = SpinConfig 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) for _ 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) self.tok_embeddings.weight = self.lm_head.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) self.register_buffer("freqs_sin", freqs_sin) self.post_init() def forward( self, input_ids: torch.Tensor, labels: torch.Tensor = None, kv_caches=None, start_pos: int = 0, return_dict: bool = True, ): B, T = input_ids.shape x = self.tok_embeddings(input_ids) freqs_cos = self.freqs_cos[start_pos : start_pos + T] freqs_sin = self.freqs_sin[start_pos : start_pos + T] new_kv_caches = [] for i, layer in enumerate(self.layers): cache_i = kv_caches[i] if kv_caches is not None else None x, new_cache = layer(x, freqs_cos, freqs_sin, kv_cache=cache_i) new_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 not return_dict: return (logits, loss, new_kv_caches) return CausalLMOutputWithPast( loss=loss, logits=logits, past_key_values=new_kv_caches )