main
py 628 lines 22 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """HyperCLOVAX V2 (32B Think Model) Implementation.
4
5 This module contains the V2 architecture that uses Qwen2.5 Vision Transformer
6 instead of CLIP/SigLIP used in V1.
7
8 Supports:
9 - HyperCLOVAX-SEED-Think-32B: Vision + Text
10 """
11
12 from collections.abc import Iterable, Mapping, Sequence
13 from functools import partial
14 from typing import Annotated, Literal, TypedDict
15
16 import torch
17 import torch.nn as nn
18 from transformers import BatchFeature
19
20 from vllm.config import VllmConfig
21 from vllm.config.multimodal import MultiModalDummyOptions
22 from vllm.forward_context import set_forward_context
23 from vllm.inputs import MultiModalDataDict
24 from vllm.multimodal import MULTIMODAL_REGISTRY
25 from vllm.multimodal.inputs import (
26 MultiModalFieldConfig,
27 MultiModalKwargsItems,
28 )
29 from vllm.multimodal.parse import ImageSize, MultiModalDataItems
30 from vllm.multimodal.processing import (
31 BaseDummyInputsBuilder,
32 BaseMultiModalProcessor,
33 BaseProcessingInfo,
34 ProcessorInputs,
35 PromptReplacement,
36 PromptUpdate,
37 )
38 from vllm.sequence import IntermediateTensors
39 from vllm.utils.tensor_schema import TensorSchema, TensorShape
40
41 from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP
42 from .qwen2_5_vl import Qwen2_5_VisionTransformer
43 from .utils import (
44 AutoWeightsLoader,
45 WeightsMapper,
46 init_vllm_registered_model,
47 maybe_prefix,
48 )
49
50 # V2 (32B Think model) uses different tokens - retrieved from config at runtime
51 # These placeholder strings must match the chat template format exactly.
52 # The chat template produces: <|image_start|><|IMAGE_PAD|><|image_end|>
53 # Similar to Qwen2-VL's <|vision_start|><|image_pad|><|vision_end|> format.
54 V2_IMAGE_TOKEN: str = "<|image_start|><|IMAGE_PAD|><|image_end|>"
55 V2_VIDEO_TOKEN: str = "<|video_start|><|VIDEO_PAD|><|video_end|>"
56
57
58 class HCXVisionV2ImagePixelInputs(TensorSchema):
59 """V2 Image inputs using Qwen2.5-VL style grid_thw format.
60
61 Dimensions:
62 - np: Number of patches
63 - ni: Number of images
64 - cps: Number of channels * patch_size * patch_size
65 """
66
67 type: Literal["pixel_values"] = "pixel_values"
68 pixel_values: Annotated[torch.Tensor, TensorShape("np", "cps")]
69 image_grid_thw: Annotated[torch.Tensor, TensorShape("ni", 3)]
70
71
72 class HCXVisionV2ImageEmbeddingInputs(TensorSchema):
73 """V2 Image embedding inputs.
74
75 Dimensions:
76 - nf: Number of image features
77 - hs: Hidden size
78 - ni: Number of images
79 """
80
81 type: Literal["image_embeds"] = "image_embeds"
82 image_embeds: Annotated[torch.Tensor, TensorShape("nf", "hs")]
83 image_grid_thw: Annotated[torch.Tensor, TensorShape("ni", 3)]
84
85
86 HCXVisionV2ImageInputs = HCXVisionV2ImagePixelInputs | HCXVisionV2ImageEmbeddingInputs
87
88
89 class HCXVisionV2VideoPixelInputs(TensorSchema):
90 """V2 Video inputs using Qwen2.5-VL style grid_thw format.
91
92 Dimensions:
93 - np: Number of patches
94 - nv: Number of videos
95 - ctps: Number of channels * temporal_patch_size * patch_size * patch_size
96 """
97
98 type: Literal["pixel_values_videos"] = "pixel_values_videos"
99 pixel_values_videos: Annotated[torch.Tensor, TensorShape("np", "ctps")]
100 video_grid_thw: Annotated[torch.Tensor, TensorShape("nv", 3)]
101
102
103 class HCXVisionV2VideoEmbeddingInputs(TensorSchema):
104 """V2 Video embedding inputs.
105
106 Dimensions:
107 - nf: Number of video features
108 - hs: Hidden size
109 - nv: Number of videos
110 """
111
112 type: Literal["video_embeds"] = "video_embeds"
113 video_embeds: Annotated[torch.Tensor, TensorShape("nf", "hs")]
114 video_grid_thw: Annotated[torch.Tensor, TensorShape("nv", 3)]
115
116
117 HCXVisionV2VideoInputs = HCXVisionV2VideoPixelInputs | HCXVisionV2VideoEmbeddingInputs
118
119
120 class HCXVisionV2MultiModalInputs(TypedDict, total=False):
121 image: HCXVisionV2ImageInputs | None
122 video: HCXVisionV2VideoInputs | None
123
124
125 class HCXVisionV2ProcessingInfo(BaseProcessingInfo):
126 """Processing info for HyperCLOVAX V2 (32B Think model)."""
127
128 def get_supported_mm_limits(self) -> Mapping[str, int | None]:
129 return {"image": None, "video": None}
130
131 def get_num_image_tokens(
132 self,
133 *,
134 image_width: int,
135 image_height: int,
136 ) -> int:
137 hf_config = self.get_hf_config()
138 vision_config = hf_config.vision_config
139 patch_size = vision_config.patch_size
140 spatial_merge_size = vision_config.spatial_merge_size
141
142 grid_h = image_height // patch_size
143 grid_w = image_width // patch_size
144
145 return (grid_h * grid_w) // (spatial_merge_size**2)
146
147 def get_num_video_tokens(
148 self,
149 *,
150 video_width: int,
151 video_height: int,
152 num_frames: int,
153 ) -> int:
154 hf_config = self.get_hf_config()
155 vision_config = hf_config.vision_config
156 patch_size = vision_config.patch_size
157 temporal_patch_size = vision_config.temporal_patch_size
158 spatial_merge_size = vision_config.spatial_merge_size
159
160 grid_t = num_frames // temporal_patch_size
161 grid_h = video_height // patch_size
162 grid_w = video_width // patch_size
163
164 return (grid_t * grid_h * grid_w) // (spatial_merge_size**2)
165
166 def get_image_size_with_most_features(self) -> ImageSize:
167 hf_config = self.get_hf_config()
168 vision_config = hf_config.vision_config
169 # Use a reasonable default size
170 size = getattr(vision_config, "image_size", 448)
171 return ImageSize(width=size, height=size)
172
173 def get_max_image_tokens(self) -> int:
174 target_width, target_height = self.get_image_size_with_most_features()
175 return self.get_num_image_tokens(
176 image_width=target_width,
177 image_height=target_height,
178 )
179
180
181 class HCXVisionV2DummyInputsBuilder(BaseDummyInputsBuilder[HCXVisionV2ProcessingInfo]):
182 """Dummy inputs builder for HyperCLOVAX V2 memory profiling."""
183
184 def get_dummy_text(
185 self,
186 mm_counts: Mapping[str, int],
187 ) -> str:
188 num_images = mm_counts.get("image", 0)
189 num_videos = mm_counts.get("video", 0)
190 return V2_IMAGE_TOKEN * num_images + V2_VIDEO_TOKEN * num_videos
191
192 def get_dummy_mm_data(
193 self,
194 seq_len: int,
195 mm_counts: Mapping[str, int],
196 mm_options: MultiModalDummyOptions | None = None,
197 mm_processor_kwargs: Mapping[str, object] | None = None,
198 ) -> MultiModalDataDict:
199 target_width, target_height = self.info.get_image_size_with_most_features()
200 target_num_frames = 16 # Default for video
201
202 result: MultiModalDataDict = {
203 "image": self._get_dummy_images(
204 width=target_width,
205 height=target_height,
206 num_images=mm_counts.get("image", 0),
207 overrides=mm_options.get("image") if mm_options else None,
208 ),
209 "video": self._get_dummy_videos(
210 width=target_width,
211 height=target_height,
212 num_frames=target_num_frames,
213 num_videos=mm_counts.get("video", 0),
214 overrides=mm_options.get("video") if mm_options else None,
215 ),
216 }
217
218 return result
219
220
221 class HCXVisionV2MultiModalProcessor(
222 BaseMultiModalProcessor[HCXVisionV2ProcessingInfo]
223 ):
224 """Multimodal processor for HyperCLOVAX V2 (32B Think model)."""
225
226 def get_dummy_inputs(
227 self,
228 seq_len: int,
229 mm_counts: Mapping[str, int],
230 mm_options: MultiModalDummyOptions,
231 ) -> ProcessorInputs:
232 num_images = mm_counts.get("image", 0)
233 num_videos = mm_counts.get("video", 0)
234 prompt_text = V2_IMAGE_TOKEN * num_images + V2_VIDEO_TOKEN * num_videos
235
236 builder = self.dummy_inputs
237 dummy_mm_data = builder.get_dummy_mm_data(seq_len, mm_counts, mm_options)
238 dummy_mm_items = self.info.parse_mm_data(dummy_mm_data, validate=False)
239
240 tokenizer = self.info.get_tokenizer()
241 prompt = tokenizer.encode(
242 prompt_text,
243 **self.info.default_tok_params.get_encode_kwargs(),
244 )
245
246 return ProcessorInputs(prompt=prompt, mm_data_items=dummy_mm_items)
247
248 def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str:
249 return self.dummy_inputs.get_dummy_text(mm_counts)
250
251 def _get_prompt_updates(
252 self,
253 mm_items: MultiModalDataItems,
254 hf_processor_mm_kwargs: Mapping[str, object],
255 out_mm_kwargs: MultiModalKwargsItems,
256 ) -> Sequence[PromptUpdate]:
257 hf_config = self.info.get_hf_config()
258
259 # Use token IDs directly from config.
260 # This matches what get_dummy_inputs uses, ensuring consistency.
261 placeholder: dict[str, int] = {
262 "image": hf_config.image_token_id, # 128060 for <|IMAGE_PAD|>
263 "video": hf_config.video_token_id, # 128061 for <|VIDEO_PAD|>
264 }
265
266 merge_size = hf_config.vision_config.spatial_merge_size
267
268 def get_replacement_v2(
269 item_idx: int,
270 modality: str,
271 out_mm_kwargs: MultiModalKwargsItems,
272 ):
273 out_item = out_mm_kwargs[modality][item_idx]
274
275 if modality == "image":
276 grid_thw_elem = out_item.get("image_grid_thw")
277 if grid_thw_elem is not None:
278 # Access .data to get the actual tensor from MultiModalFieldElem
279 grid_thw = grid_thw_elem.data
280 assert isinstance(grid_thw, torch.Tensor)
281 # Qwen2.5-VL style calculation
282 h, w = grid_thw[1].item(), grid_thw[2].item()
283 num_tokens = (h * w) // (merge_size**2)
284 else:
285 # Fallback or error
286 raise ValueError("Missing image_grid_thw for V2 model")
287 elif modality == "video":
288 grid_thw_elem = out_item.get("video_grid_thw")
289 if grid_thw_elem is not None:
290 # Access .data to get the actual tensor from MultiModalFieldElem
291 grid_thw = grid_thw_elem.data
292 assert isinstance(grid_thw, torch.Tensor)
293 t, h, w = grid_thw[0].item(), grid_thw[1].item(), grid_thw[2].item()
294 num_tokens = (t * h * w) // (merge_size**2)
295 else:
296 raise ValueError("Missing video_grid_thw for V2 model")
297 else:
298 raise NotImplementedError(modality)
299
300 return [placeholder[modality]] * num_tokens
301
302 return [
303 PromptReplacement(
304 modality=modality,
305 target=[
306 placeholder[modality],
307 ],
308 replacement=partial(
309 get_replacement_v2,
310 modality=modality,
311 out_mm_kwargs=out_mm_kwargs,
312 ),
313 )
314 for modality in ("image", "video")
315 ]
316
317 def _get_mm_fields_config(
318 self,
319 hf_inputs: BatchFeature,
320 hf_processor_mm_kwargs: Mapping[str, object],
321 ) -> Mapping[str, MultiModalFieldConfig]:
322 # HyperCLOVAX V2 uses Qwen2.5-VL style flattened pixel values where
323 # pixel_values has shape (num_patches, channels*patch_size*patch_size)
324 # while image_grid_thw has shape (num_images, 3).
325 # We need to use flat_from_sizes to correctly handle this mismatch.
326 hf_config = self.info.get_hf_config()
327 spatial_merge_size = hf_config.vision_config.spatial_merge_size
328
329 image_grid_thw = hf_inputs.get("image_grid_thw", torch.empty((0, 3)))
330 image_pixel_grid_sizes = image_grid_thw.prod(-1)
331 image_embed_grid_sizes = (
332 image_pixel_grid_sizes // spatial_merge_size // spatial_merge_size
333 )
334
335 video_grid_thw = hf_inputs.get("video_grid_thw", torch.empty((0, 3)))
336 video_pixel_grid_sizes = video_grid_thw.prod(-1)
337 video_embed_grid_sizes = (
338 video_pixel_grid_sizes // spatial_merge_size // spatial_merge_size
339 )
340
341 return dict(
342 pixel_values=MultiModalFieldConfig.flat_from_sizes(
343 "image", image_pixel_grid_sizes
344 ),
345 image_embeds=MultiModalFieldConfig.flat_from_sizes(
346 "image", image_embed_grid_sizes
347 ),
348 image_grid_thw=MultiModalFieldConfig.batched("image", keep_on_cpu=True),
349 pixel_values_videos=MultiModalFieldConfig.flat_from_sizes(
350 "video", video_pixel_grid_sizes
351 ),
352 video_embeds=MultiModalFieldConfig.flat_from_sizes(
353 "video", video_embed_grid_sizes
354 ),
355 video_grid_thw=MultiModalFieldConfig.batched("video", keep_on_cpu=True),
356 )
357
358
359 @MULTIMODAL_REGISTRY.register_processor(
360 HCXVisionV2MultiModalProcessor,
361 info=HCXVisionV2ProcessingInfo,
362 dummy_inputs=HCXVisionV2DummyInputsBuilder,
363 )
364 class HCXVisionV2ForCausalLM(nn.Module, SupportsMultiModal, SupportsPP):
365 """HyperCLOVAX-SEED Vision-Language Model (V2 architecture).
366
367 Supports:
368 - HyperCLOVAX-SEED-Think-32B: Vision + Text
369
370 Uses Qwen2.5 Vision Transformer as the vision encoder.
371 """
372
373 packed_modules_mapping = {
374 "qkv_proj": ["q_proj", "k_proj", "v_proj"],
375 "gate_up_proj": ["gate_proj", "up_proj"],
376 "qkv": ["qkv"], # For vision tower
377 }
378
379 # Weight mapping for loading HuggingFace checkpoints
380 # NOTE: Order matters! Ignores (None) should come before renames to prevent
381 # partial matches
382 hf_to_vllm_mapper = WeightsMapper(
383 orig_to_new_prefix={
384 "model.": "", # Remove model. prefix if present
385 "vision_model.": "visual.", # HF uses vision_model, we use visual
386 },
387 orig_to_new_substr={
388 # Ignore modules not implemented in vLLM
389 "discrete_vision_model": None, # TextAlignedTokenizer
390 },
391 )
392
393 def __init__(
394 self,
395 *,
396 vllm_config: VllmConfig,
397 prefix: str = "",
398 ) -> None:
399 super().__init__()
400
401 config = vllm_config.model_config.hf_config
402 quant_config = vllm_config.quant_config
403
404 # Text config
405 text_config = config.text_config
406 if text_config.model_type in ["gpt2", "hyperclovax", "llama"]:
407 text_config._attn_implementation = "sdpa"
408 if text_config.model_type != "hyperclovax":
409 text_config.logits_scaling = 1.0
410
411 # Vision config
412 vision_config = config.vision_config
413
414 self.config = config
415 self.vision_config = vision_config
416 self.text_config = text_config
417 self.vllm_config = vllm_config
418
419 # Linear projector (vision_hidden_size -> text_hidden_size)
420 # For V2 model: mm_projector_type is "linear"
421 vision_hidden_size = vision_config.hidden_size
422 text_hidden_size = text_config.hidden_size
423
424 # Check if out_hidden_size is defined (Qwen2.5-VL style)
425 # The merger in Qwen2.5 VisionTransformer handles projection to out_hidden_size
426 if hasattr(vision_config, "out_hidden_size"):
427 out_hidden = vision_config.out_hidden_size
428 else:
429 out_hidden = vision_hidden_size
430
431 with self._mark_tower_model(vllm_config, {"image", "video"}):
432 self.visual = Qwen2_5_VisionTransformer(
433 vision_config=vision_config,
434 norm_eps=getattr(config, "rms_norm_eps", 1e-6),
435 quant_config=quant_config,
436 prefix=maybe_prefix(prefix, "visual"),
437 )
438 self.mm_projector = nn.Linear(out_hidden, text_hidden_size)
439
440 with self._mark_language_model(vllm_config):
441 self.language_model = init_vllm_registered_model(
442 vllm_config=vllm_config,
443 hf_config=text_config,
444 prefix=maybe_prefix(prefix, "language_model"),
445 )
446
447 self.make_empty_intermediate_tensors = (
448 self.language_model.make_empty_intermediate_tensors
449 )
450
451 @classmethod
452 def get_placeholder_str(cls, modality: str, i: int) -> str | None:
453 if modality.startswith("image"):
454 return V2_IMAGE_TOKEN
455 if modality.startswith("video"):
456 return V2_VIDEO_TOKEN
457
458 raise ValueError("Only image or video modality is supported")
459
460 def _parse_and_validate_image_input(
461 self,
462 **kwargs: object,
463 ) -> HCXVisionV2ImageInputs | None:
464 pixel_values = kwargs.pop("pixel_values", None)
465 image_embeds = kwargs.pop("image_embeds", None)
466 image_grid_thw = kwargs.pop("image_grid_thw", None)
467
468 if pixel_values is None and image_embeds is None:
469 return None
470
471 if pixel_values is not None:
472 return HCXVisionV2ImagePixelInputs(
473 pixel_values=pixel_values,
474 image_grid_thw=image_grid_thw,
475 )
476
477 if image_embeds is not None:
478 return HCXVisionV2ImageEmbeddingInputs(
479 image_embeds=image_embeds,
480 image_grid_thw=image_grid_thw,
481 )
482
483 return None
484
485 def _parse_and_validate_video_input(
486 self,
487 **kwargs: object,
488 ) -> HCXVisionV2VideoInputs | None:
489 pixel_values_videos = kwargs.pop("pixel_values_videos", None)
490 video_embeds = kwargs.pop("video_embeds", None)
491 video_grid_thw = kwargs.pop("video_grid_thw", None)
492
493 if pixel_values_videos is None and video_embeds is None:
494 return None
495
496 if pixel_values_videos is not None:
497 return HCXVisionV2VideoPixelInputs(
498 pixel_values_videos=pixel_values_videos,
499 video_grid_thw=video_grid_thw,
500 )
501
502 if video_embeds is not None:
503 return HCXVisionV2VideoEmbeddingInputs(
504 video_embeds=video_embeds,
505 video_grid_thw=video_grid_thw,
506 )
507
508 return None
509
510 def _process_image_input(
511 self,
512 image_input: HCXVisionV2ImageInputs,
513 ) -> tuple[torch.Tensor, ...]:
514 """Process images through Qwen2.5 ViT and projector."""
515 grid_thw = image_input["image_grid_thw"]
516 assert grid_thw.ndim == 2
517 grid_thw_list = grid_thw.tolist()
518
519 if image_input["type"] == "image_embeds":
520 image_embeds = image_input["image_embeds"].type(self.visual.dtype)
521 else:
522 pixel_values = image_input["pixel_values"]
523 with set_forward_context(None, self.vllm_config):
524 image_embeds = self.visual(pixel_values, grid_thw=grid_thw_list)
525
526 # Apply projector
527 image_embeds = self.mm_projector(image_embeds)
528
529 # Split concatenated embeddings for each image
530 merge_size = self.visual.spatial_merge_size
531 sizes = (grid_thw.prod(-1) // merge_size // merge_size).tolist()
532 return image_embeds.split(sizes)
533
534 def _process_video_input(
535 self,
536 video_input: HCXVisionV2VideoInputs,
537 ) -> tuple[torch.Tensor, ...]:
538 """Process videos through Qwen2.5 ViT and projector."""
539 grid_thw = video_input["video_grid_thw"]
540 assert grid_thw.ndim == 2
541 grid_thw_list = grid_thw.tolist()
542
543 if video_input["type"] == "video_embeds":
544 video_embeds = video_input["video_embeds"].type(self.visual.dtype)
545 else:
546 pixel_values_videos = video_input["pixel_values_videos"]
547 with set_forward_context(None, self.vllm_config):
548 video_embeds = self.visual(pixel_values_videos, grid_thw=grid_thw_list)
549
550 # Apply projector
551 video_embeds = self.mm_projector(video_embeds)
552
553 # Split concatenated embeddings for each video
554 merge_size = self.visual.spatial_merge_size
555 sizes = (grid_thw.prod(-1) // merge_size // merge_size).tolist()
556 return video_embeds.split(sizes)
557
558 def _parse_and_validate_multimodal_inputs(
559 self, **kwargs: object
560 ) -> HCXVisionV2MultiModalInputs:
561 modalities: HCXVisionV2MultiModalInputs = {}
562
563 for input_key in kwargs:
564 if (
565 input_key in ("pixel_values", "image_embeds")
566 and "image" not in modalities
567 ):
568 modalities["image"] = self._parse_and_validate_image_input(**kwargs)
569 if (
570 input_key in ("pixel_values_videos", "video_embeds")
571 and "video" not in modalities
572 ):
573 modalities["video"] = self._parse_and_validate_video_input(**kwargs)
574
575 return modalities
576
577 def embed_multimodal(
578 self,
579 **kwargs: object,
580 ) -> MultiModalEmbeddings:
581 modalities = self._parse_and_validate_multimodal_inputs(**kwargs)
582 if not modalities:
583 return []
584
585 multimodal_embeddings: tuple[torch.Tensor, ...] = ()
586
587 for modality in modalities:
588 if modality == "image":
589 image_input = modalities["image"]
590 if image_input is not None:
591 image_embeddings = self._process_image_input(image_input)
592 multimodal_embeddings += tuple(image_embeddings)
593 if modality == "video":
594 video_input = modalities["video"]
595 if video_input is not None:
596 video_embeddings = self._process_video_input(video_input)
597 multimodal_embeddings += tuple(video_embeddings)
598
599 return multimodal_embeddings
600
601 def forward(
602 self,
603 input_ids: torch.Tensor,
604 positions: torch.Tensor,
605 intermediate_tensors: IntermediateTensors | None = None,
606 inputs_embeds: torch.Tensor | None = None,
607 **kwargs: object,
608 ) -> torch.Tensor | IntermediateTensors:
609 if intermediate_tensors is not None:
610 inputs_embeds = None
611
612 hidden_states = self.language_model.model(
613 input_ids, positions, intermediate_tensors, inputs_embeds=inputs_embeds
614 )
615 return hidden_states
616
617 def compute_logits(
618 self,
619 hidden_states: torch.Tensor,
620 ) -> torch.Tensor | None:
621 return self.language_model.compute_logits(hidden_states)
622
623 def load_weights(
624 self,
625 weights: Iterable[tuple[str, torch.Tensor]],
626 ) -> set[str]:
627 loader = AutoWeightsLoader(self)
628 return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)