[Model] Support Gemma2 embedding model (#9004)

2024-10-04 23:57:05 -07:00
parent 53b3a33027
commit 15986f598c
5 changed files with 99 additions and 3 deletions
--- a/vllm/model_executor/models/gemma2.py
+++ b/vllm/model_executor/models/gemma2.py
@@ -278,11 +278,14 @@ class Gemma2Model(nn.Module):
        kv_caches: List[torch.Tensor],
        attn_metadata: AttentionMetadata,
        intermediate_tensors: Optional[IntermediateTensors],
+        inputs_embeds: Optional[torch.Tensor] = None,
    ) -> Union[torch.Tensor, IntermediateTensors]:
        if get_pp_group().is_first_rank:
-            hidden_states = self.embed_tokens(input_ids)
+            if inputs_embeds is not None:
+                hidden_states = inputs_embeds
+            else:
+                hidden_states = self.embed_tokens(input_ids)
            hidden_states *= self.normalizer
-
            residual = None
        else:
            assert intermediate_tensors is not None