Update modeling_llama_nsa.py

modeling_llama_nsa.py

@@ -279,33 +279,35 @@ class LlamaNSAAttention(nn.Module):
         attention_interface: Callable = eager_attention_forward
         if self.config._attn_implementation != "eager":
             attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
-        
-        #attn_output_mha, attn_weights = attention_interface(
-        #        self,
-        #        query_states.transpose(1,2),
-        #        key_states.transpose(1,2),
-        #        value_states.transpose(1,2),
-        #        attention_mask,
-        #        dropout=0.0 if not self.training else self.attention_dropout,
-        #        scaling=self.scaling,
-        #        **kwargs,
-        #)
-        #attn_output_mha = attn_output_mha * g_slc.unsqueeze(-1) # also gated
 
-        # new for NSA
-        attn_output, _ = parallel_nsa(
-            q=query_states,
-            k=key_states,
-            v=value_states,
-            g_cmp=0,
-            g_slc=g_slc,
-            g_swa=0,
-            block_size=self.config.block_size,
-            block_counts=self.config.block_counts,
-            window_size=self.config.window_size,
-            head_first=False,
-        )
-        attn_weights = None
+        if not self.training and do_mha:
+            attn_output_mha, attn_weights = attention_interface(
+                    self,
+                    query_states.transpose(1,2),
+                    key_states.transpose(1,2),
+                    value_states.transpose(1,2),
+                    attention_mask,
+                    dropout=0.0 if not self.training else self.attention_dropout,
+                    scaling=self.scaling,
+                    **kwargs,
+            )
+            attn_output_mha = attn_output_mha * g_slc.unsqueeze(-1) # also gated
+
+        if self.training or not do_mha:
+            # new for NSA
+            attn_output, _ = parallel_nsa(
+                q=query_states,
+                k=key_states,
+                v=value_states,
+                g_cmp=0,
+                g_slc=g_slc,
+                g_swa=0,
+                block_size=self.config.block_size,
+                block_counts=self.config.block_counts,
+                window_size=self.config.window_size,
+                head_first=False,
+            )
+            attn_weights = None
         
         sa_loss = 0#torch.nn.SmoothL1Loss()(attn_output_mha, attn_output.detach()) + torch.nn.SmoothL1Loss()(attn_output_mha.detach(), attn_output)
         
@@ -313,10 +315,19 @@ class LlamaNSAAttention(nn.Module):
         #    attn_output_mha = attn_output_mha.reshape(*input_shape, -1).contiguous()
         #    attn_output_mha = self.o_proj(attn_output_mha)
         #    return attn_output_mha, attn_weights, sa_loss
-        if True:
+        if self.training:
             attn_output = attn_output.reshape(*input_shape, -1).contiguous()
             attn_output = self.o_proj(attn_output)
             return attn_output, attn_weights, sa_loss
+        else:
+            if not do_mha:
+                attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+                attn_output = self.o_proj(attn_output)
+                return attn_output, attn_weights, sa_loss
+            else:
+                attn_output_mha = attn_output_mha.reshape(*input_shape, -1).contiguous()
+                attn_output_mha = self.o_proj(attn_output_mha)
+                return attn_output_mha, attn_weights, sa_loss 
 
 
 class LlamaNSADecoderLayer(GradientCheckpointingLayer):
@@ -434,7 +445,12 @@ class LlamaNSAModel(LlamaNSAPreTrainedModel):
         use_cache: Optional[bool] = None,
         **kwargs: Unpack[TransformersKwargs],
     ) -> BaseModelOutputWithPast:
-        do_mha = False #random.random() > 0.5
+        if self.training:
+            do_mha = False
+        else:
+            if self.config.inference_mode not in ["sparse", "full"]:
+                raise ValueError
+            do_mha = False if self.config.inference_mode == "sparse" else True
            
         if (input_ids is None) ^ (inputs_embeds is not None):
             raise ValueError("You must specify exactly one of input_ids or inputs_embeds")