[`GPTNeoX`] Fix GPTNeoX + Flash Attention 2 issue (#28645)

Update modeling_gpt_neox.py

src/transformers/models/gpt_neox/modeling_gpt_neox.py

@@ -390,7 +390,7 @@ class GPTNeoXFlashAttention2(GPTNeoXAttention):
             elif hasattr(self.config, "_pre_quantization_dtype"):
                 target_dtype = self.config._pre_quantization_dtype
             else:
-                target_dtype = self.q_proj.weight.dtype
+                target_dtype = self.query_key_value.weight.dtype
 
             logger.warning_once(
                 f"The input hidden states seems to be silently casted in float32, this might be related to"