NSA 117M initial export

logs/logs_extra_keys.txt

@@ -1,49 +1 @@
-blocks.0.attn.gate.fc1.bias
-blocks.0.attn.gate.fc1.weight
-blocks.0.attn.gate.fc2.bias
-blocks.0.attn.gate.fc2.weight
-blocks.1.attn.gate.fc1.bias
-blocks.1.attn.gate.fc1.weight
-blocks.1.attn.gate.fc2.bias
-blocks.1.attn.gate.fc2.weight
-blocks.10.attn.gate.fc1.bias
-blocks.10.attn.gate.fc1.weight
-blocks.10.attn.gate.fc2.bias
-blocks.10.attn.gate.fc2.weight
-blocks.11.attn.gate.fc1.bias
-blocks.11.attn.gate.fc1.weight
-blocks.11.attn.gate.fc2.bias
-blocks.11.attn.gate.fc2.weight
-blocks.2.attn.gate.fc1.bias
-blocks.2.attn.gate.fc1.weight
-blocks.2.attn.gate.fc2.bias
-blocks.2.attn.gate.fc2.weight
-blocks.3.attn.gate.fc1.bias
-blocks.3.attn.gate.fc1.weight
-blocks.3.attn.gate.fc2.bias
-blocks.3.attn.gate.fc2.weight
-blocks.4.attn.gate.fc1.bias
-blocks.4.attn.gate.fc1.weight
-blocks.4.attn.gate.fc2.bias
-blocks.4.attn.gate.fc2.weight
-blocks.5.attn.gate.fc1.bias
-blocks.5.attn.gate.fc1.weight
-blocks.5.attn.gate.fc2.bias
-blocks.5.attn.gate.fc2.weight
-blocks.6.attn.gate.fc1.bias
-blocks.6.attn.gate.fc1.weight
-blocks.6.attn.gate.fc2.bias
-blocks.6.attn.gate.fc2.weight
-blocks.7.attn.gate.fc1.bias
-blocks.7.attn.gate.fc1.weight
-blocks.7.attn.gate.fc2.bias
-blocks.7.attn.gate.fc2.weight
-blocks.8.attn.gate.fc1.bias
-blocks.8.attn.gate.fc1.weight
-blocks.8.attn.gate.fc2.bias
-blocks.8.attn.gate.fc2.weight
-blocks.9.attn.gate.fc1.bias
-blocks.9.attn.gate.fc1.weight
-blocks.9.attn.gate.fc2.bias
-blocks.9.attn.gate.fc2.weight
 norm_f.weight

logs/logs_mapping.json

@@ -7,6 +7,10 @@
     "model.blocks.0.attn.W_V_cmp.weight",
     "model.blocks.0.attn.W_V_sel.weight",
     "model.blocks.0.attn.W_V_win.weight",
+    "model.blocks.0.attn.gate.fc1.bias",
+    "model.blocks.0.attn.gate.fc1.weight",
+    "model.blocks.0.attn.gate.fc2.bias",
+    "model.blocks.0.attn.gate.fc2.weight",
     "model.blocks.0.attn.out.weight",
     "model.blocks.0.mlp.fc1.weight",
     "model.blocks.0.mlp.fc2.weight",
@@ -19,6 +23,10 @@
     "model.blocks.1.attn.W_V_cmp.weight",
     "model.blocks.1.attn.W_V_sel.weight",
     "model.blocks.1.attn.W_V_win.weight",
+    "model.blocks.1.attn.gate.fc1.bias",
+    "model.blocks.1.attn.gate.fc1.weight",
+    "model.blocks.1.attn.gate.fc2.bias",
+    "model.blocks.1.attn.gate.fc2.weight",
     "model.blocks.1.attn.out.weight",
     "model.blocks.1.mlp.fc1.weight",
     "model.blocks.1.mlp.fc2.weight",
@@ -31,6 +39,10 @@
     "model.blocks.10.attn.W_V_cmp.weight",
     "model.blocks.10.attn.W_V_sel.weight",
     "model.blocks.10.attn.W_V_win.weight",
