main
py 619 lines 21.6 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 # adapted from https://github.com/deepseek-ai/DeepSeek-VL2/blob/faf18023f24b962b32d9f0a2d89e402a8d383a78/deepseek_vl2/models/modeling_deepseek_vl_v2.py
5 """Inference-only Deepseek-VL2 model compatible with HuggingFace weights."""
6
7 import math
8 from collections.abc import Iterable, Mapping, Sequence
9 from typing import Annotated, Literal, TypeAlias
10
11 import torch
12 import torch.nn as nn
13 import torch.nn.functional as F
14 from einops import rearrange, repeat
15 from transformers import BatchFeature
16
17 from vllm.config import VllmConfig
18 from vllm.config.multimodal import MultiModalDummyOptions
19 from vllm.distributed import get_tensor_model_parallel_world_size
20 from vllm.inputs import MultiModalDataDict
21 from vllm.model_executor.layers.quantization import QuantizationConfig
22 from vllm.model_executor.models.transformers.utils import replace_linear_class
23 from vllm.multimodal import MULTIMODAL_REGISTRY
24 from vllm.multimodal.inputs import (
25 MultiModalFieldConfig,
26 MultiModalKwargsItems,
27 )
28 from vllm.multimodal.parse import (
29 ImageEmbeddingItems,
30 ImageProcessorItems,
31 ImageSize,
32 MultiModalDataItems,
33 )
34 from vllm.multimodal.processing import BaseDummyInputsBuilder
35 from vllm.multimodal.processing.processor import (
36 BaseMultiModalProcessor,
37 BaseProcessingInfo,
38 MultiModalProcessingResult,
39 ProcessorInputs,
40 PromptReplacement,
41 PromptUpdate,
42 TimingContext,
43 )
44 from vllm.sequence import IntermediateTensors
45 from vllm.tokenizers import cached_tokenizer_from_config
46 from vllm.transformers_utils.configs.deepseek_vl2 import (
47 DeepseekVLV2Config,
48 MlpProjectorConfig,
49 VisionEncoderConfig,
50 )
51 from vllm.utils.tensor_schema import TensorSchema, TensorShape
52 from vllm.utils.torch_utils import set_default_torch_dtype
53
54 from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP
55 from .utils import (
56 AutoWeightsLoader,
57 WeightsMapper,
58 init_vllm_registered_model,
59 maybe_prefix,
60 )
61
62 # The image token id may be various
63 _IMAGE_TOKEN = "<image>"
64
65
66 class DeepseekVL2ImagePixelInputs(TensorSchema):
67 """Dimensions:
68 - bnp: Batch size * number of images * number of patches
69 - p: Number of patches
70 - c: Number of channels (3)
71 - h: Height of each image
72 - w: Width of each image
73 """
74
75 type: Literal["pixel_values"]
76 data: Annotated[torch.Tensor, TensorShape("bnp", 3, "h", "w", dynamic_dims={"bnp"})]
77 images_spatial_crop: Annotated[torch.Tensor, TensorShape("bn", 2)]
78
79
80 class DeepseekVL2VImageEmbeddingInputs(TensorSchema):
81 """Dimensions:
82 - bn: Batch size * number of images
83 - f: Image feature size
84 - h: Hidden size (must match language model backbone)
85 """
86
87 type: Literal["image_embeds"]
88 data: Annotated[torch.Tensor | list[torch.Tensor], TensorShape("bn", "f", "h")]
89
90
91 DeepseekVL2ImageInputs: TypeAlias = (
92 DeepseekVL2ImagePixelInputs | DeepseekVL2VImageEmbeddingInputs
93 )
94
95
96 class MlpProjector(nn.Module):
97 def __init__(self, cfg: MlpProjectorConfig):
98 super().__init__()
99
100 self.cfg = cfg
101 self.projector_type = cfg.projector_type
102 assert not cfg.token_pooling, "Token pooling is not supported currently."
103
104 if self.projector_type == "downsample_mlp_gelu":
105 mlp_depth = cfg.depth
106 mlp_ratio = cfg.mlp_ratio
107 modules = [
108 nn.Linear(
109 cfg.input_dim * cfg.downsample_ratio * cfg.downsample_ratio,
110 cfg.n_embed * mlp_ratio,
111 )
112 ]
113 for _ in range(1, mlp_depth - 1):
114 modules.append(nn.GELU())
115 modules.append(
116 nn.Linear(cfg.n_embed * mlp_ratio, cfg.n_embed * mlp_ratio)
117 )
118 modules.append(nn.GELU())
119 modules.append(nn.Linear(cfg.n_embed * mlp_ratio, cfg.n_embed))
120 modules = nn.Sequential(*modules)
121 elif self.projector_type == "linear":
122 modules = nn.Linear(cfg.input_dim, cfg.n_embed)
123 else:
124 raise NotImplementedError(
125 f"Unsupported projector type: {cfg.projector_type}"
126 )
127
128 self.layers = modules
129
130 def forward(self, x):
131 bs, hw, input_dim = x.shape
132 if self.projector_type == "downsample_mlp_gelu":
133 h = w = int((hw) ** 0.5)
134 """compute padding"""
135 if h % self.cfg.downsample_ratio:
136 pad = self.cfg.downsample_ratio - h % self.cfg.downsample_ratio
137 else:
138 pad = 0
139 x = x.reshape(bs, h, w, input_dim)
140 if pad > 0:
141 x = F.pad(x, (0, 0, 0, pad, 0, pad), "constant", 0)
142 """4 to 1 concat"""
143 x = x.permute(0, 3, 1, 2) # B, C, H, W
144 x = F.unfold(
145 x,
146 kernel_size=self.cfg.downsample_ratio,
147 stride=self.cfg.downsample_ratio,
148 padding=0,
149 ) # B, C*4, HW // 4
150 x = x.permute(0, 2, 1)
151
152 return self.layers(x)
153
154
155 class DeepseekVL2ProcessingInfo(BaseProcessingInfo):
156 def get_hf_config(self):
157 return self.ctx.get_hf_config(DeepseekVLV2Config)
158
159 def get_hf_processor(self, **kwargs: object):
160 return self.ctx.get_hf_processor(**kwargs)
161
162 def get_supported_mm_limits(self) -> Mapping[str, int | None]:
163 return {"image": None}
164
165 def get_num_image_tokens(
166 self, *, image_width: int, image_height: int, cropping: bool = True
167 ) -> int:
168 hf_processor = self.get_hf_processor()
169 image_size = hf_processor.image_size
170 patch_size = hf_processor.patch_size
171 downsample_ratio = hf_processor.downsample_ratio
172
173 if cropping:
174 best_width, best_height = hf_processor.select_best_resolution(
175 (image_width, image_height)
176 )
177 num_width_tiles, num_height_tiles = (
178 best_width // image_size,
179 best_height // image_size,
180 )
181 else:
182 num_width_tiles = num_height_tiles = 1
183
184 h = w = math.ceil((image_size // patch_size) / downsample_ratio)
185
186 global_views_tokens = h * (w + 1)
187 local_views_tokens = (num_height_tiles * h) * (num_width_tiles * w + 1)
188 return global_views_tokens + local_views_tokens + 1
189
190 def get_image_size_with_most_features(self) -> ImageSize:
191 hf_config = self.get_hf_config()
192 candidate_resolutions = hf_config.candidate_resolutions
193 height, width = max(
194 candidate_resolutions,
195 key=lambda x: self.get_num_image_tokens(
196 image_width=x[1], image_height=x[0]
197 ),
198 )
199 return ImageSize(width=width, height=height)
200
201
202 class DeepseekVL2DummyInputsBuilder(BaseDummyInputsBuilder[DeepseekVL2ProcessingInfo]):
203 def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
204 num_images = mm_counts.get("image", 0)
205
206 processor = self.info.get_hf_processor()
207 image_token = processor.image_token
208
209 return image_token * num_images
210
211 def get_dummy_mm_data(
212 self,
213 seq_len: int,
214 mm_counts: Mapping[str, int],
215 mm_options: MultiModalDummyOptions,
216 ) -> MultiModalDataDict:
217 max_image_size = self.info.get_image_size_with_most_features()
218
219 return {
220 "image": self._get_dummy_images(
221 width=max_image_size.width,
222 height=max_image_size.height,
223 num_images=mm_counts.get("image", 0),
224 overrides=mm_options.get("image"),
225 )
226 }
227
228
229 class DeepseekVL2MultiModalProcessor(
230 BaseMultiModalProcessor[DeepseekVL2ProcessingInfo]
231 ):
232 def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str:
233 return self.dummy_inputs.get_dummy_text(mm_counts)
234
235 def _postprocess_hf_mm_data(
236 self,
237 mm_data: Mapping[str, object],
238 hf_processor_mm_kwargs: Mapping[str, object],
239 processed_data: BatchFeature,
240 ) -> BatchFeature:
241 if not mm_data:
242 return processed_data
243
244 processed_data["num_patches"] = (
245 processed_data["images_spatial_crop"].prod(-1) + 1
246 )
247
248 return processed_data
249
250 def _get_mm_fields_config(
251 self,
252 hf_inputs: BatchFeature,
253 hf_processor_mm_kwargs: Mapping[str, object],
254 ) -> Mapping[str, MultiModalFieldConfig]:
255 num_patches = hf_inputs.get("num_patches", torch.empty(0))
256
257 return dict(
258 pixel_values=MultiModalFieldConfig.flat_from_sizes("image", num_patches),
259 images_spatial_crop=MultiModalFieldConfig.batched(
260 "image", keep_on_cpu=True
261 ),
262 image_embeds=MultiModalFieldConfig.batched("image"),
263 )
264
265 def _get_prompt_updates(
266 self,
267 mm_items: MultiModalDataItems,
268 hf_processor_mm_kwargs: Mapping[str, object],
269 out_mm_kwargs: MultiModalKwargsItems,
270 ) -> Sequence[PromptUpdate]:
271 hf_processor = self.info.get_hf_processor(**hf_processor_mm_kwargs)
272
273 image_token_id = hf_processor.image_token_id
274 assert isinstance(image_token_id, int)
275
276 def get_replacement_deepseek_vl2(item_idx: int):
277 images = mm_items.get_items(
278 "image", (ImageEmbeddingItems, ImageProcessorItems)
279 )
280
281 if isinstance(images, ImageEmbeddingItems):
282 num_image_tokens = images.get_feature_size(item_idx)
283 else:
284 assert isinstance(images, ImageProcessorItems)
285 image_size = images.get_image_size(item_idx)
286
287 num_image_tokens = self.info.get_num_image_tokens(
288 image_width=image_size.width,
289 image_height=image_size.height,
290 cropping=len(images) <= 2,
291 )
292 return [image_token_id] * num_image_tokens
293
294 return [
295 PromptReplacement(
296 modality="image",
297 target=[image_token_id],
298 replacement=get_replacement_deepseek_vl2,
299 )
300 ]
301
302 def _cached_apply_hf_processor(
303 self,
304 inputs: ProcessorInputs,
305 timing_ctx: TimingContext,
306 ) -> MultiModalProcessingResult:
307 # The processor logic is different for len(images) <= 2 vs > 2
308 # Since the processing cache assumes that the processor output is
309 # invariant of how many images are passed per prompt, we only
310 # perform caching for the most common case
311 if inputs.mm_data_items.get_count("image", strict=False) > 2:
312 return self._apply_hf_processor(inputs, timing_ctx)
313
314 return super()._cached_apply_hf_processor(inputs, timing_ctx)
315
316
317 @MULTIMODAL_REGISTRY.register_processor(
318 DeepseekVL2MultiModalProcessor,
319 info=DeepseekVL2ProcessingInfo,
320 dummy_inputs=DeepseekVL2DummyInputsBuilder,
321 )
322 class DeepseekVLV2ForCausalLM(nn.Module, SupportsMultiModal, SupportsPP):
323 hf_to_vllm_mapper = WeightsMapper(
324 orig_to_new_prefix={
325 "language.": "language_model.",
326 }
327 )
328
329 @classmethod
330 def get_placeholder_str(cls, modality: str, i: int) -> str | None:
331 if modality.startswith("image"):
332 return "<image>"
333
334 raise ValueError("Only image modality is supported")
335
336 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
337 super().__init__()
338 config: DeepseekVLV2Config = vllm_config.model_config.hf_config
339 quant_config = vllm_config.quant_config
340 multimodal_config = vllm_config.model_config.multimodal_config
341
342 self.config = config
343 self.multimodal_config = multimodal_config
344
345 self.vision_config = config.vision_config
346 self.projector_config = config.projector_config
347 self.text_config = config.text_config
348
349 model_config = vllm_config.model_config
350 tokenizer = cached_tokenizer_from_config(model_config)
351 self.image_token_id: int = tokenizer.vocab[_IMAGE_TOKEN]
352
353 with self._mark_tower_model(vllm_config, "image"):
354 self.vision = self._init_vision_module(
355 self.vision_config, quant_config, maybe_prefix(prefix, "vision")
356 )
357
358 self.projector = MlpProjector(self.projector_config)
359 self.tile_tag = config.tile_tag
360 self.global_view_pos = config.global_view_pos
361
362 # special token for image token sequence format
363 embed_std = 1 / torch.sqrt(
364 torch.tensor(self.projector_config.n_embed, dtype=torch.float32)
365 )
366 if self.tile_tag == "2D":
367 # <|view_seperator|>, <|\n|>
368 self.image_newline = nn.Parameter(
369 torch.randn(self.projector_config.n_embed) * embed_std
370 )
371 # This is a typo in original implementation
372 self.view_seperator = nn.Parameter(
373 torch.randn(self.projector_config.n_embed) * embed_std
374 )
375 else:
376 raise ValueError(
377 f"Only 2D tile_tag is supported currently, got: {self.tile_tag}"
378 )
379
380 with self._mark_language_model(vllm_config):
381 self.language_model = init_vllm_registered_model(
382 vllm_config=vllm_config,
383 hf_config=self.text_config,
384 prefix=maybe_prefix(prefix, "language"),
385 )
386
387 self.make_empty_intermediate_tensors = (
388 self.language_model.make_empty_intermediate_tensors
389 )
390
391 def _get_parent_and_attr(self, root: torch.nn.Module, dotted_name: str):
392 """Return (parent_module, final_attr_name) for a dotted module path."""
393 names = dotted_name.split(".")
394 parent = root
395 for n in names[:-1]:
396 parent = getattr(parent, n)
397 return parent, names[-1]
398
399 # patch for timm ViT instance to support tensor parallel
400 def patch_vit_for_tp(
401 self, vit: torch.nn.Module, quant_config: QuantizationConfig | None
402 ):
403 try:
404 import timm
405 except ImportError as e:
406 raise ImportError("Please install timm") from e
407
408 for name, module in vit.named_modules():
409 if isinstance(module, nn.Linear):
410 parent, attr_name = self._get_parent_and_attr(vit, name)
411 if isinstance(parent, timm.layers.Mlp) and attr_name == "fc1":
412 new_linear = replace_linear_class(
413 module, "colwise", quant_config, prefix=name
414 )
415 setattr(parent, attr_name, new_linear)
416 elif isinstance(parent, timm.layers.Mlp) and attr_name == "fc2":
417 new_linear = replace_linear_class(
418 module, "rowwise", quant_config, prefix=name
419 )
420 setattr(parent, attr_name, new_linear)
421
422 return vit
423
424 def _init_vision_module(
425 self,
426 vision_config: VisionEncoderConfig,
427 quant_config: QuantizationConfig | None,
428 prefix: str = "",
429 ) -> nn.Module:
430 # TODO: refactor vision model through timm wrapper from transformers
431 try:
432 import timm
433 except ImportError as e:
434 raise ImportError("Please install timm") from e
435
436 with set_default_torch_dtype(torch.float16):
437 model = timm.create_model(
438 "vit_so400m_patch14_siglip_384.webli",
439 pretrained=False,
440 num_classes=0,
441 dynamic_img_size=True,
442 dynamic_img_pad=True,
443 )
444
445 if get_tensor_model_parallel_world_size() > 1:
446 model = self.patch_vit_for_tp(model, quant_config)
447
448 model = model.to(dtype=torch.get_default_dtype())
449 return model
450
451 def _parse_and_validate_image_input(
452 self, **kwargs: object
453 ) -> DeepseekVL2ImageInputs | None:
454 pixel_values = kwargs.pop("pixel_values", None)
455 images_spatial_crop = kwargs.pop("images_spatial_crop", None)
456 image_embeds = kwargs.pop("image_embeds", None)
457
458 if pixel_values is None and image_embeds is None:
459 return None
460
461 if pixel_values is not None:
462 expected_h = expected_w = self.vision_config.image_size
463 return DeepseekVL2ImagePixelInputs(
464 type="pixel_values",
465 data=pixel_values,
466 images_spatial_crop=images_spatial_crop,
467 resolve_bindings={
468 "h": expected_h,
469 "w": expected_w,
470 },
471 )
472
473 if image_embeds is not None:
474 return DeepseekVL2VImageEmbeddingInputs(
475 type="image_embeds",
476 data=image_embeds,
477 )
478
479 raise AssertionError("This line should be unreachable.")
480
481 def _pixel_values_to_embedding(
482 self,
483 pixel_values: torch.Tensor,
484 images_spatial_crop: torch.Tensor,
485 ) -> list[torch.Tensor]:
486 # [batch_all_tiles, vit_seq_len, c]
487 images_feature = self.vision.forward_features(pixel_values)
488
489 # [batch_all_tiles, hw, D]
490 images_embeds = self.projector(images_feature)
491
492 _, hw, n_dim = images_embeds.shape
493 h = w = int(hw**0.5)
494
495 # fill image token based on self.tile_tag & self.global_view_pos
496 tile_index = 0
497 vision_embeddings = []
498 for jdx in range(images_spatial_crop.size(0)):
499 # extra global & local features
500 num_width_tiles, num_height_tiles = images_spatial_crop[jdx]
501 if num_width_tiles == 0 or num_height_tiles == 0:
502 break
503 num_tiles_in_image = num_width_tiles * num_height_tiles
504
505 # [hw, D]
506 global_features = images_embeds[tile_index]
507
508 # [num_height_tiles * num_width_tiles, hw, D]
509 local_features = images_embeds[
510 tile_index + 1 : tile_index + 1 + num_tiles_in_image
511 ]
512 tile_index += num_tiles_in_image + 1
513
514 # format global and local features
515 # ----------------- global view add newline -----------------
516 # [hw, D] -> [h, w, D]
517 global_features = global_features.view(h, w, n_dim)
518
519 # [D] -> [h, 1, D]
520 new_lines_in_global = repeat(self.image_newline, "d -> h 1 d", h=h)
521
522 # cat([h, w, D], [h, 1, D], dim=1) -> [h, w + 1, D]
523 global_features = torch.cat([global_features, new_lines_in_global], dim=1)
524
525 # [h, w + 1, D] -> [h * (w + 1), D]
526 global_features = global_features.view(-1, n_dim)
527
528 # ----------------- local view add newline -----------------
529 # [num_height_tiles * num_width_tiles, h * w, D] ->
530 # [num_height_tiles * h, num_width_tiles * w, D]
531 local_features = rearrange(
532 local_features,
533 "(th tw) (h w) d -> (th h) (tw w) d",
534 th=num_height_tiles,
535 tw=num_width_tiles,
536 h=h,
537 w=w,
538 )
539
540 # [D] -> [num_height_tiles * h, 1, D]
541 new_lines_in_local = repeat(
542 self.image_newline, "d -> (th h) 1 d", th=num_height_tiles, h=h
543 )
544
545 # [num_height_tiles * h, num_width_tiles * w + 1, D]
546 local_features = torch.cat([local_features, new_lines_in_local], dim=1)
547
548 # [num_height_tiles * h, num_width_tiles * w + 1, D]
549 # --> [(num_height_tiles * h) * (num_width_tiles * w + 1), D]
550 local_features = local_features.view(-1, n_dim)
551
552 # merge global and local tiles
553 if self.global_view_pos == "head":
554 global_local_features = torch.cat(
555 [
556 global_features,
557 self.view_seperator[None, :],
558 local_features,
559 ]
560 )
561 else:
562 global_local_features = torch.cat(
563 [
564 local_features,
565 self.view_seperator[None, :],
566 global_features,
567 ]
568 )
569
570 vision_embeddings.append(global_local_features)
571 return vision_embeddings
572
573 def _process_image_input(
574 self, image_input: DeepseekVL2ImageInputs
575 ) -> torch.Tensor | list[torch.Tensor]:
576 if image_input["type"] == "image_embeds":
577 return image_input["data"]
578
579 pixel_values = image_input["data"]
580 images_spatial_crop = image_input["images_spatial_crop"]
581
582 return self._pixel_values_to_embedding(
583 pixel_values=pixel_values, images_spatial_crop=images_spatial_crop
584 )
585
586 def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:
587 image_input = self._parse_and_validate_image_input(**kwargs)
588 if image_input is None:
589 return []
590 vision_embeddings = self._process_image_input(image_input)
591 return vision_embeddings
592
593 def forward(
594 self,
595 input_ids: torch.Tensor | None,
596 positions: torch.Tensor,
597 intermediate_tensors: IntermediateTensors | None = None,
598 inputs_embeds: torch.Tensor | None = None,
599 **kwargs: object,
600 ):
601 if intermediate_tensors is not None:
602 inputs_embeds = None
603
604 hidden_states = self.language_model(
605 input_ids, positions, intermediate_tensors, inputs_embeds=inputs_embeds
606 )
607
608 return hidden_states
609
610 def compute_logits(
611 self,
612 hidden_states: torch.Tensor,
613 ) -> torch.Tensor | None:
614 return self.language_model.compute_logits(hidden_states)
615
616 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
617 loader = AutoWeightsLoader(self)
618 autoloaded_weights = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
619 return autoloaded_weights