Update modeling_gemmoe.py

modeling_gemmoe.py

@@ -689,10 +689,12 @@ class GemmoeDecoderLayer(nn.Module):
     def __init__(self, config: GemmoeConfig, layer_idx: int):
         super().__init__()
         self.hidden_size = config.hidden_size
-        
+
         self.self_attn = GEMMOE_ATTENTION_CLASSES[config._attn_implementation](config, layer_idx)
+
         self.block_sparse_moe = GemmoeSparseMoeBlock(config)
         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)
 
     def forward(
         self,
@@ -705,9 +707,32 @@ class GemmoeDecoderLayer(nn.Module):
         use_cache: Optional[bool] = False,
         **kwargs,
     ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
+        if "padding_mask" in kwargs:
+            warnings.warn(
+                "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`"
+            )
+        """
+        Args:
+            hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
+            attention_mask (`torch.FloatTensor`, *optional*): attention mask of size
+                `(batch, sequence_length)` where padding elements are indicated by 0.
+            past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states
+            output_attentions (`bool`, *optional*):
+                Whether or not to return the attentions tensors of all attention layers. See `attentions` under
+                returned tensors for more detail.
+            output_router_logits (`bool`, *optional*):
+                Whether or not to return the logits of all the routers. They are useful for computing the router loss, and
+                should not be returned during inference.
+            use_cache (`bool`, *optional*):
+                If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
+                (see `past_key_values`).
+        """
+
         residual = hidden_states
+
         hidden_states = self.input_layernorm(hidden_states)
 
+        # Self Attention
         hidden_states, self_attn_weights, present_key_value = self.self_attn(
             hidden_states=hidden_states,
             attention_mask=attention_mask,
@@ -717,17 +742,21 @@ class GemmoeDecoderLayer(nn.Module):
             use_cache=use_cache,
         )
         hidden_states = residual + hidden_states
-        
+
+        # Fully Connected
         residual = hidden_states
-        hidden_states = self.input_layernorm(hidden_states)
+        hidden_states = self.post_attention_layernorm(hidden_states)
         hidden_states, router_logits = self.block_sparse_moe(hidden_states)
         hidden_states = residual + hidden_states
 
         outputs = (hidden_states,)
+
         if output_attentions:
             outputs += (self_attn_weights,)
+
         if use_cache:
             outputs += (present_key_value,)
+
         if output_router_logits:
             outputs += (router_logits,)
 
@@ -950,6 +979,14 @@ class GemmoeModel(GemmoePreTrainedModel):
         if inputs_embeds is None:
             inputs_embeds = self.embed_tokens(input_ids)
 
+        # Scale embeddings
+        # Fix for precision issue when casting to bfloat16
+        hidden_size_sqrt = math.sqrt(self.config.hidden_size)
+        if inputs_embeds.dtype == torch.bfloat16:
+            pass
+        
+        hidden_states = inputs_embeds * hidden_size_sqrt
+
         past_seen_tokens = 0
         if use_cache:  # kept for BC (cache positions)
             if not isinstance(past_key_values, StaticCache):
@@ -969,12 +1006,8 @@ class GemmoeModel(GemmoePreTrainedModel):
         # embed positions
         hidden_states = inputs_embeds
 
-        # Scale embeddings
-        hidden_size_sqrt = math.sqrt(self.config.hidden_size)
-        if inputs_embeds.dtype == torch.bfloat16:
-            hidden_states = inputs_embeds * hidden_size_sqrt
-        else:
-            hidden_states = inputs_embeds * hidden_size_sqrt
+        # normalized
+        hidden_states = hidden_states * (self.config.hidden_size**0.5)
 
         # decoder layers
         all_hidden_states = () if output_hidden_states else None