+    "model.blocks.10.attn.gate.fc1.bias",
+    "model.blocks.10.attn.gate.fc1.weight",
+    "model.blocks.10.attn.gate.fc2.bias",
+    "model.blocks.10.attn.gate.fc2.weight",
     "model.blocks.10.attn.out.weight",
     "model.blocks.10.mlp.fc1.weight",
     "model.blocks.10.mlp.fc2.weight",
@@ -43,6 +55,10 @@
     "model.blocks.11.attn.W_V_cmp.weight",
     "model.blocks.11.attn.W_V_sel.weight",
     "model.blocks.11.attn.W_V_win.weight",
+    "model.blocks.11.attn.gate.fc1.bias",
+    "model.blocks.11.attn.gate.fc1.weight",
+    "model.blocks.11.attn.gate.fc2.bias",
+    "model.blocks.11.attn.gate.fc2.weight",
     "model.blocks.11.attn.out.weight",
     "model.blocks.11.mlp.fc1.weight",
     "model.blocks.11.mlp.fc2.weight",
@@ -55,6 +71,10 @@
     "model.blocks.2.attn.W_V_cmp.weight",
     "model.blocks.2.attn.W_V_sel.weight",
     "model.blocks.2.attn.W_V_win.weight",
+    "model.blocks.2.attn.gate.fc1.bias",
+    "model.blocks.2.attn.gate.fc1.weight",
+    "model.blocks.2.attn.gate.fc2.bias",
+    "model.blocks.2.attn.gate.fc2.weight",
     "model.blocks.2.attn.out.weight",
     "model.blocks.2.mlp.fc1.weight",
     "model.blocks.2.mlp.fc2.weight",
@@ -67,6 +87,10 @@
     "model.blocks.3.attn.W_V_cmp.weight",
     "model.blocks.3.attn.W_V_sel.weight",
     "model.blocks.3.attn.W_V_win.weight",
+    "model.blocks.3.attn.gate.fc1.bias",
+    "model.blocks.3.attn.gate.fc1.weight",
+    "model.blocks.3.attn.gate.fc2.bias",
+    "model.blocks.3.attn.gate.fc2.weight",
     "model.blocks.3.attn.out.weight",
     "model.blocks.3.mlp.fc1.weight",
     "model.blocks.3.mlp.fc2.weight",
@@ -79,6 +103,10 @@
     "model.blocks.4.attn.W_V_cmp.weight",
     "model.blocks.4.attn.W_V_sel.weight",
     "model.blocks.4.attn.W_V_win.weight",
+    "model.blocks.4.attn.gate.fc1.bias",
+    "model.blocks.4.attn.gate.fc1.weight",
+    "model.blocks.4.attn.gate.fc2.bias",
+    "model.blocks.4.attn.gate.fc2.weight",
     "model.blocks.4.attn.out.weight",
     "model.blocks.4.mlp.fc1.weight",
     "model.blocks.4.mlp.fc2.weight",
@@ -91,6 +119,10 @@
     "model.blocks.5.attn.W_V_cmp.weight",
     "model.blocks.5.attn.W_V_sel.weight",
     "model.blocks.5.attn.W_V_win.weight",
+    "model.blocks.5.attn.gate.fc1.bias",
+    "model.blocks.5.attn.gate.fc1.weight",
+    "model.blocks.5.attn.gate.fc2.bias",
+    "model.blocks.5.attn.gate.fc2.weight",
     "model.blocks.5.attn.out.weight",
     "model.blocks.5.mlp.fc1.weight",
     "model.blocks.5.mlp.fc2.weight",
@@ -103,6 +135,10 @@
     "model.blocks.6.attn.W_V_cmp.weight",
     "model.blocks.6.attn.W_V_sel.weight",
     "model.blocks.6.attn.W_V_win.weight",
+    "model.blocks.6.attn.gate.fc1.bias",
+    "model.blocks.6.attn.gate.fc1.weight",
+    "model.blocks.6.attn.gate.fc2.bias",
+    "model.blocks.6.attn.gate.fc2.weight",
     "model.blocks.6.attn.out.weight",
     "model.blocks.6.mlp.fc1.weight",
     "model.blocks.6.mlp.fc2.weight",
@@ -115,6 +151,10 @@
     "model.blocks.7.attn.W_V_cmp.weight",
     "model.blocks.7.attn.W_V_sel.weight",
     "model.blocks.7.attn.W_V_win.weight",
