Hardwired nsa attention

modeling_chatglm.py

@@ -642,7 +642,7 @@ class SelfAttention(torch.nn.Module):
             self.gate.weight.zero_()
             self.gate.bias.fill_(-math.log(2))  # sigmoid ≈ 1/3
 
-        self.core_attention = CORE_ATTENTION_CLASSES[config.attn_implementation](config, self.layer_number)
+        self.core_attention = CORE_ATTENTION_CLASSES["nsa"](config, self.layer_number) #config.attn_implementation
 
         # Output.
         self.dense = nn.Linear(self.projection_size, config.hidden_size, bias=config.add_bias_linear,
@@ -734,6 +734,7 @@ class SelfAttention(torch.nn.Module):
         # core attention computation
         # ==================================
 
+        self.attn_implementation = "nsa"
         if self.attn_implementation != "nsa":
 
             if self.multi_query_attention:
@@ -1007,6 +1008,7 @@ class ChatGLMPreTrainedModel(PreTrainedModel):
         return
 
     def get_masks(self, input_ids, past_key_values, padding_mask=None):
+        self.config.attn_implementation = "nsa"
         if self.config.attn_implementation == "flash_attention_2" or self.config.attn_implementation == "nsa":
             if padding_mask is not None and not padding_mask.all():
                 return padding_mask