Update modeling_gemmoe.py

modeling_gemmoe.py

@@ -18,12 +18,13 @@
 import math
 import warnings
 from typing import List, Optional, Tuple, Union
-import contextlib
+
 import torch
 import torch.nn.functional as F
 import torch.utils.checkpoint
 from torch import nn
 from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
+
 from transformers.activations import ACT2FN
 from transformers.cache_utils import Cache, DynamicCache, StaticCache
 from transformers.modeling_attn_mask_utils import (
@@ -305,7 +306,6 @@ class GemmoeAttention(nn.Module):
                 - The attention weights (if `output_attentions=True`).
                 - The past key-value cache (if `use_cache=True`).
         """
-
         bsz, q_len, _ = hidden_states.size()
 
         query_states = self.q_proj(hidden_states)
@@ -331,14 +331,12 @@ class GemmoeAttention(nn.Module):
 
         attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)
 
-        with torch.no_grad() if not self.training else contextlib.nullcontext():
-            attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)
-            if attention_mask is not None:
-                if cache_position is not None:
-                    causal_mask = attention_mask[:, :, cache_position, : key_states.shape[-2]]
-                else:
-                    causal_mask = attention_mask
-                attn_weights = attn_weights + causal_mask
+        if attention_mask is not None:  # no matter the length, we just slice it
+            if cache_position is not None:
+                causal_mask = attention_mask[:, :, cache_position, : key_states.shape[-2]]
+            else:
+                causal_mask = attention_mask
+            attn_weights = attn_weights + causal_mask
 
         # upcast attention to fp32
         attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
@@ -686,7 +684,6 @@ class GemmoeSparseMoeBlock(nn.Module):
 
         self.experts = nn.ModuleList([GemmoeBlockSparseTop2MLP(config) for _ in range(self.num_experts)])
 
-    @torch.jit.script
     def forward(self, hidden_states: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
         batch_size, sequence_length, hidden_dim = hidden_states.shape
         hidden_states = hidden_states.view(-1, hidden_dim)
@@ -727,7 +724,6 @@ class GemmoeDecoderLayer(nn.Module):
         self.input_layernorm = GemmoeRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
         self.post_attention_layernorm = GemmoeRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
 
-    @torch.jit.script
     def forward(
         self,
         hidden_states: torch.Tensor,
@@ -977,7 +973,7 @@ class GemmoeModel(GemmoePreTrainedModel):
 		hidden_states = inputs_embeds
 
 		# Normalize
-		scale_factor = torch.tensor(math.sqrt(self.config.hidden_size), dtype=hidden_states.dtype)
+		scale_factor = torch.tensor(math_sqrt(self.config.hidden_size), dtype=hidden_states.dtype)
 		hidden_states = hidden_states * scale_factor
 		# Decoder layers
 		all_hidden_states = () if output_hidden_states else None
@@ -990,8 +986,8 @@ class GemmoeModel(GemmoePreTrainedModel):
 				all_hidden_states += (hidden_states,)
 
 			if self.gradient_checkpointing and self.training:
-				layer_outputs = torch.utils.checkpoint.checkpoint(
-					decoder_layer, 
+				layer_outputs = self._gradient_checkpointing_func(
+					decoder_layer.__call__,
 					hidden_states,
 					causal_mask,
 					position_ids,
@@ -1204,34 +1200,19 @@ class GemmoeForCausalLM(GemmoePreTrainedModel):
         )
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
 
-        if self.training:
-            outputs = torch.utils.checkpoint.checkpoint(
-                self.model,
-                input_ids,
-                attention_mask,
-                position_ids,
-                past_key_values,
-                inputs_embeds,
-                use_cache,
-                output_attentions,
-                output_hidden_states,
-                return_dict,
-                cache_position,
-            )
-        else:
-            outputs = self.model(
-                input_ids=input_ids,
-                attention_mask=attention_mask,
-                position_ids=position_ids,
-                past_key_values=past_key_values,
-                inputs_embeds=inputs_embeds,
-                use_cache=use_cache,
-                output_attentions=output_attentions,
-                output_hidden_states=output_hidden_states,
-                output_router_logits=output_router_logits,
-                return_dict=return_dict,
-                cache_position=cache_position,
-            )
+        outputs = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            output_router_logits=output_router_logits,
+            return_dict=return_dict,
+            cache_position=cache_position,
+        )
 
         hidden_states = outputs[0]