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