+    "model.blocks.7.attn.gate.fc1.bias",
+    "model.blocks.7.attn.gate.fc1.weight",
+    "model.blocks.7.attn.gate.fc2.bias",
+    "model.blocks.7.attn.gate.fc2.weight",
     "model.blocks.7.attn.out.weight",
     "model.blocks.7.mlp.fc1.weight",
     "model.blocks.7.mlp.fc2.weight",
@@ -127,6 +167,10 @@
     "model.blocks.8.attn.W_V_cmp.weight",
     "model.blocks.8.attn.W_V_sel.weight",
     "model.blocks.8.attn.W_V_win.weight",
+    "model.blocks.8.attn.gate.fc1.bias",
+    "model.blocks.8.attn.gate.fc1.weight",
+    "model.blocks.8.attn.gate.fc2.bias",
+    "model.blocks.8.attn.gate.fc2.weight",
     "model.blocks.8.attn.out.weight",
     "model.blocks.8.mlp.fc1.weight",
     "model.blocks.8.mlp.fc2.weight",
@@ -139,6 +183,10 @@
     "model.blocks.9.attn.W_V_cmp.weight",
     "model.blocks.9.attn.W_V_sel.weight",
     "model.blocks.9.attn.W_V_win.weight",
+    "model.blocks.9.attn.gate.fc1.bias",
+    "model.blocks.9.attn.gate.fc1.weight",
+    "model.blocks.9.attn.gate.fc2.bias",
+    "model.blocks.9.attn.gate.fc2.weight",
     "model.blocks.9.attn.out.weight",
     "model.blocks.9.mlp.fc1.weight",
     "model.blocks.9.mlp.fc2.weight",
@@ -148,82 +196,10 @@
     "model.lm_head.weight"
   ],
   "missing": [
-    "model.blocks.0.attn.g1.weight",
-    "model.blocks.0.attn.g2.weight",
-    "model.blocks.1.attn.g1.weight",
-    "model.blocks.1.attn.g2.weight",
-    "model.blocks.10.attn.g1.weight",
-    "model.blocks.10.attn.g2.weight",
-    "model.blocks.11.attn.g1.weight",
-    "model.blocks.11.attn.g2.weight",
-    "model.blocks.2.attn.g1.weight",
-    "model.blocks.2.attn.g2.weight",
-    "model.blocks.3.attn.g1.weight",
-    "model.blocks.3.attn.g2.weight",
-    "model.blocks.4.attn.g1.weight",
-    "model.blocks.4.attn.g2.weight",
-    "model.blocks.5.attn.g1.weight",
-    "model.blocks.5.attn.g2.weight",
-    "model.blocks.6.attn.g1.weight",
-    "model.blocks.6.attn.g2.weight",
-    "model.blocks.7.attn.g1.weight",
-    "model.blocks.7.attn.g2.weight",
-    "model.blocks.8.attn.g1.weight",
-    "model.blocks.8.attn.g2.weight",
-    "model.blocks.9.attn.g1.weight",
-    "model.blocks.9.attn.g2.weight",
     "model.norm.bias",
     "model.norm.weight"
   ],
   "extra": [
-    "blocks.0.attn.gate.fc1.bias",
-    "blocks.0.attn.gate.fc1.weight",
-    "blocks.0.attn.gate.fc2.bias",
-    "blocks.0.attn.gate.fc2.weight",
-    "blocks.1.attn.gate.fc1.bias",
-    "blocks.1.attn.gate.fc1.weight",
-    "blocks.1.attn.gate.fc2.bias",
-    "blocks.1.attn.gate.fc2.weight",
-    "blocks.10.attn.gate.fc1.bias",
-    "blocks.10.attn.gate.fc1.weight",
-    "blocks.10.attn.gate.fc2.bias",
-    "blocks.10.attn.gate.fc2.weight",
-    "blocks.11.attn.gate.fc1.bias",
-    "blocks.11.attn.gate.fc1.weight",
-    "blocks.11.attn.gate.fc2.bias",
-    "blocks.11.attn.gate.fc2.weight",
-    "blocks.2.attn.gate.fc1.bias",
-    "blocks.2.attn.gate.fc1.weight",
-    "blocks.2.attn.gate.fc2.bias",
-    "blocks.2.attn.gate.fc2.weight",
-    "blocks.3.attn.gate.fc1.bias",
-    "blocks.3.attn.gate.fc1.weight",
-    "blocks.3.attn.gate.fc2.bias",
-    "blocks.3.attn.gate.fc2.weight",
-    "blocks.4.attn.gate.fc1.bias",
-    "blocks.4.attn.gate.fc1.weight",
-    "blocks.4.attn.gate.fc2.bias",
-    "blocks.4.attn.gate.fc2.weight",
-    "blocks.5.attn.gate.fc1.bias",
-    "blocks.5.attn.gate.fc1.weight",
-    "blocks.5.attn.gate.fc2.bias",
-    "blocks.5.attn.gate.fc2.weight",
-    "blocks.6.attn.gate.fc1.bias",
-    "blocks.6.attn.gate.fc1.weight",
-    "blocks.6.attn.gate.fc2.bias",
-    "blocks.6.attn.gate.fc2.weight",
-    "blocks.7.attn.gate.fc1.bias",
-    "blocks.7.attn.gate.fc1.weight",
-    "blocks.7.attn.gate.fc2.bias",
-    "blocks.7.attn.gate.fc2.weight",
-    "blocks.8.attn.gate.fc1.bias",
-    "blocks.8.attn.gate.fc1.weight",
-    "blocks.8.attn.gate.fc2.bias",
-    "blocks.8.attn.gate.fc2.weight",
-    "blocks.9.attn.gate.fc1.bias",
-    "blocks.9.attn.gate.fc1.weight",
-    "blocks.9.attn.gate.fc2.bias",
-    "blocks.9.attn.gate.fc2.weight",
     "norm_f.weight"
   ]
 }

