attn one kv

kyegomez · Aug 28, 2023 · 72cf984 · 72cf984
1 parent c37d1cb
commit 72cf984
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 24 deletions.
diff --git a/Andromeda/configs.py b/Andromeda/configs.py
@@ -12,9 +12,8 @@
     alibi_pos_bias=True, 
     alibi_num_heads=4, 
     rotary_xpos=True,
-    attn_flash=True, 
-    # shift_tokens=1, 
-    attn_one_kv_head=True, 
+    attn_flash=True,  
+    attn_kv_heads = 2, 
     qk_norm=True, 
     attn_qk_norm=True, 
     attn_qk_norm_dim_scale=True, 
@@ -32,10 +31,8 @@
     use_abs_pos_emb=False, 
     alibi_pos_bias=True, 
     alibi_num_heads=6, 
-    rotary_xpos=True,
-    attn_flash=True, 
-    shift_tokens=1, 
-    attn_one_kv_head=True, 
+    rotary_xpos=True, 
+    attn_kv_heads = 2, 
     qk_norm=True, 
     attn_qk_norm=True, 
     attn_qk_norm_dim_scale=True, 
@@ -53,10 +50,8 @@
     use_abs_pos_emb=False, 
     alibi_pos_bias=True, 
     alibi_num_heads=8, 
-    rotary_xpos=True,
-    attn_flash=True, 
-    shift_tokens=1, 
-    attn_one_kv_head=True, 
+    rotary_xpos=True, 
+    attn_kv_heads = 2, 
     qk_norm=True, 
     attn_qk_norm=True, 
     attn_qk_norm_dim_scale=True, 
@@ -72,10 +67,8 @@
     use_abs_pos_emb=False, 
     alibi_pos_bias=True, 
     alibi_num_heads=4, 
-    rotary_xpos=True,
-    attn_flash=True, 
-    shift_tokens=1, 
-    attn_one_kv_head=True, 
+    rotary_xpos=True, 
+    attn_kv_heads = 2, 
     qk_norm=True, 
     attn_qk_norm=True, 
     attn_qk_norm_dim_scale=True, 
@@ -91,10 +84,8 @@
     use_abs_pos_emb=False, 
     alibi_pos_bias=True, 
     alibi_num_heads=4, 
-    rotary_xpos=True,
-    attn_flash=True, 
-    shift_tokens=1, 
-    attn_one_kv_head=True, 
+    rotary_xpos=True, 
+    attn_kv_heads = 2, 
     qk_norm=True, 
     attn_qk_norm=True, 
     attn_qk_norm_dim_scale=True, 
@@ -110,10 +101,8 @@
     use_abs_pos_emb=False, 
     alibi_pos_bias=True, 
     alibi_num_heads=4, 
-    rotary_xpos=True,
-    attn_flash=True, 
-    shift_tokens=1, 
-    attn_one_kv_head=True, 
+    rotary_xpos=True, 
+    attn_kv_heads = 2, 
     qk_norm=True, 
     attn_qk_norm=True, 
     attn_qk_norm_dim_scale=True, 

diff --git a/Andromeda/model.py b/Andromeda/model.py
@@ -48,7 +48,7 @@ def __init__(self,
                  alibi_num_heads=12, 
                  rotary_xpos=True,
                  attn_flash=True, 
-                attn_kv_heads = 2,
+                 attn_kv_heads = 2,
                  qk_norm=True, 
                  attn_qk_norm=True, 
                  attn_qk_norm_dim_scale=True,