Update modeling_gemmoe.py

modeling_gemmoe.py

@@ -65,7 +65,7 @@ logger = logging.get_logger(__name__)
 _CONFIG_FOR_DOC = "GemmoeConfig"
 
 def approx_gelu(x):
-    return 0.5 * x * (1 + torch.tanh(math.sqrt(2 / math.pi) * (x + 0.044715 * x**3)))
+    return x * torch.sigmoid(1.702 * x)
 
 def _get_unpad_data(attention_mask):
     seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)
@@ -348,9 +348,9 @@ class GemmoeFlashAttention2(GemmoeAttention):
                 f" {target_dtype}."
             )
 
-            query_states = query_states.to(target_dtype)
-            key_states = key_states.to(target_dtype)
-            value_states = value_states.to(target_dtype)
+            query_states = query_states.to(target_dtype, non_blocking=True)
+            key_states = key_states.to(target_dtype, non_blocking=True)
+            value_states = value_states.to(target_dtype, non_blocking=True)
 
         attn_output = self._flash_attention_forward(
             query_states, key_states, value_states, attention_mask, q_len, dropout=dropout_rate
@@ -845,12 +845,7 @@ 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
+            hidden_states = inputs_embeds * (self.config.hidden_size ** 0.5)
 
         past_seen_tokens = 0
         if use_cache:  # kept for BC (cache positions)