Update modeling_gemmoe.py

modeling_gemmoe.py

@@ -167,10 +167,8 @@ class GemmoeRMSNorm(nn.Module):
         self.weight = nn.Parameter(torch.zeros(dim))
 
     def _norm(self, x):
-        # Ensure the entire normalization is done in float32
-        x_float = x.float()  # upcast to float32
-        mean = x_float.pow(2).mean(-1, keepdim=True)
-        normed_x = x_float * torch.rsqrt(mean + self.eps)
+        x_float = x.float()
+        normed_x = x_float * torch.rsqrt(x_float.pow(2).mean(-1, keepdim=True) + self.eps)
         return normed_x
 
     def forward(self, x):
@@ -191,16 +189,14 @@ class GemmoeRotaryEmbedding(nn.Module):
 
     def _set_cos_sin_cache(self, seq_len, device, dtype):
         self.max_seq_len_cached = seq_len
-        freq_exponents = (2.0 / self.dim) * (
-            torch.arange(self.dim // 2, dtype=torch.int64, device="cpu").float()
-        )
+        freq_exponents = (2.0 / self.dim) * (torch.arange(self.dim // 2, dtype=torch.float32, device="cpu").float())
         timescale = self.base ** freq_exponents
-        positions = torch.arange(self.max_seq_len_cached, device="cpu", dtype=torch.int64).float()
+        positions = torch.arange(self.max_seq_len_cached, device="cpu", dtype=torch.float32).float()
         radians_new = positions[..., None] / timescale[None, None, :]
         radians_new = radians_new.squeeze(0)
         emb = torch.cat((radians_new, radians_new), dim=-1)
-        cos = emb.cos().to(device=device, non_blocking=True)
-        sin = emb.sin().to(device=device, non_blocking=True)
+        cos = emb.cos().to(device=device, dtype=dtype, non_blocking=True)
+        sin = emb.sin().to(device=device, dtype=dtype, non_blocking=True)
         self.register_buffer("cos_cached", cos, persistent=False)
         self.register_buffer("sin_cached", sin, persistent=False)
 
@@ -209,10 +205,7 @@ class GemmoeRotaryEmbedding(nn.Module):
             seq_len = x.size(2)
         if seq_len > self.max_seq_len_cached:
             self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype)
-        return (
-            self.cos_cached[:seq_len],
-            self.sin_cached[:seq_len],
-        )
+        return self.cos_cached[:seq_len], self.sin_cached[:seq_len]
 
 # Copied from transformers.models.llama.modeling_llama.rotate_half
 def rotate_half(x):
@@ -222,27 +215,7 @@ def rotate_half(x):
     return torch.cat((-x2, x1), dim=-1)
 
 
-# Copied from transformers.models.llama.modeling_llama.apply_rotary_pos_emb
 def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
-    """Applies Rotary Position Embedding to the query and key tensors.
-
-    Args:
-        q (`torch.Tensor`): The query tensor.
-        k (`torch.Tensor`): The key tensor.
-        cos (`torch.Tensor`): The cosine part of the rotary embedding.
-        sin (`torch.Tensor`): The sine part of the rotary embedding.
-        position_ids (`torch.Tensor`, *optional*):
-            Deprecated and unused.
-        unsqueeze_dim (`int`, *optional*, defaults to 1):
-            The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
-            sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
-            that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
-            k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
-            cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
-            the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
-    Returns:
-        `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
-    """
     seq_len, dim = q.shape[-2], q.shape[-1]
     cos = cos[:seq_len].view(1, 1, seq_len, dim)
     sin = sin[:seq_len].view(1, 1, seq_len, dim)
@@ -250,7 +223,6 @@ def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
     k_embed = (k * cos) + (rotate_half(k) * sin)
     return q_embed, k_embed
 
-
     
 # Copied from transformers.models.llama.modeling_llama.repeat_kv
 def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
@@ -662,7 +634,7 @@ class GemmoeBlockSparseTop2MLP(nn.Module):
         super().__init__()
         self.ffn_dim = config.intermediate_size
         self.hidden_dim = config.hidden_size
-
+        
         self.w1 = nn.Linear(self.hidden_dim, self.ffn_dim, bias=False)
         self.w2 = nn.Linear(self.ffn_dim, self.hidden_dim, bias=False)
         self.w3 = nn.Linear(self.hidden_dim, self.ffn_dim, bias=False)
@@ -683,16 +655,13 @@ class GemmoeSparseMoeBlock(nn.Module):
         self.num_experts = config.num_local_experts
         self.top_k = 2
 
-        # gating
         self.gate = nn.Linear(self.hidden_dim, self.num_experts, bias=False)
-
         self.experts = nn.ModuleList([GemmoeBlockSparseTop2MLP(config) for _ in range(self.num_experts)])
 
     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)
 
-        # router_logits: (batch * sequence_length, n_experts)
         router_logits = self.gate(hidden_states)
         routing_weights = F.softmax(router_logits, dim=1)
         topk_weight, topk_idx = torch.topk(routing_weights, self.top_k, dim=-1, sorted=False)
@@ -700,17 +669,18 @@ class GemmoeSparseMoeBlock(nn.Module):
 
         hidden_states = hidden_states.repeat_interleave(self.top_k, dim=0)
 
-        y = torch.empty_like(hidden_states)
-
-        flat_topk_idx = topk_idx.view(-1)
+        expert_outputs = []
         for i in range(self.num_experts):
-            expert = self.experts[i]
-            expert_output = expert(hidden_states[flat_topk_idx == i])
-            y[flat_topk_idx == i] = expert_output
-
-        y = (y.view(*topk_weight.shape, -1) * topk_weight.unsqueeze(-1)).sum(dim=1)
+            expert_input = hidden_states[topk_idx[:, i]]
+            expert_output = self.experts[i](expert_input)
+            expert_outputs.append(expert_output)
 
-        final_hidden_states = y.reshape(batch_size, sequence_length, hidden_dim)
+        expert_outputs = torch.stack(expert_outputs, dim=1)
+        expert_outputs = expert_outputs.view(batch_size * sequence_length, self.top_k, -1)
+        
+        final_hidden_states = torch.einsum("bke,bkd->bed", topk_weight, expert_outputs)
+        final_hidden_states = final_hidden_states.view(batch_size, sequence_length, hidden_dim)
+        
         return final_hidden_states.to(hidden_states.dtype), router_logits.to(hidden_states.dtype)
 
     
@@ -719,12 +689,10 @@ 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,
@@ -737,32 +705,9 @@ 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,
@@ -772,21 +717,17 @@ class GemmoeDecoderLayer(nn.Module):
             use_cache=use_cache,
         )
         hidden_states = residual + hidden_states
-
-        # Fully Connected
+        
         residual = hidden_states
-        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = self.input_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,)
 
@@ -1009,14 +950,6 @@ 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):
@@ -1036,8 +969,12 @@ class GemmoeModel(GemmoePreTrainedModel):
         # embed positions
         hidden_states = inputs_embeds
 
-        # normalized
-        hidden_states = hidden_states * (self.config.hidden_size**0.5)
+        # 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
 
         # decoder layers
         all_hidden_states = () if output_hidden_states else None