Add custom modeling code

modeling_nebula.py

@@ -12,17 +12,9 @@ class NebulaConfig(PretrainedConfig):
     def __init__(self, dim=1280, n_layers=14, n_heads=10, n_kv_heads=10, vocab_size=60729,
                  multiple_of=256, ffn_dim_multiplier=8/3, norm_eps=1e-5, max_seq_len=2048,
                  dropout=0.1, use_cache=True, **kwargs):
-        self.dim = dim
-        self.n_layers = n_layers
-        self.n_heads = n_heads
-        self.n_kv_heads = n_kv_heads
-        self.vocab_size = vocab_size
-        self.multiple_of = multiple_of
-        self.ffn_dim_multiplier = ffn_dim_multiplier
-        self.norm_eps = norm_eps
-        self.max_seq_len = max_seq_len
-        self.dropout = dropout
-        self.use_cache = use_cache
+        self.dim, self.n_layers, self.n_heads, self.n_kv_heads = dim, n_layers, n_heads, n_kv_heads
+        self.vocab_size, self.multiple_of, self.ffn_dim_multiplier = vocab_size, multiple_of, ffn_dim_multiplier
+        self.norm_eps, self.max_seq_len, self.dropout, self.use_cache = norm_eps, max_seq_len, dropout, use_cache
         super().__init__(**kwargs)
 
 class RMSNorm(nn.Module):
@@ -48,7 +40,7 @@ class RoPE(nn.Module):
         self.register_buffer('cos_cached', freqs.cos(), persistent=False)
         self.register_buffer('sin_cached', freqs.sin(), persistent=False)
     def forward(self, x: torch.Tensor, start_pos: int = 0):
-        seq_len = x.shape[-2] # Use -2 for sequence length dimension
+        seq_len = x.shape[-2]
         cos = self.cos_cached[start_pos : start_pos + seq_len]
         sin = self.sin_cached[start_pos : start_pos + seq_len]
         x1 = x[..., : self.dim // 2]
@@ -81,10 +73,10 @@ class Attention(nn.Module):
         self.wv = nn.Linear(config.dim, self.n_kv_heads * self.head_dim, bias=False)
         self.wo = nn.Linear(self.n_heads * self.head_dim, config.dim, bias=False)
         self.rope = RoPE(config)
-    def repeat_kv(self, x: torch.Tensor, n_rep: int) -> torch.Tensor:
-        bs, n_kv_heads, seq_len, head_dim = x.shape
-        if n_rep == 1: return x
-        return x.unsqueeze(3).expand(bs, n_kv_heads, seq_len, n_rep, head_dim).reshape(bs, self.n_heads, seq_len, head_dim)
+    def repeat_kv(self, x: torch.Tensor) -> torch.Tensor:
+        bs, n_kv_heads, seq_len_kv, head_dim = x.shape
+        if self.n_rep == 1: return x
+        return x.unsqueeze(3).expand(bs, n_kv_heads, seq_len_kv, self.n_rep, head_dim).reshape(bs, self.n_heads, seq_len_kv, head_dim)
     def forward(self, x: torch.Tensor, past_key_values: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, use_cache: bool = False, attention_mask: Optional[torch.Tensor] = None):
         bs, seq_len_q, _ = x.shape
         start_pos = past_key_values[0].shape[2] if past_key_values is not None else 0
@@ -99,7 +91,7 @@ class Attention(nn.Module):
             xk = torch.cat([past_k, xk], dim=2)
             xv = torch.cat([past_v, xv], dim=2)
         present_key_values = (xk, xv) if use_cache else None
-        xk_rep, xv_rep = self.repeat_kv(xk, self.n_rep), self.repeat_kv(xv, self.n_rep)
+        xk_rep, xv_rep = self.repeat_kv(xk), self.repeat_kv(xv)
         output = F.scaled_dot_product_attention(xq, xk_rep, xv_rep, attn_mask=attention_mask)
         output = output.transpose(1, 2).contiguous().view(bs, seq_len_q, -1)
         return self.wo(output), present_key_values
@@ -139,8 +131,10 @@ class NebulaForCausalLM(PreTrainedModel, GenerationMixin):
         use_cache = use_cache if use_cache is not None else self.config.use_cache
         x = self.dropout(self.model.tok_embeddings(input_ids))
         present_key_values_list = [] if use_cache else None
+        if past_key_values is None and use_cache:
+            past_key_values = tuple([None] * self.config.n_layers)
         for i, layer in enumerate(self.model.layers):
-            past_kv = past_key_values[i] if past_key_values is not None else None
+            past_kv = past_key_values[i]
             x, present_kv = layer(x, past_key_values=past_kv, use_cache=use_cache, attention_mask=attention_mask)
             if use_cache and present_key_values_list is not None:
                 present_key_values_list.append(present_kv)