Supporting embedding models

[jc9123]

Added support for BGE-M3 which is an embedded model. The embedded vector is stored in the .embed of the returned Request Object. A example test case can be found in tests/models/test_embedded.py.

FIX #9847

[robertgshaw2-neuralmagic]

Thank you for the PR.

A couple comments:

• The major thing that should be done to add BGE to vLLM is that the model should be re-implemented using the layers from vLLM. vLLM has optimized inference kernels for most layers in transformers, we should use these in the implementation. Ideally, you would also use the ColumnParallel and RowParallel layers to support multi-gpu inference as well. Currently, it looks like you are using the transformers model inside vLLM in the definition of BGE. Take a look at the implementation of some of the other models for inspiration. You will need to modify attention to be bidirectional, but they should be instructive

• I do not think that we should be adding a new input params_dict to vLLM. Additionally, we should not be tokenizing the input inside the forward method. The inputs to encoder models are tokens, just like for LLMs. We should be able to leverage the existing tokenization / sequence group / input_metadata infrastructure for encoder models

[jc9123]

I've removed param_dict and and replaced it with the existing tokenization infrastructure. I think I've changed the inference layer to use vLLM's inference layer as well, but I may have missed something so let me know if I need to fix anything. Thanks!

[robertgshaw2-neuralmagic]

@jc9123 i will take a look

[robertgshaw2-neuralmagic]

I left a few tactical nits above which should be addressed directly.

However, there is significant work to do to get this PR into a place that is ready to merge. I would suggest we split up the work into a series of PRs.

Milestone 1: Basic functionality using Hugging Face models

In this step, we will get basic functionality working, meaning:

• Model is implemented using HF transformers
• Update the LLM front end to support LLM.encode(). This will receive input text and return an embedding.

See below for a detailed workplan to address Milestone 1

Additionally, make sure there are checks that turn off incompatible features are turned off for embedding models:

• Tensor parallelism > 1
• Speculative decoding
• Automatic prefix caching
• LORA
• Quantization
• CUDAGraphs

Milestone 2: Implement the model properly

In this step, we will implement the model using the vllm layer primitives rather than using the transformers implementation of the model. This will enable us to support tensor parallelism and have more optimized inference.

The goal is to remove the need to import transformers and instead use the vLLM layer primitives that are optimized for inference like:

• ColumnParallelLinear, RowParallelLinear, Attention, etc.

Milestone 3: Support sparse vectors, colbert, normalization, etc

In this step, we will add more examples of types of embeddings that can be handled. See here for examples of what this means

Milestone 4: Add the OpenAI Server Front End

• Add support in the entrypoints/openai for the embeddings functionality.
• https://platform.openai.com/docs/api-reference/embeddings

Milestone 5: Expand feature set [optional]

• Multimodal embedding models (e.g. CLIP)
• Support quantized models
• Support quantized embeddings
• Support variable length embeddings
• [ others that I am missing ]

[robertgshaw2-neuralmagic]

This file should not be touched

[robertgshaw2-neuralmagic]

This should not be a dependency. If you need it for testing, add it to requirements-dev

[robertgshaw2-neuralmagic]

This file should not be touched

[robertgshaw2-neuralmagic]

this file should not be touched

[robertgshaw2-neuralmagic]

let's use the term embedding_model to match the terminology used by OAI

https://platform.openai.com/docs/api-reference/embeddings

[robertgshaw2-neuralmagic]

remove

[robertgshaw2-neuralmagic]

nit

[robertgshaw2-neuralmagic]

nit

[robertgshaw2-neuralmagic]

All of the logic for making stopping decisions occurs in llm_engine.check_stop

This is correct logic (in that each sequence only needs to run once), but please move it into the check_stop function

[robertgshaw2-neuralmagic]

See my feedback below on how this file needs to be reworked

[robertgshaw2-neuralmagic]

@jc9123 I made a PR on your branch to:

• fix the nits listed here
• get started implementing the XLMRobertaModel

[robertgshaw2-neuralmagic]

If you want to chat live ping me at [ rshaw at neuralmagic dot com ]

Milestone 1 Workplan

A) Accept my PR to address the nits

• Rs/embedding step 1 jc9123/vllm#1

B) Finish the XLMRobertaModel definition started

In the PR I did against your branch, I setup the basics of the model, which creates the weights using the config during init and then loads the weights during load_weights() (which is called by vLLM when the model is being loaded.

from transformers import XLMRobertaConfig
from xlm_roberta import XLMRobertaModel

model_id = "BAAI/bge-m3"
config = XLMRobertaConfig.from_pretrained(model_id)
model = XLMRobertaModel(config=config)
model.load_weights(model_id)

You need to finish off this function by implementing the forward method. forward needs to convert from the input data passed around by vllm into the formats expected by the hugging face implementation of the model (which we are using here). You should handle both the batched and non-batched case

Note: BAAI/bge-m3 uses the XLMRobertaModel architecture. So In vLLM we will implement the XLMRobertaModel architecture and then any embedding model that uses this architecture can use the definition

C) Update the LLM front end UX

