[Inference] support deepseek_vl2

paddlenlp/experimental/transformers/deepseek_v2/modeling.py

@@ -55,7 +55,7 @@
 )
 from paddlenlp.utils.log import logger
 
-__all__ = ["DeepseekV2ForCausalLMBlockInferenceModel"]
+__all__ = ["DeepseekV2ForCausalLMBlockInferenceModel", "DeepseekVLV2ForCausalLMBlockInferenceModel"]
 
 
 class DeepseekScalingRotaryEmbedding(nn.Layer):
@@ -1316,8 +1316,14 @@
         kwargs["max_input_length"] = self.max_seq_len
         kwargs["block_size"] = self.block_size
 
-        inputs_embeds = self.embed_tokens(ids_remove_padding)
-
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(ids_remove_padding)
+        else:
+            assert len(inputs_embeds.shape) == 3
+            # This is the case in the image-to-text model
+            # In the prefill phase, the language model is first fed with inputs_embeds instead of input_ids
+            # but in decoder phase, the language model is fed with input_ids just like normal text-to-text model.
+            inputs_embeds = inputs_embeds.reshape([-1, inputs_embeds.shape[2]])
         with dy2st_nocheck_guard_context():
             hidden_states, _ = self.transformer_block(
                 input_ids=input_ids,
@@ -1573,6 +1579,7 @@
     def prepare_inputs_for_generation(self, **kwargs):
         # only last token for inputs_ids if cache is defined in kwargs
         input_ids = kwargs["input_ids"]
+        inputs_embeds = kwargs.get("inputs_embeds", None)
         src_mask = kwargs.get("src_mask", None)
         block_tables = kwargs.get("block_tables", None)
 
@@ -1593,6 +1600,7 @@
 
         model_inputs = {
             "input_ids": input_ids,
+            "inputs_embeds": inputs_embeds,
             "src_mask": src_mask,
             "rope_emb": None,
             "pre_caches": pre_caches,
@@ -1613,6 +1621,7 @@
     def forward(
         self,
         input_ids,
+        inputs_embeds=None,
         src_mask=None,
         pre_caches=None,
         caches=None,
@@ -1630,6 +1639,7 @@
     ):
         outputs = self.deepseek_v2(
             input_ids,
+            inputs_embeds=inputs_embeds,
             src_mask=src_mask,
             caches=caches,
             rope_emb=None,
@@ -1805,3 +1815,19 @@
         )
 
         return logits, hidden_states
+
+
+class DeepseekVLV2ForCausalLMBlockInferenceModel(DeepseekV2ForCausalLMBlockInferenceModel):
+    def __init__(self, config: DeepseekV2Config):
+        super().__init__(config, base_model_prefix="language.model")
+
+    def get_input_embeddings(self):
+        return self.deepseek_v2.embed_tokens
+
+    @paddle.no_grad()
+    def set_state_dict(self, state_dict):
+        if "language.lm_head.weight" in state_dict:
+            self.lm_head.weight.set_value(
+                paddle.to_tensor(state_dict["language.lm_head.weight"]).cast(self.lm_head.weight.dtype)
+            )
+        self.deepseek_v2.set_state_dict({k: state_dict[k] for k in state_dict.keys()})