[RFC]: Make any vLLM model a pooling model

[DarkLight1337]

Motivation.

Currently, we have to open new PRs to add pooling functionality for existing architectures supported in vLLM. Since the code involved is basically the same for each model, there is potential to automate away this boilerplate.

Proposed Change.

Implement a pooling adapter that can be applied to any existing text generation model in vLLM. To preserve features such as LoRA, PP and multimodality, the adapter simply creates a new subclass of the original model.

• 10769

The pooling adapter to apply depends on the purpose of the model. To facilitate this, the embedding task will be split into the following tasks:

• Embedding (embed), with default pooler pooling_type=LAST, normalize=True
  • Usually one output per prompt
• Classification (classify), with default pooler pooling_type=LAST, softmax=True
  • Usually one output per prompt
  • Also used for cross-encoder models
  • The scoring module will be automatically loaded from sentence-transformers.
  • Other models should register a separate architecture so that the model contains the correct modules.
• Reward Modeling (reward), with default pooler pooling_type=ALL
  • Usually one output per token per prompt
  • If the reward model has any additional modules compared to the text generation model, a separate architecture should be registered so that the model contains the correct modules.
  • To avoid peiyi9979/math-shepherd-mistral-7b-prm conflicting with other LlamaForCausalLM pooling models, we will remove the pooler default from the text generation model. Since users already have to override step_tag_id etc. anyway to use that model, this should not be a major breaking change.

Meanwhile, current embedding-related classes will be renamed to avoid confusion between embed and other pooling tasks:

• [Internal] VllmModelForEmbedding -> VllmModelForPooling
• [Internal] is_embedding_model -> is_pooling_model
• [Internal] EmbeddingModelRunner -> PoolingModelRunner
• [API] EmbeddingOutput -> PoolingOutput
• [API] EmbeddingRequestOutput -> PoolingRequestOutput

Feedback Period.

1-2 weeks

CC List.

@youkaichao @mgoin @robertgshaw2-neuralmagic @maxdebayser

Any Other Things.

Note that we can still directly map to pooling models in the model registry. This is used when the model architecture has different pooling defaults (e.g. pooling_type=CLS for BERT) or additional modules (e.g. score in Qwen2ForRewardModel). For models that already support pooling, the adapter returns the original model without modifications.

Before submitting a new issue...

• [X] Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.

[maxdebayser]

I think this makes sense, but we should add a verification to prevent models that are already loaded as embedding models to be wrapped again. Would this wrapper assume a default Pooling configuration or should we raise an error if the user doesn't provide this information on the command line?

[DarkLight1337]

I think this makes sense, but we should add a verification to prevent models that are already loaded as embedding models to be wrapped again. Would this wrapper assume a default Pooling configuration or should we raise an error if the user doesn't provide this information on the command line?

I'm thinking of removing the existing embedding model implementations and replacing them with the wrapper. The embedding wrapper's default pooler config is shown in the above example. I'm thinking of resolving the pooler config in the same way as current code (i.e.: user > sentence-transformers > model), so the user doesn't have to provide this information in most cases.

[maxdebayser]

What would happen if the user loads a BertModel without --task-embedding? Would there be a validation to prevent this from happening our would we allow the decoder-only model runner to try to run the model for text generation?

[DarkLight1337]

In is_text_generation_model, we check for the existence of compute_logits and sample methods. Since those methods aren't in BertModel, --task generation is not allowed for that model.

[maxdebayser]

Ok, I think this makes sense.

[DarkLight1337]

Update: I have expanded the scope of this RFC to cover all pooling-based models.