Enhance device compatibility by adding MPS support and ensuring tensor operations respect device context

modeling_deepseekv2.py

@@ -109,10 +109,11 @@ class DeepseekV2RMSNorm(nn.Module):
 
     def forward(self, hidden_states):
         input_dtype = hidden_states.dtype
+        device = hidden_states.device
         hidden_states = hidden_states.to(torch.float32)
         variance = hidden_states.pow(2).mean(-1, keepdim=True)
         hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
-        return self.weight * hidden_states.to(input_dtype)
+        return (self.weight.to(device) * hidden_states).to(input_dtype)
 
 
 ALL_LAYERNORM_LAYERS.append(DeepseekV2RMSNorm)
@@ -1468,6 +1469,8 @@ class DeepseekV2Model(DeepseekV2PreTrainedModel):
         super().__init__(config)
         self.padding_idx = config.pad_token_id
         self.vocab_size = config.vocab_size
+        # Set device to MPS if available, otherwise fallback to CPU
+        self.device = torch.device("mps") if torch.backends.mps.is_available() else torch.device("cpu")
 
         self.embed_tokens = nn.Embedding(
             config.vocab_size, config.hidden_size, self.padding_idx
@@ -1652,6 +1655,9 @@ class DeepseekV2ForCausalLM(DeepseekV2PreTrainedModel):
         self.model = DeepseekV2Model(config)
         self.vocab_size = config.vocab_size
         self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        # Move model to MPS if available
+        if torch.backends.mps.is_available():
+            self.to("mps")
 
         # Initialize weights and apply final processing
         self.post_init()
@@ -1782,6 +1788,15 @@ class DeepseekV2ForCausalLM(DeepseekV2PreTrainedModel):
             inputs_embeds=None,
             **kwargs,
     ):
+        # Move inputs to MPS device if available
+        if torch.backends.mps.is_available():
+            if input_ids is not None:
+                input_ids = input_ids.to("mps")
+            if attention_mask is not None:
+                attention_mask = attention_mask.to("mps")
+            if inputs_embeds is not None:
+                inputs_embeds = inputs_embeds.to("mps")
+                
         past_length = 0
         if past_key_values is not None:
             if isinstance(past_key_values, Cache):