Fix KV cache handling in Attention layer

modeling_nebula.py

@@ -12,9 +12,17 @@ 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, 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
+        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
         super().__init__(**kwargs)
 
 class RMSNorm(nn.Module):
@@ -40,11 +48,9 @@ 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[1]
+        seq_len = x.shape[-2] # Use -2 for sequence length dimension
         cos = self.cos_cached[start_pos : start_pos + seq_len]
         sin = self.sin_cached[start_pos : start_pos + seq_len]
-        cos = cos.unsqueeze(0).unsqueeze(2)
-        sin = sin.unsqueeze(0).unsqueeze(2)
         x1 = x[..., : self.dim // 2]
         x2 = x[..., self.dim // 2 :]
         rotated_x1 = x1 * cos - x2 * sin
@@ -75,10 +81,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) -> torch.Tensor:
-        bs, seq_len_kv, n_kv_heads, head_dim = x.shape
-        if self.n_rep == 1: return x
-        return x.unsqueeze(3).expand(bs, seq_len_kv, n_kv_heads, self.n_rep, head_dim).reshape(bs, seq_len_kv, self.n_heads, head_dim)
+    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 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
@@ -93,9 +99,8 @@ 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.repeat_kv(xv)
-        is_causal = False if use_cache and past_key_values is not None else True
-        output = F.scaled_dot_product_attention(xq, xk_rep, xv_rep, attn_mask=attention_mask, is_causal=is_causal)
+        xk_rep, xv_rep = self.repeat_kv(xk, self.n_rep), self.repeat_kv(xv, self.n_rep)
+        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
 
@@ -130,18 +135,21 @@ class NebulaForCausalLM(PreTrainedModel, GenerationMixin):
     def _init_weights(self, module):
         if isinstance(module, (nn.Linear, nn.Embedding)): torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
         if hasattr(module, 'is_residual_output'): torch.nn.init.normal_(module.weight, mean=0.0, std=(0.02 / math.sqrt(2 * self.config.n_layers)))
-    def forward(self, input_ids: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, past_key_values: Optional[list] = None, use_cache: Optional[bool] = None, labels: Optional[torch.Tensor] = None, **kwargs) -> CausalLMOutputWithPast:
+    def forward(self, input_ids: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, past_key_values: Optional[Tuple[Tuple[torch.Tensor, torch.Tensor], ...]] = None, use_cache: Optional[bool] = None, labels: Optional[torch.Tensor] = None, **kwargs) -> CausalLMOutputWithPast:
         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 = [] if use_cache else None
+        present_key_values_list = [] if use_cache else None
         for i, layer in enumerate(self.model.layers):
             past_kv = past_key_values[i] if past_key_values is not None else None
             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 is not None: present_key_values.append(present_kv)
+            if use_cache and present_key_values_list is not None:
+                present_key_values_list.append(present_kv)
         logits = self.model.output(self.model.norm(x))
         loss = None
-        if labels is not None: loss = nn.CrossEntropyLoss()(logits.view(-1, self.config.vocab_size), labels.view(-1))
-        return CausalLMOutputWithPast(loss=loss, logits=logits, past_key_values=tuple(present_key_values) if present_key_values else None)
-    def prepare_inputs_for_generation(self, input_ids: torch.Tensor, past_key_values: Optional[list] = None, attention_mask: Optional[torch.Tensor] = None, **kwargs) -> Dict[str, Any]:
-        if past_key_values: input_ids = input_ids[:, -1:]
+        if labels is not None:
+            loss = nn.CrossEntropyLoss()(logits.view(-1, self.config.vocab_size), labels.view(-1))
+        return CausalLMOutputWithPast(loss=loss, logits=logits, past_key_values=tuple(present_key_values_list) if present_key_values_list else None)
+    def prepare_inputs_for_generation(self, input_ids: torch.Tensor, past_key_values: Optional[Tuple[Tuple[torch.Tensor, torch.Tensor], ...]] = None, attention_mask: Optional[torch.Tensor] = None, **kwargs) -> Dict[str, Any]:
+        if past_key_values:
+            input_ids = input_ids[:, -1:]
         return {"input_ids": input_ids, "past_key_values": past_key_values, "use_cache": kwargs.get("use_cache", True), "attention_mask": attention_mask}