logs/logs_missing_keys.txt

@@ -1,26 +1,2 @@
-model.blocks.0.attn.g1.weight
-model.blocks.0.attn.g2.weight
-model.blocks.1.attn.g1.weight
-model.blocks.1.attn.g2.weight
-model.blocks.10.attn.g1.weight
-model.blocks.10.attn.g2.weight
-model.blocks.11.attn.g1.weight
-model.blocks.11.attn.g2.weight
-model.blocks.2.attn.g1.weight
-model.blocks.2.attn.g2.weight
-model.blocks.3.attn.g1.weight
-model.blocks.3.attn.g2.weight
-model.blocks.4.attn.g1.weight
-model.blocks.4.attn.g2.weight
-model.blocks.5.attn.g1.weight
-model.blocks.5.attn.g2.weight
-model.blocks.6.attn.g1.weight
-model.blocks.6.attn.g2.weight
-model.blocks.7.attn.g1.weight
-model.blocks.7.attn.g2.weight
-model.blocks.8.attn.g1.weight
-model.blocks.8.attn.g2.weight
-model.blocks.9.attn.g1.weight
-model.blocks.9.attn.g2.weight
 model.norm.bias
 model.norm.weight

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:171c4c893cc17e5a5f4ff0ab2f43cc0a7f48e50b47b48b4401b0e21b7b0eacb4
-size 320203152
+oid sha256:21d3ac54cadd49cc11ea0e88d37874aa3a9391e7b47a2704b6557a0e9229640c
+size 313204736

modeling_nsa.py

@@ -9,7 +9,12 @@ from transformers.generation.utils import GenerationMixin
 from transformers.modeling_outputs import CausalLMOutput
 
 from .configuration_nsa import NSAConfig
-_HAS_NSA = False  # Embedded NSA is provided below; no external import required.
+_HAS_NSA = False
+try:
+    from .nsa.model.llama_block_nsa import LlamaBlockNSA as _VendorNSABlock
+    _HAS_NSA = True
+except Exception:
+    _VendorNSABlock = None  # type: ignore
 
 
 class RMSNorm(nn.Module):
@@ -100,9 +105,12 @@ class EmbeddedNSAAttention(nn.Module):
         self.W_V_sel = nn.Linear(dim, n_kv_groups * d_v, bias=False)
         self.W_K_win = nn.Linear(dim, n_kv_groups * d_k, bias=False)
         self.W_V_win = nn.Linear(dim, n_kv_groups * d_v, bias=False)
