27
from collections import defaultdict
28
from collections.abc import Iterable, Mapping, Sequence
29
from functools import partial
30
+from itertools import chain
31
from typing import Annotated, Any, Callable, Literal, Optional, Union
32
33
import numpy as np
34
import torch
35
import torch.types
36
from torch import nn
37
+from torch.nn.init import trunc_normal_
38
from transformers import BatchFeature, PretrainedConfig
39
from typing_extensions import TypeVar
40
49
from vllm.model_executor.models.minicpm import MiniCPMForCausalLM
50
from vllm.model_executor.models.module_mapping import MultiModelKeys
51
from vllm.model_executor.models.qwen2 import Qwen2ForCausalLM
52
+from vllm.model_executor.models.qwen3 import Qwen3ForCausalLM
53
from vllm.model_executor.sampling_metadata import SamplingMetadata
51
-from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalKwargsItems
54
+from vllm.multimodal import MULTIMODAL_REGISTRY
55
from vllm.multimodal.inputs import (MultiModalDataDict, MultiModalFieldConfig,
53
- NestedTensors)
56
+ MultiModalKwargsItems, NestedTensors)
57
from vllm.multimodal.parse import (DictEmbeddingItems, ImageItem,
58
ImageProcessorItems, ImageSize,
59
ModalityData, ModalityDataItems,
221
return x
222
223
224
+class Resampler4_5(Resampler2_5):
225
+
226
+ def __init__(self,
227
+ num_queries: int,
228
+ embed_dim: int,
229
+ num_heads: int,
230
+ kv_dim: Optional[int] = None,
231
+ norm_layer: Callable[[int], nn.LayerNorm] = DEFAULT_LN,
232
+ max_size: tuple[int, int] = (70, 70),
233
+ max_temporal_size: int = 36000,
234
+ quant_config: Optional[QuantizationConfig] = None,
235
+ prefix: str = "") -> None:
236
+ super().__init__(num_queries,
237
+ embed_dim,
238
+ num_heads,
239
+ kv_dim,
240
+ norm_layer,
241
+ max_size,
242
+ quant_config=quant_config,
243
+ prefix=prefix)
244
+
245
+ trunc_normal_(self.query, std=.02)
246
+ self.max_temporal_size = max_temporal_size
247
+ self._set_temporal_pos_cache(self.max_temporal_size)
248
+ self.apply(self._init_weights)
249
+
250
+ def get_1d_sincos_pos_embed_from_temporal_size(self, embed_dim: int,
251
+ pos: np.ndarray):
252
+ """
253
+ embed_dim: output dimension for each position
254
+ pos: a list of positions to be encoded: size (M,)
255
+ out: (M, D)
256
+ """
257
+ assert embed_dim % 2 == 0
258
+ omega = np.arange(embed_dim // 2, dtype=np.float32)
259
+ omega /= embed_dim / 2.
260
+ omega = 1. / 10000**omega # (D/2,)
261
+
262
+ pos = pos.reshape(-1) # (M,)
263
+ out = np.einsum('m,d->md', pos, omega) # (M, D/2), outer product
264
+
265
+ emb_sin = np.sin(out) # (M, D/2)
266
+ emb_cos = np.cos(out) # (M, D/2)
267
+
268
+ emb = np.concatenate([emb_sin, emb_cos], axis=1) # (M, D)
269
+ return emb
270
+
271
+ def _set_temporal_pos_cache(self,
272
+ max_temporal_size: int,
273
+ device: torch.types.Device = "cpu") -> None:
274
+ temporal_size = np.arange(max_temporal_size, dtype=np.float32)
275
+ pos_embed = torch.from_numpy(
276
+ self.get_1d_sincos_pos_embed_from_temporal_size(
277
+ self.embed_dim, temporal_size)).float().to(device)
278
+ self.register_buffer("temporal_pos_embed", pos_embed, persistent=False)
279
+
280
+ def _adjust_temporal_pos_cache(self,
281
+ max_temporal_size: int,
282
+ device: torch.types.Device = "cpu"):
283
+ if max_temporal_size > self.max_temporal_size:
284
+ self.max_temporal_size = max_temporal_size
285
+ self._set_temporal_pos_cache(self.max_temporal_size, device)
286
+
287
+ def _init_weights(self, m: Union[nn.Linear, nn.LayerNorm]):
288
+ if isinstance(m, nn.Linear):
289
+ trunc_normal_(m.weight, std=.02)
290
+ if isinstance(m, nn.Linear) and m.bias is not None:
291
+ nn.init.constant_(m.bias, 0)
292
+ elif isinstance(m, nn.LayerNorm):
293
+ nn.init.constant_(m.bias, 0)
294
+ nn.init.constant_(m.weight, 1.0)
295
+
296
+ def forward(
297
+ self,
298
+ x: torch.Tensor,
299
+ tgt_sizes: torch.Tensor,
300
+ # temporal_ids for high refresh rate videos
301
+ temporal_ids=None
302
+ ) -> torch.Tensor:
303
+ assert x.shape[0] == tgt_sizes.shape[0]
304
+ bs = x.shape[0]
305
+
306
+ device = x.device
307
+ dtype = x.dtype
308
+
309
+ patch_len = tgt_sizes[:, 0] * tgt_sizes[:, 1]
310
+
311
+ self._adjust_pos_cache(tgt_sizes, device=device)
312
+
313
+ temporal_pos_emb = False
314
+ temporal_ids_flatten = None
315
+ if temporal_ids is not None:
316
+ # example: [[-1], [-1], [2, 6, 9]]
317
+ temporal_ids_flatten = list(chain.from_iterable(temporal_ids))
318
+ max_temporal_size = max(temporal_ids_flatten, default=0)
319
+ if max_temporal_size > -1:
320
+ temporal_pos_emb = True
321
+ if max_temporal_size > self.max_temporal_size:
322
+ self._adjust_temporal_pos_cache(max_temporal_size, device)
323
+
324
+ max_patch_len = patch_len.max().item()
325
+ assert isinstance(max_patch_len, int)
326
+
327
+ key_padding_mask = torch.zeros((bs, max_patch_len),
328
+ dtype=torch.bool,
329
+ device=device)
330
+
331
+ x, _ = self.kv_proj(x) # B * L * D
332
+ x = self.ln_kv(x).permute(1, 0, 2) # L * B * D
333
+ q = self.ln_q(self.query) # Q * D
334
+
335
+ pos_embed_2d = []
336
+ pos_embed_temporal = []
337
+ for i in range(bs):
338
+ tgt_h, tgt_w = tgt_sizes[i]
339
+ if temporal_pos_emb:
340
+ if temporal_ids_flatten[i] == -1:
341
+ pos_embed_temporal.append(
342
+ torch.zeros(self.embed_dim, dtype=dtype,
343
+ device=device))
344
+ else:
345
+ pos_embed_temporal.append(self.temporal_pos_embed[
346
+ temporal_ids_flatten[i]].to(dtype)) # D
347
+
348
+ pos_embed_2d.append(self.pos_embed[:tgt_h, :tgt_w, :].reshape(
349
+ (tgt_h * tgt_w, -1)).to(dtype)) # patches * D
350
+ key_padding_mask[i, patch_len[i]:] = True
351
+
352
+ pos_embed_2d = torch.nn.utils.rnn.pad_sequence(
353
+ pos_embed_2d, batch_first=True,
354
+ padding_value=0.0).permute(1, 0, 2) # BLD => L * B * D
355
+
356
+ k = x
357
+ v = x + pos_embed_2d
358
+ if pos_embed_temporal:
359
+ k += torch.stack(pos_embed_temporal, dim=0)
360
+ bs = len(temporal_ids)
361
+ merge_k = []
362
+ merge_v = []
363
+ merge_key_padding_mask = []
364
+
365
+ start = 0
366
+ for tp in temporal_ids:
367
+ end = start + len(tp)
368
+ # L * (end-start) * D -> (end-start) * L * D
369
+ # -> 1 * L*(end-start) * D
370
+ merge_k.append(k[:, start:end, :].permute(1, 0, 2).reshape(
371
+ -1, self.embed_dim))
372
+ merge_v.append(v[:, start:end, :].permute(1, 0, 2).reshape(
373
+ -1, self.embed_dim))
374
+ merge_key_padding_mask.append(
375
+ key_padding_mask[start:end, :].reshape(-1, 1))
376
+
377
+ start = end
378
+
379
+ k = torch.nn.utils.rnn.pad_sequence(merge_k,
380
+ batch_first=True,
381
+ padding_value=0.0).permute(
382
+ 1, 0, 2) # L*(end-start)
383
+ v = torch.nn.utils.rnn.pad_sequence(merge_v,
384
+ batch_first=True,
385
+ padding_value=0.0).permute(
386
+ 1, 0, 2) # L*(end-start)
387
+ key_padding_mask = torch.nn.utils.rnn.pad_sequence(
388
+ merge_key_padding_mask, batch_first=True,
389
+ padding_value=True).squeeze(-1)
390
+
391
+ out = self.attn(
392
+ self._repeat(q, bs), # Q * B * D
393
+ k, # L * B * D + L * B * D
394
+ v,
395
+ key_padding_mask=key_padding_mask,
396
+ )[0]
397
+ # out: Q * B * D
398
+ x = out.permute(1, 0, 2) # B * Q * D
399
+
400
+ x = self.ln_post(x)
401
+ x = x @ self.proj
402
+ return x
403
+
404
+
405
def get_version_by_config(config: PretrainedConfig) -> tuple[int, ...]:
406
version_float = getattr(config, "version", None)
407
538
539
def get_supported_mm_limits(self) -> Mapping[str, Optional[int]]:
540
mm_limits = {"image": None}
357
- if self.get_model_version() == (2,
358
- 6) or self.get_model_version() == (4,
359
- 0):
541
+ if self.get_model_version() in {(2, 6), (4, 0), (4, 5)}:
542
mm_limits["video"] = None
543
544
return mm_limits
819
out_keys: set[str],
820
) -> dict[str, NestedTensors]:
821
# This processor supports zipping prompt and mm_data together
640
- if self.info.get_model_version() == (
641
- 2, 6) or self.info.get_model_version() == (4, 0):
822
+ if self.info.get_model_version() in {(2, 6), (4, 0), (4, 5)}:
823
inputs = super()._call_hf_processor(
824
prompt=prompts, # type: ignore
825
mm_data=mm_data,
997
# and config class
998
self.config = config
999
self.multimodal_config = multimodal_config
819
- self.use_data_parallel = multimodal_config.mm_encoder_tp_mode == "data"
1000
1001
self.version = get_version_by_config(self.config)
1002
self.llm = self.init_llm(vllm_config=vllm_config,
1544
prefix: str = "",
1545
) -> nn.Module:
1546
quant_config = self._maybe_ignore_quant_config(quant_config)
1367
- model = Idefics2VisionTransformer(
1368
- config.vision_config,
1369
- quant_config=quant_config,
1370
- prefix=prefix,
1371
- use_data_parallel=self.use_data_parallel)
1547
+ model = Idefics2VisionTransformer(config.vision_config,
1548
+ quant_config=quant_config,
1549
+ prefix=prefix)
1550
if self.config.drop_vision_last_layer:
1551
model.encoder.layers = model.encoder.layers[:-1]
1552
return model
1614
return loader.load_weights(weights)
1615
1616
1617
+class MiniCPMV4_5(MiniCPMVBaseModel, SupportsLoRA):
1618
+ packed_modules_mapping = {
1619
+ "qkv_proj": [
1620
+ "q_proj",
1621
+ "k_proj",
1622
+ "v_proj",
1623
+ ],
1624
+ "gate_up_proj": [
1625
+ "gate_proj",
1626
+ "up_proj",
1627
+ ],
1628
+ }
1629
+
1630
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
1631
+ super().__init__(vllm_config=vllm_config, prefix=prefix)
1632
+ assert self.version == (4, 5)
1633
+
1634
+ def _maybe_ignore_quant_config(self, quant_config: QuantizationConfig):
1635
+ if isinstance(quant_config, (AWQConfig, AWQMarlinConfig)):
1636
+ return None
1637
+ return quant_config
1638
+
1639
+ def init_llm(
1640
+ self,
1641
+ vllm_config: VllmConfig,
1642
+ prefix: str = "",
1643
+ ) -> nn.Module:
1644
+ return Qwen3ForCausalLM(vllm_config=vllm_config, prefix=prefix)
1645
+
1646
+ def init_vision_module(
1647
+ self,
1648
+ config: PretrainedConfig,
1649
+ quant_config: Optional[QuantizationConfig] = None,
1650
+ prefix: str = "",
1651
+ ) -> nn.Module:
1652
+ quant_config = self._maybe_ignore_quant_config(quant_config)
1653
+ model = Idefics2VisionTransformer(config.vision_config,
1654
+ quant_config=quant_config,
1655
+ prefix=prefix)
1656
+ if self.config.drop_vision_last_layer:
1657
+ model.encoder.layers = model.encoder.layers[:-1]
1658
+ return model
1659
+
1660
+ def init_resampler(
1661
+ self,
1662
+ embed_dim: int,
1663
+ vision_dim: int,
1664
+ quant_config: Optional[QuantizationConfig] = None,
1665
+ prefix: str = "",
1666
+ ) -> nn.Module:
1667
+ quant_config = self._maybe_ignore_quant_config(quant_config)
1668
+ with set_default_torch_dtype(torch.float16):
1669
+ # The resampler in 4.0 remains consistent with the one in 2.5/2.6.
1670
+ resampler = Resampler4_5(num_queries=self.config.query_num,
1671
+ embed_dim=embed_dim,
1672
+ num_heads=embed_dim // 128,
1673
+ kv_dim=vision_dim,
1674
+ quant_config=quant_config,
1675
+ prefix=prefix)
1676
+
1677
+ return resampler.to(device=current_platform.device_type,
1678
+ dtype=torch.get_default_dtype())
1679
+
1680
+ def get_vision_hidden_states(
1681
+ self, data: MiniCPMVImagePixelInputs) -> torch.Tensor:
1682
+ pixel_values = data["pixel_values"]
1683
+ tgt_sizes = data["tgt_sizes"]
1684
+ temporal_ids = data.get('temporal_ids', None)
1685
+
1686
+ B = len(pixel_values)
1687
+ P = pixel_values[0].shape[-2]
1688
+ L = max(item.shape[-1] for item in pixel_values)
1689
+ device = pixel_values[0].device
1690
+ dtype = pixel_values[0].dtype
1691
+
1692
+ all_pixel_values = torch.zeros((B, 3, P, L),
1693
+ dtype=dtype,
1694
+ device=device)
1695
+ all_temporal_ids = None if temporal_ids is None else flatten_2d_lists(
1696
+ temporal_ids)
1697
+ for i, pixel_values_item in enumerate(pixel_values):
1698
+ L_item = pixel_values_item.shape[-1]
1699
+ all_pixel_values[i, ..., :L_item] = pixel_values_item
1700
+
1701
+ num_patches = tgt_sizes.prod(-1)
1702
+ max_patches = num_patches.max().item()
1703
+ assert isinstance(max_patches, int)
1704
+
1705
+ patch_attn_mask = torch.zeros((B, max_patches),
1706
+ dtype=torch.bool,
1707
+ device=device)
1708
+ for i, num_patches_item in enumerate(num_patches):
1709
+ patch_attn_mask[i, :num_patches_item] = True
1710
+
1711
+ vision_embedding = self.vpm(
1712
+ all_pixel_values,
1713
+ patch_attention_mask=patch_attn_mask.unsqueeze(1),
1714
+ tgt_sizes=tgt_sizes,
1715
+ )
1716
+
1717
+ return self.resampler(vision_embedding, tgt_sizes, all_temporal_ids)
1718
+
1719
+ def load_weights(self, weights: Iterable[tuple[str,
1720
+ torch.Tensor]]) -> set[str]:
1721
+ loader = AutoWeightsLoader(self,
1722
+ skip_prefixes=["apm.", "audio", "tts"])
1723
+ return loader.load_weights(weights)
1724
+
1725
+
1726
_SUPPORT_VERSION = {
1727
(2, 0): MiniCPMV2_0,
1728
(2, 5): MiniCPMV2_5,
1729
(2, 6): MiniCPMV2_6,
1730
(4, 0): MiniCPMV4_0,
1731
+ (4, 5): MiniCPMV4_5,
1732
}
1733
1734