A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 17:15 UTC
[Model] Support multi-image for Molmo (#15438)
Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>
Cyrus Leung committed
Mar 26, 2025 at 11:26 UTC
997c8811d6aadf92dc299e0c2a8d274117308880
4 files changed
+39
-35
docs/source/models/supported_models.md
+1
-1
@@ -853,7 +853,7 @@ See [this page](#generative-models) for more information on how to use generativ
853
*
854
- * `MolmoForCausalLM`
855
* Molmo
856
- * T + I
856
+ * T + I<sup>+</sup>
857
* `allenai/Molmo-7B-D-0924`, `allenai/Molmo-7B-O-0924`, etc.
858
* ✅︎
859
* ✅︎
tests/models/decoder_only/vision_language/test_models.py
+1
-1
@@ -431,7 +431,7 @@ VLM_TEST_SETTINGS = {
431
),
432
"molmo": VLMTestInfo(
433
models=["allenai/Molmo-7B-D-0924"],
434
- test_type=(VLMTestType.IMAGE),
434
+ test_type=(VLMTestType.IMAGE, VLMTestType.MULTI_IMAGE),
435
prompt_formatter=identity,
436
max_model_len=4096,
437
max_num_seqs=2,
vllm/model_executor/models/molmo.py
+28
-29
@@ -57,7 +57,7 @@ from .utils import (AutoWeightsLoader, WeightsMapper, flatten_bn,
57
is_pp_missing_parameter,
58
make_empty_intermediate_tensors_factory, make_layers,
59
maybe_prefix, merge_multimodal_embeddings)
60
-from .vision import select_patch_features
60
+from .vision import scatter_patch_features, select_patch_features
61
62
# TODO: hard-coded for now. Consider making it configurable.
63
VIT_LAYERS = [-2, -9]
@@ -71,13 +71,13 @@ POOLING_SIZE = 2
71
72
73
class MolmoImageInputs(TypedDict):
74
- images: Union[torch.Tensor, List[torch.Tensor]]
74
+ images: Union[torch.Tensor, list[torch.Tensor]]
75
"""Shape: `(batch_size, num_crops, num_patch, patch_dim)`"""
76
77
- image_masks: Optional[Union[torch.Tensor, List[torch.Tensor]]]
77
+ image_masks: Optional[Union[torch.Tensor, list[torch.Tensor]]]
78
"""Shape: `(batch_size, num_crops, num_patch)`"""
79
80
- feat_is_patch: Union[torch.Tensor, List[torch.Tensor]]
80
+ feat_is_patch: Union[torch.Tensor, list[torch.Tensor]]
81
"""
82
A boolean mask indicating which image features correspond
83
to patch tokens.
@@ -85,7 +85,7 @@ class MolmoImageInputs(TypedDict):
85
Shape: `(batch_size, num_crops, num_patch)`
86
"""
87
88
- embed_is_patch: Union[torch.Tensor, List[torch.Tensor]]
88
+ embed_is_patch: Union[torch.Tensor, list[torch.Tensor]]
89
"""
90
A boolean mask indicating which image embeddings correspond
91
to patch tokens.
@@ -93,7 +93,7 @@ class MolmoImageInputs(TypedDict):
93
Shape: `(batch_size, num_embeds)`
94
"""
95
96
- num_crops: torch.Tensor
96
+ num_crops: Union[torch.Tensor, list[torch.Tensor]]
97
"""Shape: `(batch_size, num_images)`"""
98
99
@@ -1144,13 +1144,7 @@ class MolmoProcessorWrapper:
1144
1145
image_input_idx = outputs.pop("image_input_idx", None)
1146
if image_input_idx is not None:
1147
- input_is_patch = input_ids == self.image_patch_id
1148
- image_input_idx_flat: torch.Tensor = image_input_idx.view(-1)
1149
- image_valid_flat = image_input_idx_flat >= 0
1150
- feat_is_patch_flat = image_valid_flat.clone()
1151
- feat_is_patch_flat[image_valid_flat] = (
1152
- input_is_patch[image_input_idx_flat[image_valid_flat]])
1153
- feat_is_patch = feat_is_patch_flat.view(*image_input_idx.shape)
1147
+ feat_is_patch = image_input_idx >= 0
1148
1149
input_is_embed = torch.isin(
1150
input_ids,
@@ -1165,6 +1159,17 @@ class MolmoProcessorWrapper:
1159
embed_is_patch = embed_ids == self.image_patch_id
1160
assert embed_is_patch.sum() == feat_is_patch.sum()
1161
1162
+ # image_tokens = extra_joint + joint
1163
+ # Both `extra_joint` and `joint` have `im_start_id` and `im_end_id`
1164
+ embed_start = torch.nonzero(embed_ids == self.im_start_id)[::2, 0]
1165
+ embed_end = torch.nonzero(embed_ids == self.im_end_id)[1::2, 0]
1166
+ assert len(embed_start) == len(embed_end) == len(images)
1167
+
1168
+ embed_is_patch = [
1169
+ embed_is_patch[start:end + 1]
1170
+ for start, end in zip(embed_start, embed_end)
1171
+ ]
1172
+
1173
tilings = [
1174
self.select_tiling(
1175
image_width=image.size[0],
@@ -1180,7 +1185,7 @@ class MolmoProcessorWrapper:
1185
outputs["num_crops"] = num_crops
1186
outputs["img_patch_id"] = self.image_patch_id
1187
1183
- return BatchFeature(outputs, tensor_type=return_tensors)
1188
+ return BatchFeature(outputs)
1189
1190
1191
class MolmoProcessingInfo(BaseProcessingInfo):
@@ -1190,9 +1195,7 @@ class MolmoProcessingInfo(BaseProcessingInfo):
1195
return MolmoProcessorWrapper(processor)
1196
1197
def get_supported_mm_limits(self) -> Mapping[str, Optional[int]]:
1193
- # TODO: Investigate different `embed_is_patch` between cache/no-cache
1194
- # in multi-image case
1195
- return {"image": 1}
1198
+ return {"image": None}
1199
1200
def get_mm_max_tokens_per_item(
1201
self,
@@ -1325,7 +1328,7 @@ class MolmoMultiModalProcessor(BaseMultiModalProcessor[MolmoProcessingInfo]):
1328
"image", num_crops),
1329
feat_is_patch=MultiModalFieldConfig.flat_from_sizes(
1330
"image", num_crops),
1328
- embed_is_patch=MultiModalFieldConfig.shared("image", num_images),
1331
+ embed_is_patch=MultiModalFieldConfig.batched("image"),
1332
num_crops=MultiModalFieldConfig.batched("image"),
1333
img_patch_id=MultiModalFieldConfig.shared("image", num_images),
1334
)
@@ -1499,7 +1502,7 @@ class MolmoForCausalLM(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA,
1502
def _process_image_input(
1503
self,
1504
image_input: MolmoImageInputs,
1502
- ) -> Union[torch.Tensor, List[torch.Tensor]]:
1505
+ ) -> Union[torch.Tensor, list[torch.Tensor]]:
1506
if isinstance(image_input["images"], list):
1507
# Call the vision backbone on the whole batch at once
1508
images_flat = flatten_bn(image_input["images"], concat=True)
@@ -1530,7 +1533,7 @@ class MolmoForCausalLM(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA,
1533
feat_is_patch: torch.Tensor, # Shape: (num_crop, num_patch)
1534
num_crops: torch.Tensor, # Shape: (num_images,)
1535
embed_is_patch: torch.Tensor, # Shape: (num_embeds,)
1533
- ) -> list[torch.Tensor]:
1536
+ ) -> tuple[torch.Tensor, ...]:
1537
"""
1538
Scatter the patch features into a contiguous tensor that corresponds
1539
to the embedding tokens defined by the multimodal processor.
@@ -1565,16 +1568,12 @@ class MolmoForCausalLM(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA,
1568
feats_per_image = features.split(num_crops_per_image)
1569
f_is_patch_per_image = feat_is_patch.split(num_crops_per_image)
1570
1568
- _, _, embed_dim = features.shape
1569
- (num_embeds, ) = embed_is_patch.shape
1570
-
1571
- embeds_in_batch = list[torch.Tensor]()
1572
- for feats, f_is_patch in zip(feats_per_image, f_is_patch_per_image):
1573
- embeds = feats.new_full((num_embeds, embed_dim), torch.nan)
1574
- embeds[embed_is_patch] = feats[f_is_patch]
1575
- embeds_in_batch.append(embeds)
1571
+ features = torch.cat([
1572
+ feats[f_is_patch]
1573
+ for feats, f_is_patch in zip(feats_per_image, f_is_patch_per_image)
1574
+ ])
1575
1577
- return embeds_in_batch
1576
+ return scatter_patch_features(features, embed_is_patch)
1577
1578
def get_multimodal_embeddings(
1579
self, **kwargs: object) -> Optional[MultiModalEmbeddings]:
vllm/model_executor/models/vision.py
+9
-4
@@ -155,7 +155,7 @@ def resolve_visual_encoder_outputs(
155
156
def scatter_patch_features(
157
features: torch.Tensor,
158
- embed_is_patch: torch.Tensor,
158
+ embed_is_patch: Union[torch.Tensor, list[torch.Tensor]],
159
) -> tuple[torch.Tensor, ...]:
160
"""
161
Scatter the patch features into a contiguous tensor that corresponds
@@ -194,14 +194,19 @@ def scatter_patch_features(
194
The resulting embedding tensor is:
195
[ nan p1 p2 nan p3 p4 nan nan ]
196
"""
197
- num_images, num_embeds = embed_is_patch.shape
198
- num_embeds_per_image = [num_embeds] * num_images
197
+ num_embeds_per_image = [
198
+ e_is_patch.numel() for e_is_patch in embed_is_patch
199
+ ]
200
+ if isinstance(embed_is_patch, torch.Tensor):
201
+ embed_is_patch_flat = embed_is_patch.view(-1)
202
+ else:
203
+ embed_is_patch_flat = torch.cat(embed_is_patch)
204
205
embeds_flat = features.new_full(
206
(sum(num_embeds_per_image), features.shape[-1]),
207
fill_value=torch.nan,
208
)
204
- embeds_flat[embed_is_patch.view(-1)] = features.flatten(0, -2)
209
+ embeds_flat[embed_is_patch_flat] = features.flatten(0, -2)
210
211
return embeds_flat.split(num_embeds_per_image)
212