Skip to content

[Model] Skip unused Jina V5 output layers #52037

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
noooop merged 2 commits into from
Aug 13, 2026
+21 −3
Merged

[Model] Skip unused Jina V5 output layers #52037

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
c60ad48
[Model] Skip unused Jina V5 output layers
BabyDrangoner Aug 12, 2026
2f5766a
[Model] Remove unnecessary Jina V5 regression test
BabyDrangoner Aug 13, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
24 changes: 21 additions & 3 deletions vllm/model_executor/models/jina.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,8 @@
from torch import nn

from vllm.config import VllmConfig
from vllm.model_executor.layers.logits_processor import LogitsProcessor
from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
from vllm.sequence import IntermediateTensors
from vllm.tasks import PoolingTask
from vllm.transformers_utils.repo_utils import get_hf_file_bytes
Expand All @@ -26,7 +28,13 @@
from .interfaces_base import VllmModelForPooling
from .llama import LlamaForCausalLM
from .qwen3 import Qwen3ForCausalLM, Qwen3Model
from .utils import AutoWeightsLoader, WeightsMapper, maybe_prefix
from .utils import (
AutoWeightsLoader,
StageMissingLayer,
WeightsMapper,
maybe_prefix,
no_init_weights,
)

logger = logging.getLogger(__name__)

Expand Down Expand Up @@ -267,7 +275,12 @@ class JinaEmbeddingsV5DecoderModel(Qwen3ForCausalLM, VllmModelForPooling):
)

def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
super().__init__(vllm_config=vllm_config, prefix=prefix)
with no_init_weights(
self,
lambda mod: StageMissingLayer("output", mod),
targets=(LogitsProcessor, ParallelLMHead),
):
super().__init__(vllm_config=vllm_config, prefix=prefix)
_setup_jina_v5_task_and_pooler(self, vllm_config)

def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
Expand All @@ -289,7 +302,12 @@ class JinaEmbeddingsV5EncoderModel(LlamaForCausalLM, VllmModelForPooling):
)

def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
super().__init__(vllm_config=vllm_config, prefix=prefix)
with no_init_weights(
self,
lambda mod: StageMissingLayer("output", mod),
targets=(LogitsProcessor, ParallelLMHead),
):
super().__init__(vllm_config=vllm_config, prefix=prefix)
_setup_jina_v5_task_and_pooler(self, vllm_config)

def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
Expand Down
Loading