| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Inference-only Jurassic model.""" |
| 4 | |
| 5 | import torch |
| 6 | |
| 7 | from vllm.config import VllmConfig |
| 8 | from vllm.logger import init_logger |
| 9 | from vllm.model_executor.layers.activation import get_act_fn |
| 10 | from vllm.model_executor.layers.linear import ColumnParallelLinear |
| 11 | |
| 12 | from .interfaces import supports_pp |
| 13 | from .step3_vl import Step3VLForConditionalGeneration |
| 14 | from .step_vl import PerceptionEncoder |
| 15 | from .utils import WeightsMapper, init_vllm_registered_model, maybe_prefix |
| 16 | from .vision import run_dp_sharded_vision_model |
| 17 | |
| 18 | logger = init_logger(__name__) |
| 19 | |
| 20 | |
| 21 | class Step3p7ForConditionalGeneration(Step3VLForConditionalGeneration): |
| 22 | hf_to_vllm_mapper = WeightsMapper( |
| 23 | orig_to_new_prefix={ |
| 24 | "model.vision_model.": "vision_model.", |
| 25 | "model.vit_large_projector.": "vit_large_projector.", |
| 26 | "model.vit_large_projector": "vit_large_projector", |
| 27 | "model.language_model.": "language_model.model.", |
| 28 | "model.language_model": "language_model.model", |
| 29 | "model.": "language_model.model.", |
| 30 | "lm_head.": "language_model.lm_head.", |
| 31 | "lm_head": "language_model.lm_head", |
| 32 | }, |
| 33 | orig_to_new_substr={ |
| 34 | ".attn.in_proj_weight": ".attn.qkv_proj.weight", |
| 35 | ".attn.in_proj_bias": ".attn.qkv_proj.bias", |
| 36 | ".mlp.c_fc": ".mlp.fc1", |
| 37 | ".mlp.c_proj": ".mlp.fc2", |
| 38 | }, |
| 39 | ) |
| 40 | |
| 41 | def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: |
| 42 | super(Step3VLForConditionalGeneration, self).__init__() |
| 43 | |
| 44 | config = vllm_config.model_config.hf_config |
| 45 | multimodal_config = vllm_config.model_config.get_multimodal_config() |
| 46 | quant_config = vllm_config.quant_config |
| 47 | |
| 48 | self.config = config |
| 49 | self.multimodal_config = multimodal_config |
| 50 | self.use_data_parallel = multimodal_config.mm_encoder_tp_mode == "data" |
| 51 | |
| 52 | with self._mark_tower_model(vllm_config, "image"): |
| 53 | self.vision_model = PerceptionEncoder( |
| 54 | config.vision_config, |
| 55 | get_act_fn(config.vision_config.hidden_act), |
| 56 | quant_config=quant_config, |
| 57 | prefix=maybe_prefix(prefix, "vision_model"), |
| 58 | ) |
| 59 | self.vit_large_projector = ColumnParallelLinear( |
| 60 | config.vision_config.width * 4, |
| 61 | config.text_config.hidden_size, |
| 62 | bias=config.projector_bias, |
| 63 | gather_output=True, |
| 64 | quant_config=quant_config, |
| 65 | prefix=maybe_prefix(prefix, "vit_large_projector"), |
| 66 | disable_tp=self.use_data_parallel, |
| 67 | ) |
| 68 | |
| 69 | with self._mark_language_model(vllm_config): |
| 70 | language_model = init_vllm_registered_model( |
| 71 | vllm_config=vllm_config, |
| 72 | hf_config=config.text_config, |
| 73 | prefix=maybe_prefix(prefix, "language_model"), |
| 74 | ) |
| 75 | self.language_model = language_model |
| 76 | assert supports_pp(language_model) |
| 77 | |
| 78 | self.make_empty_intermediate_tensors = ( |
| 79 | language_model.make_empty_intermediate_tensors |
| 80 | ) |
| 81 | |
| 82 | def _get_vision_model_output( |
| 83 | self, input_tensor: torch.Tensor | None |
| 84 | ) -> torch.Tensor | None: |
| 85 | if input_tensor is None: |
| 86 | return None |
| 87 | if self.use_data_parallel: |
| 88 | return run_dp_sharded_vision_model(input_tensor, self.vision_model) |
| 89 | return self.vision_model(input_tensor) |
| 90 | |
| 91 | def _process_image_features(self, image_features: torch.Tensor) -> torch.Tensor: |
| 92 | image_features, _ = self.vit_large_projector(image_features) |
| 93 | return image_features |