fix split error for attention

modeling_llama.py

@@ -456,7 +456,7 @@ class LlamaAttention(nn.Module):
             bsz, q_len, self.num_heads + 2 * self.num_key_value_heads, self.head_dim
         ).transpose(1, 2)
         query_states, key_states, value_states = torch.split(
-            qkv_states, [self.num_heads, self.num_heads + self.num_key_value_heads], dim=1
+            qkv_states, [self.num_heads, self.num_key_value_heads, self.num_key_value_heads], dim=1
         )
 
         if position_embeddings is None: