updates for error: . in config and main file ( no trainable param)

modeling_spin.py

@@ -1,13 +1,13 @@
-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 import PreTrainedModel, GenerationMixin
 from transformers.modeling_outputs import CausalLMOutputWithPast
+from .configuration_spin import SpinConfig
 
-class RMSNorm(nn.Module):
 
+class RMSNorm(nn.Module):
     def __init__(self, dim: int, eps: float = 1e-5):
         super().__init__()
         self.eps = eps
@@ -17,28 +17,18 @@ class RMSNorm(nn.Module):
         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)
-    )
+    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,
-):
-    
+    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)
 
@@ -51,8 +41,8 @@ def apply_rotary_emb(
     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):
 
+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)
@@ -62,8 +52,8 @@ class SwiGLU(nn.Module):
     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):
 
+class CausalSelfAttention(nn.Module):
     def __init__(self, config: SpinConfig):
         super().__init__()
         self.n_heads = config.n_heads
@@ -74,30 +64,17 @@ class CausalSelfAttention(nn.Module):
         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)
+        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: 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]]:
+    def forward(self, x, freqs_cos, freqs_sin, kv_cache=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 = q.transpose(1, 2)
-        k = k.transpose(1, 2)
-        v = v.transpose(1, 2)
+        q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)
 
         if kv_cache is not None:
             prev_k, prev_v = kv_cache
@@ -105,19 +82,16 @@ class CausalSelfAttention(nn.Module):
             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]
+            scores = scores + self.causal_mask[:T, : k.size(2)]
 
         attn_weights = F.softmax(scores, dim=-1)
-        out = attn_weights @ v
-        out = out.transpose(1, 2).contiguous().view(B, T, C)
+        out = (attn_weights @ v).transpose(1, 2).contiguous().view(B, T, C)
         return self.out_proj(out), new_kv_cache
-        
-class TransformerBlock(nn.Module):
 
+
+class TransformerBlock(nn.Module):
     def __init__(self, config: SpinConfig):
         super().__init__()
         self.attn_norm = RMSNorm(config.d_model, eps=config.norm_eps)
@@ -125,64 +99,69 @@ class TransformerBlock(nn.Module):
         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
-        )
+    def forward(self, x, freqs_cos, freqs_sin, kv_cache=None):
+        attn_out, next_kv = self.attn(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):
+class SpinForCausalLM(PreTrainedModel, GenerationMixin):
     config_class = SpinConfig
+    _tied_weights_keys = ["lm_head.weight"]
 
     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.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
+        # Weight tying
+        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)
-        self.register_buffer("freqs_sin", freqs_sin)
+        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,
+        input_ids: torch.Tensor = None,
         labels: torch.Tensor = None,
-        kv_caches=None,
-        start_pos: int = 0,
+        past_key_values: list = None,
+        use_cache: bool = False,
         return_dict: bool = True,
+        **kwargs,
     ):
         B, T = input_ids.shape
         x = self.tok_embeddings(input_ids)
 
+        start_pos = 0
+        if past_key_values is not None 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]
 
         new_kv_caches = []
         for i, layer in enumerate(self.layers):
-            cache_i = kv_caches[i] if kv_caches is not None else None
+            cache_i = past_key_values[i] if past_key_values is not None else None
             x, new_cache = layer(x, freqs_cos, freqs_sin, kv_cache=cache_i)
             new_kv_caches.append(new_cache)
 
@@ -191,17 +170,23 @@ class SpinForCausalLM(PreTrainedModel):
 
         loss = None
         if labels is not None:
-            loss = F.cross_entropy(
-                logits.view(-1, self.config.vocab_size),
-                labels.view(-1),
-                ignore_index=-100,
-            )
+            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
+            loss=loss,
+            logits=logits,
+            past_key_values=new_kv_caches if use_cache else None,
         )
 
+    def prepare_inputs_for_generation(self, input_ids, past_key_values=None, **kwargs):
+        if past_key_values is not None and len(past_key_values) > 0 and past_key_values[0] is not None:
+            input_ids = input_ids[:, -1:]
+        return {
+            "input_ids": input_ids,
+            "past_key_values": past_key_values,
+            "use_cache": True,
+        }