Update modeling_llama_nsa.py

modeling_llama_nsa.py

@@ -279,8 +279,9 @@ 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(
+       
+        if self.training or do_mha:
+            attn_output_mha, attn_weights = attention_interface(
                 self,
                 query_states.transpose(1,2),
                 key_states.transpose(1,2),
@@ -289,26 +290,29 @@ class LlamaNSAAttention(nn.Module):
                 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_output_mha = attn_output_mha * g_slc.unsqueeze(-1) # also gated
+
+        if self.training or not do_mha:
+            # new for Sparse Attention
+            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 = torch.nn.SmoothL1Loss()(attn_output_mha, attn_output.detach()) + torch.nn.SmoothL1Loss()(attn_output_mha.detach(), attn_output)
-        
+        if self.training:
+            sa_loss = torch.nn.SmoothL1Loss()(attn_output_mha, attn_output.detach()) + torch.nn.SmoothL1Loss()(attn_output_mha.detach(), attn_output)
+        else:
+            sa_loss = 0
         if do_mha:
             attn_output_mha = attn_output_mha.reshape(*input_shape, -1).contiguous()
             attn_output_mha = self.o_proj(attn_output_mha)
@@ -434,8 +438,13 @@ class LlamaNSAModel(LlamaNSAPreTrainedModel):
         use_cache: Optional[bool] = None,
         **kwargs: Unpack[TransformersKwargs],
     ) -> BaseModelOutputWithPast:
-        do_mha = random.random() > 0.5
-           
+        if self.training:
+            do_mha = random.random() > 0.5
+        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")
 
@@ -560,12 +569,13 @@ class LlamaNSAForCausalLM(LlamaNSAPreTrainedModel, GenerationMixin):
             #loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs)
             loss = ForCausalLMLoss(hidden_states=hidden_states[:, slice_indices, :], labels=labels, lm_head_weights=self.lm_head.weight, hidden_size=self.config.hidden_size, vocab_size=self.config.vocab_size, **kwargs)
 
-        outputs.sa_loss = outputs.sa_loss*10
         if self.training:
+            outputs.sa_loss = outputs.sa_loss*10
+            loss = loss + outputs.sa_loss
             print(f"main={loss.item():.4f}, sa={outputs.sa_loss.item():.4f}")
 
         return CausalLMOutputWithPast(
-            loss=loss + outputs.sa_loss,
+            loss=loss,
             logits=logits,
             past_key_values=outputs.past_key_values,
             hidden_states=outputs.hidden_states,
@@ -590,4 +600,4 @@ __all__ = [
     "LlamaNSAForSequenceClassification",
     "LlamaNSAForQuestionAnswering",
     "LlamaNSAForTokenClassification",
-]
+]