-        self.g1 = nn.Linear(dim, max(1, dim // 4), bias=False)
-        self.g2 = nn.Linear(max(1, dim // 4), 3, bias=False)
-        nn.init.zeros_(self.g2.weight)
+        # Gate MLP operates on per-group pooled Q with width d_k (matches training)
+        gate_hidden = max(1, d_k // 2)
+        self.gate_fc1 = nn.Linear(d_k, gate_hidden, bias=True)
+        self.gate_fc2 = nn.Linear(gate_hidden, 3, bias=True)
+        nn.init.xavier_uniform_(self.gate_fc2.weight, gain=0.1)
+        nn.init.zeros_(self.gate_fc2.bias)
         self.out = nn.Linear(n_heads * d_v, dim, bias=False)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
@@ -159,11 +167,25 @@ class EmbeddedNSAAttention(nn.Module):
         P_w = torch.nn.functional.softmax(logits_w, dim=-1)
         O_win = torch.matmul(P_w, Vw)
 
-        # Gate & mix
-        gate = self.g2(torch.nn.functional.silu(self.g1(x)))  # [B,S,3]
-        gate = torch.nn.functional.softmax(gate, dim=-1)
-        gc, gs, gw = gate[..., 0:1], gate[..., 1:2], gate[..., 2:3]
-        O = gc.unsqueeze(1) * O_cmp + gs.unsqueeze(1) * O_sel + gw.unsqueeze(1) * O_win
+        # Gate & mix: compute per-token, per-group gate from pooled Q
+        # Pool Q across heads within each kv-group
+        # Qr: [B,h,S,dk] -> reshape to [B,G,h_per_group,S,dk] then mean over h_per_group
+        G = max(1, self.n_kv_groups)
+        h_per_group = max(1, h // G)
+        Qg = Qr.view(B, G, h_per_group, S, dk).mean(dim=2)  # [B,G,S,dk]
+        Qg = Qg.permute(0, 2, 1, 3)  # [B,S,G,dk]
+        g1 = torch.nn.functional.silu(self.gate_fc1(Qg))
+        gate = torch.nn.functional.softmax(self.gate_fc2(g1), dim=-1)  # [B,S,G,3]
+        gc = gate[..., 0:1].unsqueeze(-1)  # [B,S,G,1,1]
+        gs = gate[..., 1:2].unsqueeze(-1)
+        gw = gate[..., 2:3].unsqueeze(-1)
+        # Broadcast group gates to heads within the group
+        # Reshape branch outputs to [B,S,G,h_per_group,dv]
+        Oc = O_cmp.permute(0,2,1,3).view(B, S, G, h_per_group, dv)
+        Os = O_sel.permute(0,2,1,3).view(B, S, G, h_per_group, dv)
+        Ow = O_win.permute(0,2,1,3).view(B, S, G, h_per_group, dv)
+        O = gc * Oc + gs * Os + gw * Ow
+        O = O.reshape(B, S, h, dv).permute(0, 2, 1, 3)
         O = O.transpose(1, 2).reshape(B, S, h * dv)
         return self.out(O)
 
@@ -240,7 +262,24 @@ class NSATinyLM(nn.Module):
         import os as _os
         # Allow forcing simple fallback via env for integration tests
         _force_simple = _os.getenv('NSA_REMOTE_FORCE_SIMPLE', '0').lower() in ('1','true','yes')
-        if _force_simple == False:
+        if not _force_simple and _HAS_NSA and _VendorNSABlock is not None:
+            # Prefer vendored NSA block to match training semantics and map gate weights
+            self.blocks = nn.ModuleList([
+                _VendorNSABlock(
+                    dim=self.hidden_size,
+                    n_heads=self.num_attention_heads,
+                    n_kv_groups=self.n_kv_groups,
+                    d_k=self.d_k,
+                    d_v=self.d_v,
+                    l=self.l,
+                    d=self.d,
+                    l_sel=self.l_sel,
+                    n_sel=self.n_sel,
+                    w=self.w,
+                ) for _ in range(self.num_hidden_layers)
+            ])
+        elif not _force_simple:
+            # Fallback to embedded minimal NSA if vendor import failed
             self.blocks = nn.ModuleList([
                 NSABlockRemote(
                     self.hidden_size,