[Bugfix] Fix embedding to support 2D inputs (#5829)

2024-06-26 00:15:22 -07:00
parent 515080ad2f
commit 6806998bf9
1 changed files with 2 additions and 2 deletions
--- a/vllm/model_executor/layers/vocab_parallel_embedding.py
+++ b/vllm/model_executor/layers/vocab_parallel_embedding.py
@@ -310,7 +310,7 @@ class VocabParallelEmbedding(torch.nn.Module):
        output_parallel = F.embedding(masked_input.long(), self.weight)
        # Mask the output embedding.
        if self.tp_size > 1:
-            output_parallel.masked_fill_(input_mask.unsqueeze(1), 0)
+            output_parallel.masked_fill_(input_mask.unsqueeze(-1), 0)
        # Reduce across all the model parallel GPUs.
        output = tensor_model_parallel_all_reduce(output_parallel)
        return output