main
py 527 lines 18.4 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Inference-only Bailing MoE V3 vision-language model."""
4
5 import copy
6 from collections.abc import Iterable, Mapping, Sequence
7 from typing import Any
8
9 import numpy as np
10 import regex as re
11 import torch
12 import torch.nn as nn
13 from transformers import AutoProcessor, BatchFeature
14
15 from vllm.config import VllmConfig
16 from vllm.config.multimodal import MultiModalDummyOptions
17 from vllm.inputs import MultiModalDataDict
18 from vllm.model_executor.layers.linear import (
19 ColumnParallelLinear,
20 RowParallelLinear,
21 )
22 from vllm.model_executor.layers.quantization import QuantizationConfig
23 from vllm.model_executor.layers.quantization.fp8 import Fp8Config
24 from vllm.model_executor.models.module_mapping import MultiModelKeys
25 from vllm.multimodal import MULTIMODAL_REGISTRY
26 from vllm.multimodal.inputs import (
27 MultiModalFeatureSpec,
28 MultiModalFieldConfig,
29 MultiModalKwargsItems,
30 )
31 from vllm.multimodal.parse import MultiModalDataItems
32 from vllm.multimodal.processing import (
33 BaseDummyInputsBuilder,
34 BaseMultiModalProcessor,
35 PromptReplacement,
36 PromptUpdate,
37 )
38 from vllm.multimodal.processing.processor import (
39 MultiModalPromptUpdates,
40 PlaceholderFeaturesInfo,
41 )
42 from vllm.sequence import IntermediateTensors
43
44 from .bailing_moe_v3 import BailingMoeV3ForCausalLM
45 from .interfaces import (
46 HasInnerState,
47 IsHybrid,
48 MultiModalEmbeddings,
49 SupportsMRoPE,
50 SupportsMultiModal,
51 SupportsPP,
52 _require_is_multimodal,
53 )
54 from .qwen2_5_vl import (
55 Qwen2_5_VLImageEmbeddingInputs,
56 Qwen2_5_VLImageInputs,
57 Qwen2_5_VLImagePixelInputs,
58 )
59 from .qwen2_vl import (
60 Qwen2VLProcessingInfo,
61 _create_qwen2vl_field_factory,
62 )
63 from .qwen3_vl import Qwen3_VisionTransformer
64 from .utils import (
65 AutoWeightsLoader,
66 WeightsMapper,
67 _merge_multimodal_embeddings,
68 maybe_prefix,
69 )
70
71 IMAGE_TOKEN = "<|image_pad|>"
72 IMAGE_PLACEHOLDER = f"<|vision_start|>{IMAGE_TOKEN}<|vision_end|>"
73
74
75 class BailingMoeV3VLProcessingInfo(Qwen2VLProcessingInfo):
76 """Processing metadata for the remote Bailing multimodal processor."""
77
78 def get_hf_config(self):
79 return self.ctx.get_hf_config()
80
81 def get_hf_processor(self, **kwargs: object) -> AutoProcessor:
82 # The checkpoint supplies its processor through trust_remote_code.
83 return self.ctx.get_hf_processor(**kwargs)
84
85 def get_supported_mm_limits(self) -> Mapping[str, int | None]:
86 return {"image": None}
87
88 def get_mm_max_tokens_per_item(
89 self,
90 seq_len: int,
91 mm_counts: Mapping[str, int],
92 ) -> Mapping[str, int]:
93 return {"image": self.get_max_image_tokens()}
94
95
96 class BailingMoeV3VLDummyInputsBuilder(
97 BaseDummyInputsBuilder[BailingMoeV3VLProcessingInfo]
98 ):
99 def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
100 return IMAGE_PLACEHOLDER * mm_counts.get("image", 0)
101
102 def get_dummy_mm_data(
103 self,
104 seq_len: int,
105 mm_counts: Mapping[str, int],
106 mm_options: MultiModalDummyOptions,
107 ) -> MultiModalDataDict:
108 target_width, target_height = self.info.get_image_size_with_most_features()
109 return {
110 "image": self._get_dummy_images(
111 width=target_width,
112 height=target_height,
113 num_images=mm_counts.get("image", 0),
114 overrides=mm_options.get("image"),
115 )
116 }
117
118
119 class BailingMoeV3VLMultiModalProcessor(
120 BaseMultiModalProcessor[BailingMoeV3VLProcessingInfo]
121 ):
122 def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str:
123 return self.dummy_inputs.get_dummy_text(mm_counts)
124
125 def _apply_prompt_updates(
126 self,
127 token_ids: list[int],
128 mm_prompt_updates: MultiModalPromptUpdates,
129 ) -> tuple[list[int], Mapping[str, list[PlaceholderFeaturesInfo]]]:
130 if image_updates := mm_prompt_updates.get("image"):
131 tokenizer = self.info.get_tokenizer()
132 # HF adds a newline after each image before tokenization. Re-encode
133 # the text so it can merge with existing whitespace at the boundary.
134 text = re.sub(
135 f"{re.escape(IMAGE_PLACEHOLDER)}|{re.escape(IMAGE_TOKEN)}",
136 lambda _: IMAGE_PLACEHOLDER + "\n",
137 tokenizer.decode(token_ids),
138 count=len(image_updates),
139 )
140 token_ids = tokenizer.encode(text, add_special_tokens=False)
141
142 return super()._apply_prompt_updates(token_ids, mm_prompt_updates)
143
144 def _get_mm_fields_config(
145 self,
146 hf_inputs: BatchFeature,
147 hf_processor_mm_kwargs: Mapping[str, object],
148 ) -> Mapping[str, MultiModalFieldConfig]:
149 merge_size = self.info.get_hf_config().vision_config.spatial_merge_size
150 fields = _create_qwen2vl_field_factory(merge_size)(hf_inputs)
151 return {
152 key: fields[key]
153 for key in ("pixel_values", "image_embeds", "image_grid_thw")
154 }
155
156 def _get_prompt_updates(
157 self,
158 mm_items: MultiModalDataItems,
159 hf_processor_mm_kwargs: Mapping[str, Any],
160 out_mm_kwargs: MultiModalKwargsItems,
161 ) -> Sequence[PromptUpdate]:
162 tokenizer = self.info.get_tokenizer()
163 image_token_id = tokenizer.get_vocab()[IMAGE_TOKEN]
164 merge_size = self.info.get_hf_config().vision_config.spatial_merge_size
165 merge_unit = merge_size**2
166
167 def get_image_replacement(item_idx: int) -> list[int]:
168 out_item = out_mm_kwargs["image"][item_idx]
169 image_grid_thw = out_item["image_grid_thw"].data
170 assert isinstance(image_grid_thw, torch.Tensor)
171 num_image_tokens = int(image_grid_thw.prod()) // merge_unit
172 return [image_token_id] * num_image_tokens
173
174 return [
175 PromptReplacement(
176 modality="image",
177 target=[image_token_id],
178 replacement=get_image_replacement,
179 )
180 ]
181
182
183 class BailingMoeV3VisionPatchMerger(nn.Module):
184 """Normalize patches and merge each 2x2 group without an MLP."""
185
186 def __init__(
187 self,
188 context_dim: int,
189 spatial_merge_size: int,
190 norm_eps: float = 1e-6,
191 ) -> None:
192 super().__init__()
193 self.hidden_size = context_dim * spatial_merge_size**2
194 self.norm = nn.LayerNorm(context_dim, eps=norm_eps)
195
196 def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
197 hidden_states = self.norm(hidden_states)
198 return hidden_states.reshape(-1, self.hidden_size)
199
200
201 class BailingMoeV3VisionTransformer(Qwen3_VisionTransformer):
202 """Qwen3 ViT with Bailing's parameter-free patch merger."""
203
204 def __init__(
205 self,
206 vision_config: Any,
207 norm_eps: float = 1e-6,
208 quant_config: QuantizationConfig | None = None,
209 prefix: str = "",
210 ) -> None:
211 # Honor whole-tower exclusions before the parent creates linear layers.
212 if (
213 isinstance(quant_config, Fp8Config)
214 and prefix in quant_config.ignored_layers
215 ):
216 quant_config = None
217
218 vision_config = copy.deepcopy(vision_config)
219 vision_config.deepstack_visual_indexes = []
220 merged_hidden_size = (
221 vision_config.hidden_size * vision_config.spatial_merge_size**2
222 )
223 vision_config.out_hidden_size = merged_hidden_size
224
225 super().__init__(
226 vision_config,
227 norm_eps=norm_eps,
228 quant_config=quant_config,
229 prefix=prefix,
230 )
231 self.merger = BailingMoeV3VisionPatchMerger(
232 context_dim=vision_config.hidden_size,
233 spatial_merge_size=vision_config.spatial_merge_size,
234 norm_eps=norm_eps,
235 )
236 self.deepstack_visual_indexes = []
237 self.deepstack_merger_list = nn.ModuleList()
238 self.out_hidden_size = merged_hidden_size
239
240
241 class BailingMoeV3VLProjector(nn.Module):
242 """Tensor-parallel MLP mapping merged vision features to text width."""
243
244 def __init__(
245 self,
246 vision_hidden_size: int,
247 text_hidden_size: int,
248 quant_config: QuantizationConfig | None = None,
249 prefix: str = "",
250 ) -> None:
251 super().__init__()
252 if (
253 isinstance(quant_config, Fp8Config)
254 and prefix in quant_config.ignored_layers
255 ):
256 quant_config = None
257
258 self.linear_fc1 = ColumnParallelLinear(
259 vision_hidden_size,
260 text_hidden_size,
261 bias=True,
262 quant_config=quant_config,
263 prefix=maybe_prefix(prefix, "linear_fc1"),
264 )
265 self.act_fn = nn.GELU()
266 self.linear_fc2 = RowParallelLinear(
267 text_hidden_size,
268 text_hidden_size,
269 bias=True,
270 quant_config=quant_config,
271 prefix=maybe_prefix(prefix, "linear_fc2"),
272 )
273
274 def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
275 hidden_states, _ = self.linear_fc1(hidden_states)
276 hidden_states = self.act_fn(hidden_states)
277 hidden_states, _ = self.linear_fc2(hidden_states)
278 return hidden_states
279
280
281 @MULTIMODAL_REGISTRY.register_processor(
282 BailingMoeV3VLMultiModalProcessor,
283 info=BailingMoeV3VLProcessingInfo,
284 dummy_inputs=BailingMoeV3VLDummyInputsBuilder,
285 )
286 class BailingMoeV3VLForConditionalGeneration(
287 nn.Module,
288 HasInnerState,
289 IsHybrid,
290 SupportsMultiModal,
291 SupportsPP,
292 SupportsMRoPE,
293 ):
294 """Native vLLM wrapper for ``BailingMoeV3VLConfig`` checkpoints."""
295
296 packed_modules_mapping = {
297 **BailingMoeV3ForCausalLM.packed_modules_mapping,
298 # The vision checkpoint already stores Q/K/V in one tensor.
299 "qkv": ["qkv"],
300 }
301
302 hf_to_vllm_mapper = BailingMoeV3ForCausalLM.hf_to_vllm_mapper | WeightsMapper(
303 orig_to_new_prefix={
304 # Match module names as well as parameter names for quantization.
305 "model.visual": "visual",
306 "model.linear_proj": "linear_proj",
307 "linear_proj.0": "linear_proj.linear_fc1",
308 "linear_proj.2": "linear_proj.linear_fc2",
309 "lm_head": "language_model.lm_head",
310 "model.": "language_model.model.",
311 }
312 )
313
314 @classmethod
315 def get_placeholder_str(cls, modality: str, i: int) -> str | None:
316 if modality.startswith("image"):
317 return IMAGE_PLACEHOLDER
318 raise ValueError("Only image modality is supported")
319
320 @classmethod
321 def get_mamba_state_shape_from_config(
322 cls, vllm_config: VllmConfig
323 ) -> tuple[tuple[int, int], tuple[int, int, int]]:
324 text_config = vllm_config.model_config.hf_config.text_config
325 return BailingMoeV3ForCausalLM.get_mamba_state_shape_from_config(
326 vllm_config.with_hf_config(text_config)
327 )
328
329 @classmethod
330 def get_mamba_state_dtype_from_config(
331 cls, vllm_config: VllmConfig
332 ) -> tuple[torch.dtype, torch.dtype]:
333 text_config = vllm_config.model_config.hf_config.text_config
334 return BailingMoeV3ForCausalLM.get_mamba_state_dtype_from_config(
335 vllm_config.with_hf_config(text_config)
336 )
337
338 @classmethod
339 def get_mamba_state_copy_func(cls) -> tuple:
340 return BailingMoeV3ForCausalLM.get_mamba_state_copy_func()
341
342 def __init__(self, *, vllm_config: VllmConfig, prefix: str = "model") -> None:
343 super().__init__()
344 config = vllm_config.model_config.hf_config
345 vision_config = config.vision_config
346 text_config = config.text_config
347 quant_config = vllm_config.quant_config
348
349 self.config = config
350 self.model_config = vllm_config.model_config
351
352 merged_vision_size = vision_config.hidden_size * (
353 vision_config.spatial_merge_size**2
354 )
355 with self._mark_tower_model(vllm_config, "image"):
356 self.visual = BailingMoeV3VisionTransformer(
357 vision_config,
358 norm_eps=getattr(config, "rms_norm_eps", 1e-6),
359 quant_config=quant_config,
360 prefix=maybe_prefix(prefix, "visual"),
361 )
362 self.linear_proj = BailingMoeV3VLProjector(
363 vision_hidden_size=merged_vision_size,
364 text_hidden_size=text_config.hidden_size,
365 quant_config=quant_config,
366 prefix=maybe_prefix(prefix, "linear_proj"),
367 )
368
369 with self._mark_language_model(vllm_config):
370 self.language_model = BailingMoeV3ForCausalLM(
371 vllm_config=vllm_config.with_hf_config(text_config),
372 prefix=maybe_prefix(prefix, "language_model"),
373 )
374
375 self.make_empty_intermediate_tensors = (
376 self.language_model.make_empty_intermediate_tensors
377 )
378
379 def _parse_and_validate_image_input(
380 self, **kwargs: object
381 ) -> Qwen2_5_VLImageInputs | None:
382 pixel_values = kwargs.pop("pixel_values", None)
383 image_embeds = kwargs.pop("image_embeds", None)
384 image_grid_thw = kwargs.pop("image_grid_thw", None)
385
386 if pixel_values is None and image_embeds is None:
387 return None
388 if pixel_values is not None:
389 return Qwen2_5_VLImagePixelInputs(
390 type="pixel_values",
391 pixel_values=pixel_values,
392 image_grid_thw=image_grid_thw,
393 )
394 return Qwen2_5_VLImageEmbeddingInputs(
395 type="image_embeds",
396 image_embeds=image_embeds,
397 image_grid_thw=image_grid_thw,
398 )
399
400 def _process_image_input(
401 self, image_input: Qwen2_5_VLImageInputs
402 ) -> tuple[torch.Tensor, ...]:
403 image_grid_thw = image_input["image_grid_thw"]
404 assert image_grid_thw.ndim == 2
405
406 if image_input["type"] == "image_embeds":
407 image_embeds = image_input["image_embeds"]
408 else:
409 pixel_values = image_input["pixel_values"].type(self.visual.dtype)
410 vision_features = self.visual(
411 pixel_values,
412 grid_thw=image_grid_thw,
413 )
414 image_embeds = self.linear_proj(vision_features)
415
416 image_embeds = image_embeds.to(dtype=self.visual.dtype)
417 merge_unit = self.visual.spatial_merge_size**2
418 sizes = (image_grid_thw.prod(-1) // merge_unit).tolist()
419 return image_embeds.split(sizes)
420
421 def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:
422 image_input = self._parse_and_validate_image_input(**kwargs)
423 if image_input is None:
424 return []
425 return self._process_image_input(image_input)
426
427 def embed_input_ids(
428 self,
429 input_ids: torch.Tensor,
430 multimodal_embeddings: MultiModalEmbeddings | None = None,
431 *,
432 is_multimodal: torch.Tensor | None = None,
433 ) -> torch.Tensor:
434 inputs_embeds = self._embed_text_input_ids(
435 input_ids,
436 self.language_model.embed_input_ids,
437 is_multimodal=is_multimodal,
438 )
439 if multimodal_embeddings is None or len(multimodal_embeddings) == 0:
440 return inputs_embeds
441
442 return _merge_multimodal_embeddings(
443 inputs_embeds=inputs_embeds,
444 multimodal_embeddings=multimodal_embeddings,
445 is_multimodal=_require_is_multimodal(is_multimodal),
446 )
447
448 def get_mrope_input_positions(
449 self,
450 input_tokens: list[int],
451 mm_features: list[MultiModalFeatureSpec],
452 ) -> tuple[torch.Tensor, int]:
453 position_blocks: list[np.ndarray] = []
454 consumed = 0
455 merge_size = self.config.vision_config.spatial_merge_size
456
457 for mm_feature in sorted(mm_features, key=lambda f: f.mm_position.offset):
458 if mm_feature.modality != "image":
459 raise ValueError(f"Unsupported modality: {mm_feature.modality}")
460
461 offset = mm_feature.mm_position.offset
462 text_len = offset - consumed
463 start = position_blocks[-1].max() + 1 if position_blocks else 0
464 position_blocks.append(
465 np.broadcast_to(np.arange(text_len), (3, text_len)) + start
466 )
467
468 assert mm_feature.data is not None
469 image_grid_thw = mm_feature.data["image_grid_thw"].data
470 assert isinstance(image_grid_thw, torch.Tensor)
471 t, h, w = image_grid_thw.tolist()
472 assert t == 1, f"Image must have one temporal grid, got {t}"
473 grid = np.indices((t, h // merge_size, w // merge_size))
474 position_blocks.append(grid.reshape(3, -1) + text_len + start)
475 consumed = offset + t * (h // merge_size) * (w // merge_size)
476
477 if consumed < len(input_tokens):
478 start = position_blocks[-1].max() + 1 if position_blocks else 0
479 text_len = len(input_tokens) - consumed
480 position_blocks.append(
481 np.broadcast_to(np.arange(text_len), (3, text_len)) + start
482 )
483
484 if not position_blocks:
485 positions = np.empty((3, 0), dtype=np.int64)
486 return torch.from_numpy(positions), 0
487
488 positions = np.concatenate(position_blocks, axis=1).reshape(3, -1)
489 position_delta = int(positions.max() + 1 - len(input_tokens))
490 return torch.from_numpy(positions), position_delta
491
492 def forward(
493 self,
494 input_ids: torch.Tensor | None,
495 positions: torch.Tensor,
496 intermediate_tensors: IntermediateTensors | None = None,
497 inputs_embeds: torch.Tensor | None = None,
498 **kwargs: object,
499 ) -> torch.Tensor | IntermediateTensors:
500 if intermediate_tensors is not None:
501 inputs_embeds = None
502 return self.language_model(
503 input_ids=input_ids,
504 positions=positions,
505 intermediate_tensors=intermediate_tensors,
506 inputs_embeds=inputs_embeds,
507 )
508
509 def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor:
510 return self.language_model.compute_logits(hidden_states)
511
512 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
513 loader = AutoWeightsLoader(self)
514 return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
515
516 def get_mm_mapping(self) -> MultiModelKeys:
517 return MultiModelKeys.from_string_field(
518 language_model="language_model",
519 connector="linear_proj",
520 tower_model="visual.",
521 )
522
523 def get_num_mm_encoder_tokens(self, num_image_tokens: int) -> int:
524 return num_image_tokens * self.visual.spatial_merge_size**2
525
526 def get_num_mm_connector_tokens(self, num_vision_tokens: int) -> int:
527 return num_vision_tokens // self.visual.spatial_merge_size**2