[Model] Skip unused Jina V5 output layers (#52037)
Signed-off-by: kiroxu <148877251+BabyDrangoner@users.noreply.github.com> Co-authored-by: Codex <noreply@openai.com>
kiroxu committed
Aug 13, 2026 at 14:14 UTC
89c8401c8aeb53ccee87060e913514068e0a1e82
1 file changed
+21
-3
vllm/model_executor/models/jina.py
+21
-3
@@ -11,6 +11,8 @@ from safetensors.torch import load as safetensors_load
11
from torch import nn
12
13
from vllm.config import VllmConfig
14
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
15
+from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
16
from vllm.sequence import IntermediateTensors
17
from vllm.tasks import PoolingTask
18
from vllm.transformers_utils.repo_utils import get_hf_file_bytes
@@ -26,7 +28,13 @@ from .interfaces import SupportsLateInteraction
28
from .interfaces_base import VllmModelForPooling
29
from .llama import LlamaForCausalLM
30
from .qwen3 import Qwen3ForCausalLM, Qwen3Model
29
-from .utils import AutoWeightsLoader, WeightsMapper, maybe_prefix
31
+from .utils import (
32
+ AutoWeightsLoader,
33
+ StageMissingLayer,
34
+ WeightsMapper,
35
+ maybe_prefix,
36
+ no_init_weights,
37
+)
38
39
logger = logging.getLogger(__name__)
40
@@ -267,7 +275,12 @@ class JinaEmbeddingsV5DecoderModel(Qwen3ForCausalLM, VllmModelForPooling):
275
)
276
277
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
270
- super().__init__(vllm_config=vllm_config, prefix=prefix)
278
+ with no_init_weights(
279
+ self,
280
+ lambda mod: StageMissingLayer("output", mod),
281
+ targets=(LogitsProcessor, ParallelLMHead),
282
+ ):
283
+ super().__init__(vllm_config=vllm_config, prefix=prefix)
284
_setup_jina_v5_task_and_pooler(self, vllm_config)
285
286
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
@@ -289,7 +302,12 @@ class JinaEmbeddingsV5EncoderModel(LlamaForCausalLM, VllmModelForPooling):
302
)
303
304
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
292
- super().__init__(vllm_config=vllm_config, prefix=prefix)
305
+ with no_init_weights(
306
+ self,
307
+ lambda mod: StageMissingLayer("output", mod),
308
+ targets=(LogitsProcessor, ParallelLMHead),
309
+ ):
310
+ super().__init__(vllm_config=vllm_config, prefix=prefix)
311
_setup_jina_v5_task_and_pooler(self, vllm_config)
312
313
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: