[Model] Pooling models default to using chunked prefill & prefix caching if supported. #20930
Conversation
|
👋 Hi! Thank you for contributing to the vLLM project. 💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels. Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging. To run CI, PR reviewers can either: Add 🚀 |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Summary of Changes
Hello @noooop, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!
This pull request significantly improves how vLLM handles Hugging Face ForSequenceClassification models. It introduces an automatic conversion mechanism that allows any ForCausalLM model to be used for sequence classification without requiring explicit registration or custom wrapper classes. This change streamlines model integration, centralizes classification task detection, and enhances the flexibility of vLLM's model loading and execution pipeline, while also addressing a reported issue with the TRANSFORMERS implementation for these models.
Highlights
- Automatic
ForSequenceClassificationSupport: I've implemented a new mechanism to automatically convertForCausalLMmodels intoForSequenceClassificationmodels by dynamically applying an adapter. This removes the need for explicit registration and custom wrapper classes for each specific model type, streamlining model integration. - Centralized Classification Task Logic: I've introduced a new
_is_classify_taskmethod and updated the task resolution invllm/config.py. This provides a more robust way to identify and handle classification tasks, ensuring they are correctly routed to the 'pooling' runner. - Refactored Model Registry: The model registry has been cleaned up by removing explicit entries for automatically convertible
ForSequenceClassificationmodels. The registry's ability to inspect and normalize these architectures on the fly has been enhanced, making it more flexible. - Improved
ScoreModelBias Handling: The_ScoreModeladapter now correctly respects thescore_biasconfiguration from the Hugging Face model config. This allows classification heads to have a bias if specified by the original model. - Enhanced Testing Coverage: The test suite for model initialization (
test_initialization.py) now includes automatically convertedForSequenceClassificationmodels, ensuring the new automatic support mechanism functions as expected.
Using Gemini Code Assist
The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.
Invoking Gemini
You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.
| Feature | Command | Description |
|---|---|---|
| Code Review | /gemini review |
Performs a code review for the current pull request in its current state. |
| Pull Request Summary | /gemini summary |
Provides a summary of the current pull request in its current state. |
| Comment | @gemini-code-assist | Responds in comments when explicitly tagged, both in issue comments and review comments. |
| Help | /gemini help |
Displays a list of available commands. |
Customization
To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.
Limitations & Feedback
Gemini Code Assist is currently in preview and may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments to provide feedback.
You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.
Footnotes
-
Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩
There was a problem hiding this comment.
Code Review
This pull request refactors the model loading and registration logic to automatically support ForSequenceClassification models, which is a great improvement. The changes involve removing hardcoded model registrations and introducing dynamic conversion logic.
My review has identified a critical issue in the new conversion logic in vllm/model_executor/model_loader/utils.py that could lead to incorrect model loading. I've also pointed out a few medium-severity issues related to code clarity, maintainability, and a typo. Addressing these points will improve the robustness and readability of the new implementation.
DebertaV2ForSequenceClassification uses classifier, while this pr uses score, so it is not supported. Is the title not well chosen? In fact, this pr only implements a small amount of functionality. |
noooop
changed the title
|
This pull request has merge conflicts that must be resolved before it can be |
noooop
changed the title
noooop
changed the title
noooop
requested review from
ywang96,
simon-mo,
WoosukKwon,
youkaichao,
robertgshaw2-redhat,
mgoin and
tlrmchlsmth
as code owners
|
Looks like Prithvi failure in https://buildkite.com/vllm/ci/builds/26638/steps/canvas?jid=01989bd9-f941-48bf-b60a-6b25eb361cb9 is caused by this PR, please fix |
|
|
Looks like for this model |
…ing if supported. (vllm-project#20930) Signed-off-by: wang.yuqi <noooop@126.com>
…ing if supported. (vllm-project#20930) Signed-off-by: wang.yuqi <noooop@126.com> Signed-off-by: Paul Pak <paulpak58@gmail.com>
…ing if supported. (vllm-project#20930) Signed-off-by: wang.yuqi <noooop@126.com>
…ing if supported. (vllm-project#20930) Signed-off-by: wang.yuqi <noooop@126.com> Signed-off-by: Boyuan Feng <boyuan@meta.com>
…ing if supported. (vllm-project#20930) Signed-off-by: wang.yuqi <noooop@126.com> Signed-off-by: Diego-Castan <diego.castan@ibm.com>
…ing if supported. (vllm-project#20930) Signed-off-by: wang.yuqi <noooop@126.com>
…ing if supported. (vllm-project#20930) Signed-off-by: wang.yuqi <noooop@126.com>
…ing if supported. (vllm-project#20930) Signed-off-by: wang.yuqi <noooop@126.com>
…ing if supported. (vllm-project#20930) Signed-off-by: wang.yuqi <noooop@126.com> Signed-off-by: Xiao Yu <xiao.yu@amd.com>
…ing if supported. (vllm-project#20930) Signed-off-by: wang.yuqi <noooop@126.com> Signed-off-by: Xiao Yu <xiao.yu@amd.com>
…ing if supported. (vllm-project#20930) Signed-off-by: wang.yuqi <noooop@126.com>
…ing if supported. (vllm-project#20930) Signed-off-by: wang.yuqi <noooop@126.com>
…ing if supported. (vllm-project#20930) Signed-off-by: wang.yuqi <noooop@126.com>
|
@noooop where is As far as I know |
Here is some detailed background Alibaba-NLP/gte-Qwen2-1.5B-instruct uses the methods mentioned in llm2vec (As far as I know, this is the first model that uses
vllm/vllm/model_executor/models/qwen2.py Lines 211 to 218 in e08a3a3 #24636 ( Sorry, I should have added a comment I hope to use |
|
Thanks for the context. The problem is that all legacy encoder models (where |
in #24636 Lines 568 to 580 in e08a3a3 I thought self.model_config.pooler_config is not None means It is a pooling model, means not self.model_config.is_encoder_decoder. Perhaps this taxonomy is wrong. I'm sorry about it I think the following code should not be run in the encoder_decoder model, but I am not familiar with this code and don't know where to put it. The current vllm model config is spaghetti code, We need to find a way to refactor it |
|
Sorry I should clarify, I mean encoder-only models. vLLM doesn't support encoder-decoder models anymore so they aren't relevant here. I noticed this issue when using the Transformers backend to load Albert. if self.model_config: # True
if self.model_config.pooler_config: # True
pooling_type = self.model_config.pooler_config.pooling_type
if pooling_type is None or pooling_type.lower() != "last": # False
disable_chunked_prefill_reasons.append(
"Only \"last\" pooling supports chunked "
"prefill and prefix caching; disabling both.")
if not getattr(self.model_config.hf_config, "is_causal", True): # False! But should be True
disable_chunked_prefill_reasons.append(
"Only models using causal attention supports chunked "
"prefill and prefix caching; disabling both.") This is the problem because |
|
|
||
| def default_pooling_type(pooling_type: str) -> object: | ||
| """Set default_pooling_type decorator. """ | ||
|
|
||
| def func(model: object): | ||
| model.default_pooling_type = pooling_type | ||
| return model | ||
|
|
||
| return func | ||
|
|
||
|
|
||
| def get_default_pooling_type(model: Union[type[object], object]) -> str: | ||
| return getattr(model, "default_pooling_type", "LAST") | ||
|
|
||
|
|
There was a problem hiding this comment.
This PR introduces default_pooling_type to avoid enable chunked_prefill for bidirectional attention models such as BERT.
However, the Transformers backend does not have default_pooling_type. (This results in the default being last pooling)
How about, disable chunked_prefill for all pooling models using Transformers backend. Until we have a way to distinguish whether the Transformers backend uses bidirectional attention or causal attention
There was a problem hiding this comment.
Thanks for the pointer.
However, the Transformers backend does not have default_pooling_type. (This results in the default being last pooling)
For the Transformers backend I'm not sure it makes sense to have a static default pooling type because it could load models that should have different defaults. When is this checked? Could we set default_pooling_type dynamically on init of TransformersModel?
How about, disable chunked_prefill for all pooling models using Transformers backend. Until we have a way to distinguish whether the Transformers backend uses bidirectional attention or causal attention
TransformersModel also works with causal embeddings models so I'd rather not disable these features globally for the feature.
There was a problem hiding this comment.
TransformersModelalso works with causal embeddings models so I'd rather not disable these features globally for the feature.
I also want TransformersModel to support chunked_prefill, but there is no good way to distinguish whether TransformersModel is causal attention or bidirectional attention
In particular, Transformers allow loading arbitrary modules via trust_remote_code. (
I'm not 100% sure vllm currently supports loading arbitrary modules, but it will eventually.)
┓( ´∀` )┏
TL;DR
Essential Elements of an Effective PR Description Checklist
supported_models.mdandexamplesfor a new model.Purpose
vllm/vllm/model_executor/models/adapters.py
Lines 175 to 180 in 3fc9644
This piece of code may not affect the main process…
refer to #20012
set default_pooling_type in vllm.config
Keep Decode-only SequenceClassification models support auto prefix cache
The pooler_config in the startup log shows the correct pooling_type instead of None, which will help with debugging
LLM.encode() and pooling server uses ALL pooling, which does not support chunked prefill. e.g. jason9693/Qwen2.5-1.5B-apeach
cc @DarkLight1337 @maxdebayser
Test Plan
pytest -s -vvv tests/test_config.py::test_default_pooling_type
pytest -s -vvv tests/models/language/pooling/test_auto_prefix_cache_support.py
pytest -s -vvv tests/entrypoints/llm/test_classify.py::test_encode_api
pytest -s -vvv tests/entrypoints/openai/test_classification.py::test_pooling
Add test for default_pooling_type in tests/models/language/pooling/mteb_utils.py, double check all implementations use the correct default_pooling_type
Test Result
passed
(Optional) Documentation Update
Fix #20894
Fix #19950
implicit conversion part split to #21103