Update default configs for 340M models

Author: yzhangcs

configs/nsa_340M.json

@@ -9,8 +9,9 @@
     "initializer_range": 0.006,
     "max_position_embeddings": 8192,
     "model_type": "nsa",
-    "num_heads": 64,
-    "num_kv_heads": 4,
+    "num_heads": 32,
+    "num_kv_heads": 2,
+    "head_dim": 64,
     "block_size": 64,
     "block_counts": 16,
     "window_size": 512,
@@ -19,5 +20,4 @@
     "tie_word_embeddings": false,
     "use_cache": true,
     "vocab_size": 32000
-}
-
+}

native_sparse_attention/configuration_nsa.py

@@ -16,7 +16,7 @@ def __init__(
         num_hidden_layers: int = 24,
         num_heads: int = 64,
         num_kv_heads: int = 4,
-        head_dim: int = 64,
+        head_dim: int = 32,
         qkv_bias: bool = False,
         block_size: int = 64,
         block_counts: Optional[int] = 16,

native_sparse_attention/modeling_nsa.py

@@ -70,7 +70,7 @@ def __init__(
         self.k_proj = nn.Linear(self.hidden_size, self.kv_dim, bias=self.qkv_bias)
         self.v_proj = nn.Linear(self.hidden_size, self.kv_dim, bias=self.qkv_bias)
         self.g_proj = nn.Linear(self.hidden_size, self.num_heads * 3, bias=False)
-        self.o_proj = nn.Linear(self.kv_dim, self.hidden_size, bias=False)
+        self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False)
 
         self.rotary = RotaryEmbedding(dim=self.head_dim, base=self.rope_theta)