I would like the UX to look something like this:

from vllm import LLM, EmbeddingParams

# LLM should autodetect this is an embedding model since it uses `XLMRobertaModel`
model = LLM("BAAI/bge-m3")
parameters = EmbeddingParams()
sentences = ["The quick...", "brown fox...", ... "the lazy..."]

output = model.encode(sentences, parameters=parameters)

Then, we can add things like normalize, return_sparse, return_dense, return_colbert_vecs via the EmbeddingParams in the future. [ This does not need to be done for this PR ]

We can look at the following repos for potential features to support via encode

• SentenceTransformers https://www.sbert.net/docs/package_reference/SentenceTransformer.html
• FlagEmbedding https://github.com/FlagOpen/FlagEmbedding/tree/master/FlagEmbedding/BGE_M3

D) Clean up internals

Update RequestOutput in vllm/vllm/outputs.py

• Create an EmbeddingOutput class
• Modify RequestOutput to look like the following

class RequestOutput:
        request_id: str,
        prompt: str,
        prompt_token_ids: List[int],
        prompt_logprobs: Optional[PromptLogprobs],
        outputs: Union[List[CompletionOutput], List[EmbeddingOutput],  # << RequestOutput can be either Completion or Embedding
        finished: bool,
        metrics: Optional[RequestMetrics] = None,
        lora_request: Optional[LoRARequest] = None,

Fix terminology

• embedded_model --> embedding_model
• embed --> embedding

Move finish sequence logic to check_stop

• Currently the logic is in llm_engine._process_model_output()
• The logic should be in llm_engine._check_stop

Automatically detect that the model is an embedding model

• The user should not have to specify that it is an embedding model
• Somewhere in the vllm code, create a registry that selects which models are embedding and which models are decoders

^^^ Once this is all done, we can take another check-in

E) Turn Off KV Cache Memory usage

• note: we still allocate the KV cache even when its not being used for embedding models

F) Update Pass EmbeddingParams around instead of SamplingParams

• note: this is going to require a lot of work

G) Update parameters for max batching

• We should be able to support many more tokens in a batch than we currently do for LLMs

[jc9123]

Hi Robert, thank you for the detail feedback! I will begin working on milestone 1 and let you know if I have any questions.

[robertgshaw2-neuralmagic]

@jc9123 updated comment above

[jc9123]

@jc9123 updated comment above

Thanks! I will check in with you after finishing steps a-d.

[ywang96]

IMO it's a better idea to create a separate class for output of embedding models. Many attributes of RequestOutput are tied to language models but not relevant to embedding models. We can discuss this in the later stages of this PR series.

[robertgshaw2-neuralmagic]

Agree @ywang96

I think we will want to refactor RequestOutput and SamplingParams to be abstract with separate versions for Completion vs Embedding

It will be tricky to do this without breaking changes to the UX

[robertgshaw2-neuralmagic]

@jc9123 any update?

[jc9123]

@jc9123 any update?

Sorry for the delay but I will need a bit more time to finish milestone 1. I should be able to finish it by early next week.

[brandonsorensen]

I don't want to step on @jc9123's toes, but I'd be happy to help out on this issue. For example, I could already start work on the front end bits.

[jc9123]

@jc9123 any update?

Hi Robert, sorry for the delay but I've finished checkpoint 1.a->1.d. Could you take a look; right now, I've mostly copied EmbeddingParams and EmbeddingOutput from their Sampling counterpart to ensure that everything still work but was wondering if there is something specific I should add/change.

Edit : Found some errors in implementation, will fix by EOD.
Edit2: Fixed!

[CatherineSue]

@jc9123 @robertgshaw2-neuralmagic , it seems this PR has a few milestones that are overlapped with my PR #3734. I have also separated EmbeddingOutput from CompletionOutput. Should we collaborate so we don't do duplicate work?

[jc9123]

@jc9123 @robertgshaw2-neuralmagic , it seems this PR has a few milestones that are overlapped with my PR #3734. I have also separated EmbeddingOutput from CompletionOutput. Should we collaborate so we don't do duplicate work?

Hi Catherine, we should definitely collaborate to avoid duplicate work. Between our PR, I think we should have the same implementation for the vLLM front end UX and also the EmbeddingOutput and EmbeddingParams class.

[KylinMountain]

Any update? Can we use other embedding model now?

[whybeyoung]

lgtm, but any update ?

[simon-mo]

Bart has been added, cross attention as well, i believe there should be some small incrementally work in hooking this models with the embedding API

If @jc9123 is not actively working on this anymore, please feel free to open a new PR base off this.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. @jc9123 please rebase it. https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork