File tree Expand file tree Collapse file tree 2 files changed +22
-1
lines changed Expand file tree Collapse file tree 2 files changed +22
-1
lines changed Original file line number Diff line number Diff line change @@ -399,6 +399,9 @@ def as_reward_model(cls: _T) -> _T:
399399 # Lazy import
400400 from vllm .model_executor .layers .pooler import DispatchPooler , Pooler
401401
402+ from .interfaces_base import default_pooling_type
403+
404+ @default_pooling_type ("ALL" )
402405 class ModelForReward (_create_pooling_model_cls (cls )):
403406 def _init_pooler (self , vllm_config : "VllmConfig" , prefix : str = "" ):
404407 pooler_config = vllm_config .model_config .pooler_config
Original file line number Diff line number Diff line change @@ -3587,8 +3587,26 @@ def _dummy_pooler_run(
35873587 hidden_states : torch .Tensor ,
35883588 ) -> PoolerOutput :
35893589 # Find the task that has the largest output for subsequent steps
3590+ supported_pooling_tasks = self .get_supported_pooling_tasks ()
3591+
3592+ if not supported_pooling_tasks :
3593+ if self .scheduler_config .chunked_prefill_enabled :
3594+ raise RuntimeError (
3595+ f"Model { self .model_config .model } does not support "
3596+ "any pooling tasks with chunked prefill enabled. "
3597+ "Please add --no-enable-chunked-prefill to your "
3598+ "config or CLI args. See "
3599+ "https://docs.vllm.ai/en/latest/models/pooling_models.html "
3600+ "to learn more." )
3601+ else :
3602+ raise RuntimeError (
3603+ f"Model { self .model_config .model } does not support "
3604+ "any pooling tasks. See "
3605+ "https://docs.vllm.ai/en/latest/models/pooling_models.html "
3606+ "to learn more." )
3607+
35903608 output_size = dict [PoolingTask , float ]()
3591- for task in self . get_supported_pooling_tasks () :
3609+ for task in supported_pooling_tasks :
35923610 # Run a full batch with each task to ensure none of them OOMs
35933611 output = self ._dummy_pooler_run_task (hidden_states , task )
35943612 output_size [task ] = sum (o .nbytes for o in output )
You can’t perform that action at this time.
0 commit comments