Update modeling_gemmoe.py

modeling_gemmoe.py

@@ -169,8 +169,8 @@ class GemmoeRMSNorm(nn.Module):
 
     def forward(self, x):
         output = self._norm(x.float()).type_as(x)
-        return output * (self.weight.to(x.device) + 1)  # Move self.weight to the same device as x
-    
+        return output * (self.weight + 1)
+
 ALL_LAYERNORM_LAYERS.append(GemmoeRMSNorm)
 
 class GemmoeRotaryEmbedding(nn.Module):
@@ -271,10 +271,10 @@ class GemmoeAttention(nn.Module):
         self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias)
         self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False)
         self.rotary_emb = GemmoeRotaryEmbedding(
-            self.head_dim,
-            max_position_embeddings=self.max_position_embeddings,
-            base=self.rope_theta,
-        )
+				self.head_dim,
+				max_position_embeddings=self.max_position_embeddings,
+				base=self.rope_theta,
+			)
 
     def forward(
         self,
@@ -312,11 +312,6 @@ class GemmoeAttention(nn.Module):
         key_states = self.k_proj(hidden_states)
         value_states = self.v_proj(hidden_states)
 
-        # Move query_states, key_states, and value_states to the same device as hidden_states
-        query_states = query_states.to(hidden_states.device)
-        key_states = key_states.to(hidden_states.device)
-        value_states = value_states.to(hidden_states.device)
-
         query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
         key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
         value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
@@ -1205,14 +1200,6 @@ class GemmoeForCausalLM(GemmoePreTrainedModel):
         )
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
 
-        device = input_ids.device if input_ids is not None else inputs_embeds.device
-        attention_mask = attention_mask.to(device) if attention_mask is not None else None
-        position_ids = position_ids.to(device) if position_ids is not None else None
-        past_key_values = [t.to(device) for t in past_key_values] if past_key_values is not None else None
-        inputs_embeds = inputs_embeds.to(device) if inputs_embeds is not None else None
-        labels = labels.to(device) if labels is not None else None
-        cache_position = cache_position.to(device) if cache_position is not None else None
-
         outputs = self.model(
             input_ids=input_ids,
             attention_mask=attention_mask,