@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 61c6a5a796

Fork 0

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, 18:30 UTC

[VLM] Merged multi-modal processor for Pixtral (#12211)

Signed-off-by: remi <remi@mistral.ai> Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk> Co-authored-by: DarkLight1337 <tlleungac@connect.ust.hk>

Rémi Delacourt committed Mar 15, 2025 at 14:28 UTC 61c6a5a79664882a8ab1c9af3ff78677911516dc

9 files changed +622 -360

examples/offline_inference/pixtral.py

+18 -6

@@ -43,12 +43,18 @@ from vllm.sampling_params import SamplingParams

43 # python demo.py advanced

46 -def run_simple_demo():

46 +def run_simple_demo(args: argparse.Namespace):

47 model_name = "mistralai/Pixtral-12B-2409"

48 sampling_params = SamplingParams(max_tokens=8192)

50 - # Lower max_num_seqs or max_model_len on low-VRAM GPUs.

51 - llm = LLM(model=model_name, tokenizer_mode="mistral")

50 + # Lower max_model_len and/or max_num_seqs on low-VRAM GPUs.

51 + llm = LLM(

52 + model=model_name,

53 + tokenizer_mode="mistral",

54 + max_model_len=4096,

55 + max_num_seqs=2,

56 + disable_mm_preprocessor_cache=args.disable_mm_preprocessor_cache,

57 + )

59 prompt = "Describe this image in one sentence."

60 image_url = "https://picsum.photos/id/237/200/300"

@@ -76,7 +82,7 @@ def run_simple_demo():

82 print(outputs[0].outputs[0].text)

79 -def run_advanced_demo():

85 +def run_advanced_demo(args: argparse.Namespace):

86 model_name = "mistralai/Pixtral-12B-2409"

87 max_img_per_msg = 5

88 max_tokens_per_img = 4096

@@ -87,6 +93,7 @@ def run_advanced_demo():

93 tokenizer_mode="mistral",

94 limit_mm_per_prompt={"image": max_img_per_msg},

95 max_model_len=max_img_per_msg * max_tokens_per_img,

96 + disable_mm_preprocessor_cache=args.disable_mm_preprocessor_cache,

97 )

99 prompt = "Describe the following image."

@@ -153,14 +160,19 @@ def main():

160 help="Specify the demo mode: 'simple' or 'advanced'",

161 )

162

163 + parser.add_argument(

164 + '--disable-mm-preprocessor-cache',

165 + action='store_true',

166 + help='If True, disables caching of multi-modal preprocessor/mapper.')

167 +

168 args = parser.parse_args()

169

170 if args.mode == "simple":

171 print("Running simple demo...")

160 - run_simple_demo()

172 + run_simple_demo(args)

173 elif args.mode == "advanced":

174 print("Running advanced demo...")

163 - run_advanced_demo()

175 + run_advanced_demo(args)

176

177

178 if __name__ == "__main__":

tests/models/multimodal/processing/test_common.py

+149 -51

@@ -2,17 +2,23 @@

3 import copy

4 from functools import partial

5 -from typing import Optional

5 +from typing import Optional, Union

7 import numpy as np

8 import pytest

9 +from mistral_common.protocol.instruct.messages import (ImageChunk, TextChunk,

10 + UserMessage)

11 +from mistral_common.protocol.instruct.request import ChatCompletionRequest

12 from PIL import Image

13 +from transformers import PreTrainedTokenizer, PreTrainedTokenizerFast

15 from vllm.config import ModelConfig

16 from vllm.inputs import InputProcessingContext

13 -from vllm.multimodal import MULTIMODAL_REGISTRY

14 -from vllm.multimodal.processing import ProcessingCache

15 -from vllm.transformers_utils.tokenizer import cached_tokenizer_from_config

17 +from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalDataDict

18 +from vllm.multimodal.inputs import MultiModalInputs

19 +from vllm.multimodal.processing import BaseMultiModalProcessor, ProcessingCache

20 +from vllm.transformers_utils.tokenizer import (MistralTokenizer,

21 + cached_tokenizer_from_config)

23 from ....multimodal.utils import random_audio, random_image, random_video

24 from ...registry import HF_EXAMPLE_MODELS

@@ -85,14 +91,6 @@ def _test_processing_correctness(

91 partial(random_audio, rng, min_len=512, max_len=1024, sr=16000),

92 }

88 - tokenizer_encode_kwargs = {}

89 - if model_config.hf_config.model_type in ("mllama", "whisper", "ultravox"):

90 - # For some multimodal models, tokenizer will always add bos_token

91 - # at the beginning of prompt by default, causing hf_processor outputs

92 - # incorrect token ids. So we need use `add_special_tokens=False` here

93 - # to leave bos_token to be added by the processor.

94 - tokenizer_encode_kwargs = {"add_special_tokens": False}

95 -

94 for batch_idx in range(num_batches):

95 mm_data = {

96 k:

@@ -115,43 +113,131 @@ def _test_processing_correctness(

113 elif len(mm_data[k]) == 1:

114 mm_data[k] = mm_data[k][0]

115

118 - baseline_result = baseline_processor.apply(

119 - prompt,

120 - mm_data=mm_data,

121 - hf_processor_mm_kwargs={},

122 - )

123 - cached_result = cached_processor.apply(

124 - prompt,

125 - mm_data=mm_data,

126 - hf_processor_mm_kwargs={},

127 - )

128 -

129 - assert _drop_mm_kwargs_keys(

130 - baseline_result, ignore_mm_keys) == _drop_mm_kwargs_keys(

131 - cached_result, ignore_mm_keys), (

132 - f"Failed ({batch_idx=}, {prompt=}, {mm_data=})")

133 -

134 - baseline_tokenized_result = baseline_processor.apply(

135 - tokenizer.encode(prompt, **tokenizer_encode_kwargs),

136 - mm_data=mm_data,

137 - hf_processor_mm_kwargs={},

138 - )

139 -

140 - assert _drop_mm_kwargs_keys(

141 - baseline_result, ignore_mm_keys) == _drop_mm_kwargs_keys(

142 - baseline_tokenized_result, ignore_mm_keys), (

143 - f"Failed ({batch_idx=}, {prompt=}, {mm_data=})")

144 -

145 - cached_tokenized_result = cached_processor.apply(

146 - tokenizer.encode(prompt, **tokenizer_encode_kwargs),

147 - mm_data=mm_data,

148 - hf_processor_mm_kwargs={},

149 - )

150 -

151 - assert _drop_mm_kwargs_keys(

152 - cached_result, ignore_mm_keys) == _drop_mm_kwargs_keys(

153 - cached_tokenized_result, ignore_mm_keys), (

154 - f"Failed ({batch_idx=}, {prompt=}, {mm_data=})")

116 + if isinstance(tokenizer, MistralTokenizer):

117 + _test_processing_correctness_mistral(

118 + model_config,

119 + tokenizer,

120 + prompt,

121 + mm_data,

122 + baseline_processor,

123 + cached_processor,

124 + batch_idx,

125 + ignore_mm_keys=ignore_mm_keys,

126 + )

127 + else:

128 + _test_processing_correctness_hf(

129 + model_config,

130 + tokenizer,

131 + prompt,

132 + mm_data,

133 + baseline_processor,

134 + cached_processor,

135 + batch_idx,

136 + ignore_mm_keys=ignore_mm_keys,

137 + )

138 +

139 +

140 +def _test_processing_correctness_hf(

141 + model_config: ModelConfig,

142 + tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast],

143 + prompt: str,

144 + mm_data: MultiModalDataDict,

145 + baseline_processor: BaseMultiModalProcessor,

146 + cached_processor: BaseMultiModalProcessor,

147 + batch_idx: int,

148 + ignore_mm_keys: Optional[list[str]] = None,

149 +):

150 + if model_config.hf_config.model_type in ("mllama", "whisper", "ultravox"):

151 + # For some multimodal models, tokenizer will always add bos_token

152 + # at the beginning of prompt by default, causing hf_processor outputs

153 + # incorrect token ids. So we need use `add_special_tokens=False` here

154 + # to leave bos_token to be added by the processor.

155 + token_prompt = tokenizer.encode(prompt, add_special_tokens=False)

156 + else:

157 + token_prompt = tokenizer.encode(prompt)

158 +

159 + baseline_result = baseline_processor.apply(

160 + prompt,

161 + mm_data=mm_data,

162 + hf_processor_mm_kwargs={},

163 + )

164 + cached_result = cached_processor.apply(

165 + prompt,

166 + mm_data=mm_data,

167 + hf_processor_mm_kwargs={},

168 + )

169 +

170 + assert _inputs_equal(

171 + baseline_result,

172 + cached_result,

173 + ignore_mm_keys,

174 + ), f"Failed ({batch_idx=}, {prompt=}, {mm_data=})"

175 +

176 + baseline_tokenized_result = baseline_processor.apply(

177 + token_prompt,

178 + mm_data=mm_data,

179 + hf_processor_mm_kwargs={},

180 + )

181 +

182 + assert _inputs_equal(

183 + baseline_result,

184 + baseline_tokenized_result,

185 + ignore_mm_keys,

186 + ), f"Failed ({batch_idx=}, {prompt=}, {mm_data=})"

187 +

188 + cached_tokenized_result = cached_processor.apply(

189 + token_prompt,

190 + mm_data=mm_data,

191 + hf_processor_mm_kwargs={},

192 + )

193 +

194 + assert _inputs_equal(

195 + cached_result,

196 + cached_tokenized_result,

197 + ignore_mm_keys,

198 + ), f"Failed ({batch_idx=}, {prompt=}, {mm_data=})"

199 +

200 +

201 +def _test_processing_correctness_mistral(

202 + model_config: ModelConfig,

203 + tokenizer: MistralTokenizer,

204 + prompt: str,

205 + mm_data: MultiModalDataDict,

206 + baseline_processor: BaseMultiModalProcessor,

207 + cached_processor: BaseMultiModalProcessor,

208 + batch_idx: int,

209 + ignore_mm_keys: Optional[list[str]] = None,

210 +):

211 + images = mm_data.get("image", [])

212 + if not isinstance(images, list):

213 + images = [images]

214 +

215 + request = ChatCompletionRequest(messages=[

216 + UserMessage(content=[

217 + TextChunk(text=prompt),

218 + *(ImageChunk(image=image) for image in images),

219 + ]),

220 + ])

221 + res = tokenizer.mistral.encode_chat_completion(request)

222 + token_prompt = res.tokens

223 +

224 + # Mistral chat outputs tokens directly, rather than text prompts

225 + baseline_tokenized_result = baseline_processor.apply(

226 + token_prompt,

227 + mm_data=mm_data,

228 + hf_processor_mm_kwargs={},

229 + )

230 + cached_tokenized_result = cached_processor.apply(

231 + token_prompt,

232 + mm_data=mm_data,

233 + hf_processor_mm_kwargs={},

234 + )

235 +

236 + assert _inputs_equal(

237 + baseline_tokenized_result,

238 + cached_tokenized_result,

239 + ignore_mm_keys,

240 + ), f"Failed ({batch_idx=}, {prompt=}, {mm_data=})"

241

242

243 # yapf: disable

@@ -173,6 +259,7 @@ def _test_processing_correctness(

259 "llava-hf/llava-onevision-qwen2-0.5b-ov-hf",

260 "meta-llama/Llama-3.2-11B-Vision-Instruct",

261 "TIGER-Lab/Mantis-8B-siglip-llama3",

262 + "mistralai/Pixtral-12B-2409",

263 "mistral-community/pixtral-12b",

264 "openbmb/MiniCPM-o-2_6",

265 "openbmb/MiniCPM-V-2_6",

@@ -241,8 +328,19 @@ def test_processing_correctness_phi3v(

328 )

329

330

244 -def _drop_mm_kwargs_keys(result: dict,

245 - ignore_mm_keys: Optional[list[str]] = None) -> dict:

331 +def _inputs_equal(

332 + a: MultiModalInputs,

333 + b: MultiModalInputs,

334 + ignore_mm_keys: Optional[list[str]] = None,

335 +):

336 + return _drop_mm_kwargs_keys(a, ignore_mm_keys) == _drop_mm_kwargs_keys(

337 + b, ignore_mm_keys)

338 +

339 +

340 +def _drop_mm_kwargs_keys(

341 + result: MultiModalInputs,

342 + ignore_mm_keys: Optional[list[str]] = None,

343 +) -> MultiModalInputs:

344 """Drop specified keys from result['mm_kwargs'].

345

346 This is mainly to avoid doing exact match of audio_features in ultravox.

vllm/model_executor/models/llava.py

+50 -76

@@ -68,23 +68,15 @@ class PixtralHFImagePixelInputs(TypedDict):

68 in which case the data is passed as a list instead of a batched tensor.

69 """

71 - feat_is_patch: Union[torch.Tensor, list[torch.Tensor]]

72 - """

73 - A boolean mask indicating which image features correspond

74 - to patch tokens.

75 -

76 - Shape: `(batch_size, num_crops, num_patch)`

77 - """

78 -

71 embed_is_patch: Union[torch.Tensor, list[torch.Tensor]]

72 """

73 A boolean mask indicating which image embeddings correspond

74 to patch tokens.

84 - Shape: `(batch_size, num_embeds)`

76 + Shape: `(batch_size, num_images, num_embeds)`

77 """

87 - num_crops: Union[torch.Tensor, list[torch.Tensor]]

79 + num_patches: Union[torch.Tensor, list[torch.Tensor]]

80 """Shape: `(batch_size, num_images)`"""

@@ -360,16 +352,16 @@ class PixtralHFMultiModalProcessor(

352 image_height=pixel_value.shape[-2],

353 ) for pixel_value in processed_outputs["pixel_values"]

354 ]

363 - num_crops = torch.tensor([(ncols + 1) * nrows

364 - for ncols, nrows in tile_sizes])

355 + num_patches = torch.tensor([(ncols + 1) * nrows

356 + for ncols, nrows in tile_sizes])

357 # Each image may result to masks of different sizes, so we need to

366 - # flatten the list and later use `num_crops` to get per-image masks.

367 - embed_is_patch = torch.tensor(

368 - flatten_2d_lists([([True] * ncols + [False]) * nrows

369 - for ncols, nrows in tile_sizes]))

370 - processed_outputs["num_crops"] = num_crops

358 + # later use `num_patches` to get per-image masks.

359 + embed_is_patch = [

360 + torch.tensor(([True] * ncols + [False]) * nrows)

361 + for ncols, nrows in tile_sizes

362 + ]

363 + processed_outputs["num_patches"] = num_patches

364 processed_outputs["embed_is_patch"] = embed_is_patch

372 - processed_outputs["feat_is_patch"] = embed_is_patch

365

366 return processed_outputs

367

@@ -378,14 +370,10 @@ class PixtralHFMultiModalProcessor(

370 hf_inputs: BatchFeature,

371 hf_processor_mm_kwargs: Mapping[str, object],

372 ) -> Mapping[str, MultiModalFieldConfig]:

381 - num_crops = hf_inputs.get("num_crops", torch.empty(0)).view(-1)

373 return dict(

383 - feat_is_patch=MultiModalFieldConfig.flat_from_sizes(

384 - "image", num_crops),

385 - embed_is_patch=MultiModalFieldConfig.flat_from_sizes(

386 - "image", num_crops),

387 - num_crops=MultiModalFieldConfig.batched("image"),

374 pixel_values=MultiModalFieldConfig.batched("image"),

375 + num_patches=MultiModalFieldConfig.batched("image"),

376 + embed_is_patch=MultiModalFieldConfig.batched("image"),

377 image_embeds=MultiModalFieldConfig.batched("image"),

378 )

379

@@ -628,27 +616,21 @@ class LlavaForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP):

616 f"Got type: {type(pixel_values)}")

617

618 if self.config.vision_config.model_type == "pixtral":

631 - feat_is_patch = kwargs.pop("feat_is_patch")

632 - if not isinstance(feat_is_patch, (torch.Tensor, list)):

633 - raise ValueError("Incorrect type of feat_is_patch. "

634 - f"Got type: {type(feat_is_patch)}")

635 -

619 embed_is_patch = kwargs.pop("embed_is_patch")

620 if not isinstance(embed_is_patch, (torch.Tensor, list)):

621 raise ValueError("Incorrect type of embed_is_patch. "

622 f"Got type: {type(embed_is_patch)}")

623

641 - num_crops = kwargs.pop("num_crops")

642 - if not isinstance(num_crops, (torch.Tensor, list)):

643 - raise ValueError("Incorrect type of num_crops. "

644 - f"Got type: {type(num_crops)}")

624 + num_patches = kwargs.pop("num_patches")

625 + if not isinstance(num_patches, (torch.Tensor, list)):

626 + raise ValueError("Incorrect type of num_patches. "

627 + f"Got type: {type(num_patches)}")

628

629 return PixtralHFImagePixelInputs(

630 type="pixel_values_pixtral",

631 pixel_values=flatten_bn(pixel_values),

649 - feat_is_patch=feat_is_patch,

632 embed_is_patch=embed_is_patch,

651 - num_crops=num_crops,

633 + num_patches=num_patches,

634 )

635

636 return LlavaImagePixelInputs(

@@ -687,21 +669,26 @@ class LlavaForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP):

669 vision_tower: Union[CLIPVisionModel, SiglipVisionModel,

670 PixtralHFVisionModel],

671 pixel_values: Union[torch.Tensor, list[torch.Tensor]],

690 - ) -> torch.Tensor:

691 -

672 + ) -> Union[torch.Tensor, tuple[torch.Tensor, ...]]:

673 # NOTE: we skip the step to select the vision feature layer since

674 # this is already done inside the vision tower

675 image_features = vision_tower(pixel_values)

676

696 - return self._select_image_features(

697 - image_features,

698 - strategy=self.config.vision_feature_select_strategy,

677 + def select_features(leaf: torch.Tensor):

678 + return self._select_image_features(

679 + leaf,

680 + strategy=self.config.vision_feature_select_strategy,

681 + )

682 +

683 + return cast(

684 + Union[torch.Tensor, tuple[torch.Tensor, ...]],

685 + json_map_leaves(select_features, image_features),

686 )

687

688 def _process_image_pixels(

689 self,

690 inputs: Union[LlavaImagePixelInputs, PixtralHFImagePixelInputs],

704 - ) -> torch.Tensor:

691 + ) -> Union[torch.Tensor, tuple[torch.Tensor, ...]]:

692 assert self.vision_tower is not None

693

694 pixel_values = inputs["pixel_values"]

@@ -731,45 +718,30 @@ class LlavaForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP):

718

719 def _get_mm_embeds(

720 self,

734 - features: torch.Tensor, # Shape: (num_crop, num_patch, d)

735 - feat_is_patch: torch.Tensor, # Shape: (num_crop, num_patch)

736 - num_crops: torch.Tensor, # Shape: (num_images,)

737 - embed_is_patch: torch.Tensor, # Shape: (num_embeds,)

738 - ) -> list[torch.Tensor]:

721 + features: torch.Tensor, # Shape: (num_patch, d)

722 + num_patches: torch.Tensor, # Shape: (num_images,)

723 + embed_is_patch: torch.Tensor, # Shape: (num_images, num_embeds)

724 + ) -> tuple[torch.Tensor, ...]:

725 """Scatter the patch features into a contiguous tensor that corresponds

726 to the embedding tokens defined by the multimodal processor.

727

728 Mostly copied from `Molmo._get_mm_embeds`. See following fixme comment.

729 """

744 -

745 - # Insert columns of nan values according to `feat_is_patch`. This work

730 + # Insert columns of nan values according to `embed_is_patch`. This work

731 # ideally should be done in `_process_image_input`, but

732 # `_process_image_input` is used in both V0 and V1 path. It's safer to

733 # put the logic here.

734 # FIXME: Move this logic to `_process_image_input` when v0 is

735 # deprecated. Merge this function with `Molmo._get_mm_embeds`.

751 - feat_is_patch = feat_is_patch.view(-1)

752 - embed_is_patch = embed_is_patch.view(-1)

753 - expanded_embedding = torch.full(

754 - (sum(num_crops), *features.shape[1:]),

755 - torch.nan,

756 - dtype=features.dtype).to(features.device)

757 - expanded_embedding[feat_is_patch] = features

736 + num_patches_per_image: list[int] = num_patches.tolist()

737

759 - num_crops_per_image = num_crops.tolist()

760 - feats_per_image = expanded_embedding.split(num_crops_per_image)

761 - f_is_patch_per_image = feat_is_patch.split(num_crops_per_image)

762 -

763 - embed_dim = expanded_embedding.shape[-1]

764 - num_embeds = embed_is_patch.shape[0]

765 -

766 - embeds_in_batch = list[torch.Tensor]()

767 - for feats, f_is_patch in zip(feats_per_image, f_is_patch_per_image):

768 - embeds = feats.new_full((num_embeds, embed_dim), torch.nan)

769 - embeds[embed_is_patch] = feats[f_is_patch]

770 - embeds_in_batch.append(embeds)

738 + embeds_flat = features.new_full(

739 + (sum(num_patches_per_image), *features.shape[1:]),

740 + fill_value=torch.nan,

741 + )

742 + embeds_flat[embed_is_patch.view(-1)] = features

743

772 - return embeds_in_batch

744 + return embeds_flat.split(num_patches_per_image)

745

746 def get_multimodal_embeddings(

747 self, **kwargs: object) -> Optional[MultiModalEmbeddings]:

@@ -784,12 +756,12 @@ class LlavaForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP):

756 # The path is used for pixtral (V0 only) and llava (V0/V1)

757 return vision_embeddings

758

787 - nested_emb = [

759 + return flatten_2d_lists(

760 self._get_mm_embeds(*args) for args in zip(

789 - vision_embeddings, image_input["feat_is_patch"],

790 - image_input["num_crops"], image_input["embed_is_patch"])

791 - ]

792 - return flatten_2d_lists(nested_emb)

761 + vision_embeddings,

762 + image_input["num_patches"],

763 + image_input["embed_is_patch"],

764 + ))

765

766 def get_input_embeddings(

767 self,

@@ -805,9 +777,11 @@ class LlavaForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP):

777 )

778

779 inputs_embeds = merge_multimodal_embeddings(

808 - input_ids, inputs_embeds, cast(NestedTensors,

809 - patch_embeddings),

810 - self.config.image_token_index)

780 + input_ids,

781 + inputs_embeds,

782 + cast(NestedTensors, patch_embeddings),

783 + self.config.image_token_index,

784 + )

785 return inputs_embeds

786

787 def forward(

vllm/model_executor/models/molmo.py

+2 -4

@@ -1585,15 +1585,13 @@ class MolmoForCausalLM(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA,

1585

1586 image_features = self._process_image_input(image_input)

1587

1588 - nested_embeds = [

1588 + return flatten_2d_lists(

1589 self._get_mm_embeds(*args) for args in zip(

1590 image_features,

1591 image_input["feat_is_patch"],

1592 image_input["num_crops"],

1593 image_input["embed_is_patch"],

1594 - )

1595 - ]

1596 - return flatten_2d_lists(nested_embeds)

1594 + ))

1595

1596 def get_input_embeddings(

1597 self,

vllm/model_executor/models/paligemma.py

+4 -5

@@ -1,7 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 -

3 -from typing import (Iterable, Literal, Mapping, Optional, Set, Tuple,

4 - TypedDict, Union)

2 +from collections.abc import Iterable, Mapping, Sequence

3 +from typing import Literal, Optional, Set, Tuple, TypedDict, Union

5 import torch

6 from torch import nn

@@ -17,7 +16,7 @@ from vllm.multimodal.inputs import (MultiModalDataDict, MultiModalFieldConfig,

16 from vllm.multimodal.parse import MultiModalDataItems

17 from vllm.multimodal.processing import (BaseMultiModalProcessor,

18 BaseProcessingInfo, PromptIndexTargets,

20 - PromptInsertion, PromptReplacement,

19 + PromptInsertion, PromptUpdate,

20 PromptUpdateDetails)

21 from vllm.multimodal.profiling import BaseDummyInputsBuilder, ProcessorInputs

22 from vllm.sequence import IntermediateTensors

@@ -144,7 +143,7 @@ class PaliGemmaMultiModalProcessor(

143 mm_items: MultiModalDataItems,

144 hf_processor_mm_kwargs: Mapping[str, object],

145 out_mm_kwargs: MultiModalKwargs,

147 - ) -> list[PromptReplacement]:

146 + ) -> Sequence[PromptUpdate]:

147 hf_config = self.info.get_hf_config()

148 image_token_id = hf_config.image_token_index

149

vllm/model_executor/models/pixtral.py

+374 -208

@@ -1,26 +1,28 @@

1 # SPDX-License-Identifier: Apache-2.0

3 import math

4 -from collections.abc import Iterable, Mapping

4 +from collections.abc import Iterable, Mapping, Sequence

5 from dataclasses import dataclass, fields

6 from functools import cached_property

7 -from typing import List, Optional, Set, Tuple, Union

7 +from typing import List, Literal, Optional, Set, Tuple, TypedDict, Union, cast

9 import torch

10 import torch.nn as nn

11 import torch.nn.functional as F

12 from mistral_common.protocol.instruct.messages import ImageChunk

13 +from mistral_common.tokens.tokenizers.multimodal import ImageEncoder

14 from PIL import Image

14 -from transformers import PixtralVisionConfig

15 +from transformers import PixtralVisionConfig, TensorType

16 +from transformers.image_utils import ImageInput

17 from transformers.models.pixtral.image_processing_pixtral import (

18 _num_image_tokens as _get_pixtral_hf_num_image_tokens)

19 from transformers.models.pixtral.modeling_pixtral import (

20 PixtralRotaryEmbedding, apply_rotary_pos_emb, position_ids_in_meshgrid)

21 +from transformers.tokenization_utils_base import TextInput

23 from vllm.config import VllmConfig

24 from vllm.distributed import divide, get_tensor_model_parallel_world_size

22 -from vllm.inputs import (INPUT_REGISTRY, DecoderOnlyInputs, DummyData,

23 - InputContext, token_inputs)

25 +from vllm.jsontree import JSONTree, json_map_leaves

26 from vllm.model_executor.layers.activation import get_act_and_mul_fn

27 from vllm.model_executor.layers.layernorm import RMSNorm

28 from vllm.model_executor.layers.linear import (MergedColumnParallelLinear,

@@ -31,13 +33,20 @@ from vllm.model_executor.layers.sampler import SamplerOutput, get_sampler

33 from vllm.model_executor.model_loader.weight_utils import default_weight_loader

34 from vllm.model_executor.sampling_metadata import SamplingMetadata

35 from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalKwargs

34 -from vllm.multimodal.inputs import PlaceholderRange

35 -from vllm.multimodal.utils import consecutive_placeholder_ranges

36 -from vllm.sequence import IntermediateTensors, SequenceData

37 -from vllm.transformers_utils.tokenizer import cached_tokenizer_from_config

36 +from vllm.multimodal.inputs import MultiModalFieldConfig, NestedTensors

37 +from vllm.multimodal.parse import (ImageProcessorItems, ImageSize,

38 + MultiModalDataItems)

39 +from vllm.multimodal.processing import (BaseMultiModalProcessor,

40 + BaseProcessingInfo, PromptReplacement,

41 + PromptUpdate)

42 +from vllm.multimodal.profiling import BaseDummyInputsBuilder, ProcessorInputs

43 +from vllm.sequence import IntermediateTensors

44 +from vllm.transformers_utils.tokenizer import (MistralTokenizer,

45 + cached_tokenizer_from_config)

46 +from vllm.utils import flatten_2d_lists

48 from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP

40 -from .utils import (init_vllm_registered_model, maybe_prefix,

49 +from .utils import (flatten_bn, init_vllm_registered_model, maybe_prefix,

50 merge_multimodal_embeddings)

51 from .vision import VisionEncoderInfo, resolve_visual_encoder_outputs

@@ -48,132 +57,275 @@ except ImportError:

57 USE_XFORMERS_OPS = False

51 -def get_max_pixtral_image_tokens(ctx: InputContext):

52 - tokenizer = cached_tokenizer_from_config(ctx.model_config)

53 - mm_encoder = tokenizer.instruct.mm_encoder

60 +class PixtralImagePixelInputs(TypedDict):

61 + type: Literal["pixel_values"]

55 - image_config = mm_encoder.mm_config if hasattr(

56 - mm_encoder, "mm_config") else mm_encoder.image_config

63 + images: Union[torch.Tensor, list[torch.Tensor]]

64 + """

65 + Shape: `(batch_size * num_images, num_channels, image_width, image_height)`

58 - max_image_size = image_config.max_image_size

59 - image_patch_size = image_config.image_patch_size

67 + The result of stacking :attr:`ImageEncoding.tokens` from each prompt.

68 + """

61 - return ((max_image_size // image_patch_size)**2)

70 + embed_is_patch: Union[torch.Tensor, list[torch.Tensor]]

71 + """

72 + A boolean mask indicating which image embeddings correspond

73 + to patch tokens.

74 +

75 + Shape: `(batch_size, num_images, num_embeds)`

76 + """

78 + num_patches: Union[torch.Tensor, list[torch.Tensor]]

79 + """Shape: `(batch_size, num_images)`"""

64 -def dummy_data_for_pixtral(ctx: InputContext, seq_len: int,

65 - mm_counts: Mapping[str, int]):

66 - tokenizer = cached_tokenizer_from_config(ctx.model_config)

68 - mm_encoder = tokenizer.mistral.instruct_tokenizer.mm_encoder

69 - image_token_id = mm_encoder.special_ids.img

82 +class PixtralProcessorAdapter:

83 + """

84 + Provide a HF-compatible interface for

85 + :class:`mistral_common.tokens.tokenizers.multimodal.ImageEncoder`.

86 + """

71 - mm_config = ctx.get_mm_config()

72 - num_images = mm_config.get_limit_per_prompt("image")

88 + def __init__(self, tokenizer: MistralTokenizer) -> None:

89 + super().__init__()

74 - # dummy size

75 - size = 256

76 - image = Image.new("RGB", (size, size), color=0)

91 + self.tokenizer = tokenizer

78 - encoding = tokenizer.instruct.mm_encoder(ImageChunk(image=image))

79 - image_feature_size = len(encoding.tokens)

80 - num_image_tokens = image_feature_size * num_images

81 - seq_data = SequenceData.from_prompt_token_counts(

82 - (image_token_id, num_image_tokens),

83 - (0, seq_len - num_image_tokens),

84 - )

93 + @property

94 + def image_processor(self) -> ImageEncoder:

95 + image_encoder = self.tokenizer.instruct.mm_encoder

96 + assert isinstance(image_encoder, ImageEncoder)

97 + return image_encoder

86 - mm_data = {"image": num_images * [image]}

87 - mm_placeholders = {

88 - "image":

89 - consecutive_placeholder_ranges(num_items=num_images,

90 - item_size=image_feature_size)

91 - }

92 - return DummyData(seq_data, mm_data, mm_placeholders)

99 + @cached_property

100 + def image_break_id(self) -> int:

101 + return self.image_processor.special_ids.img_break

102

103 + @cached_property

104 + def image_token_id(self) -> int:

105 + return self.image_processor.special_ids.img

106

95 -def input_mapper_for_pixtral(ctx: InputContext,

96 - data: object) -> MultiModalKwargs:

97 - """Maps the input data to its MultiModalKwargs (if any).

107 + @cached_property

108 + def image_end_id(self) -> int:

109 + return self.image_processor.special_ids.img_end

110

99 - Args:

100 - ctx: Context of the loaded model.

101 - data: data potentially containing PIL images to be processed

102 - and mapped to `images`.

111 + @cached_property

112 + def image_size(self) -> int:

113 + return self.image_processor.mm_config.max_image_size

114

104 - Returns:

105 - MultiModalKwargs containing the stacked normalized images tensor or

106 - image embeddings.

107 - """

108 - tokenizer = cached_tokenizer_from_config(ctx.model_config)

109 -

110 - data_list = data if isinstance(data, list) else [data]

111 -

112 - images = []

113 - image_tokens_list = []

114 - for image_data in data_list:

115 - image = ImageChunk(image=image_data)

116 - encoding = tokenizer.instruct.mm_encoder(image)

117 - image = torch.from_numpy(encoding.image).to(dtype=torch.float16)

118 - images.append(image)

119 - image_tokens_list.append(encoding.tokens)

120 -

121 - image_tokens = torch.tensor([

122 - token_id for image_tokens in image_tokens_list

123 - for token_id in image_tokens

124 - ])

125 - return MultiModalKwargs({"images": images, "image_tokens": image_tokens})

126 -

127 -

128 -def input_processor_for_pixtral(ctx: InputContext, inputs: DecoderOnlyInputs):

129 - multi_modal_data = inputs.get("multi_modal_data")

130 - if multi_modal_data is None or "image" not in multi_modal_data:

131 - return inputs

132 -

133 - prompt_token_ids = inputs.get("prompt_token_ids")

134 - prompt = inputs.get("prompt")

135 - tokenizer = cached_tokenizer_from_config(ctx.model_config)

136 -

137 - mm_encoder = tokenizer.mistral.instruct_tokenizer.mm_encoder

138 - image_token_id = mm_encoder.special_ids.img

139 - image_break_id = mm_encoder.special_ids.img_break

140 - image_end_id = mm_encoder.special_ids.img_end

141 -

142 - if image_token_id not in inputs['prompt_token_ids']:

143 - raise ValueError(

144 - f"You've passed {inputs=} without {image_token_id=}"

145 - " Make sure to process your input via mistral_common's"

146 - " tokenizer or pass a chat completion request. For more"

147 - " For more info, see: "

148 - "https://github.com/vllm-project/vllm/issues/8411.")

149 -

150 - # Get precise tracking of placeholder positions

151 - placeholder_ranges = []

152 - curr_offset = -1

153 - curr_length = 0

154 - for i in range(len(prompt_token_ids)):

155 - if prompt_token_ids[i] in (image_token_id, image_break_id):

156 - if curr_offset < 0:

157 - curr_offset = i

158 - curr_length += 1

159 - elif prompt_token_ids[i] == image_end_id:

160 - curr_length += 1

161 - placeholder_ranges.append(

162 - PlaceholderRange(offset=curr_offset, length=curr_length))

163 - curr_offset = -1

164 - curr_length = 0

165 - else:

166 - pass

167 - return token_inputs(prompt=prompt,

168 - prompt_token_ids=prompt_token_ids,

169 - multi_modal_data=multi_modal_data,

170 - multi_modal_placeholders={"image": placeholder_ranges})

115 + @cached_property

116 + def patch_size(self) -> int:

117 + return self.image_processor.mm_config.image_patch_size

118 +

119 + def __call__(

120 + self,

121 + text: Optional[Union[TextInput, list[TextInput]]] = None,

122 + images: Optional[Union[ImageInput, list[ImageInput]]] = None,

123 + return_tensors: Optional[Union[str, TensorType]] = None,

124 + **kwargs,

125 + ) -> Mapping[str, NestedTensors]:

126 + if text is None:

127 + text = []

128 + if not isinstance(text, list):

129 + text = [text]

130 + if images is None:

131 + images = []

132 + if not isinstance(images, list):

133 + images = [images]

134 +

135 + if not images:

136 + input_ids = self.tokenizer(text).input_ids

137 +

138 + return {"input_ids": torch.tensor(input_ids)}

139 +

140 + # Allow dummy text, which is used for profiling as well as token inputs

141 + if any(len(t) > 0 for t in text):

142 + raise ValueError(

143 + "You've passed text inputs instead of token inputs. "

144 + "Make sure to process your input via `mistral_common`'s "

145 + "tokenizer or pass a chat completion request. "

146 + "For more info, see: "

147 + "https://github.com/vllm-project/vllm/issues/8411.")

148 +

149 + image_token_id = self.image_token_id

150 +

151 + images_processed = list[torch.Tensor]()

152 + images_tokens = list[torch.Tensor]()

153 + images_embed_is_patch = list[torch.Tensor]()

154 + images_num_patches = list[int]()

155 +

156 + for image in images:

157 + image_inputs = self.image_processor(ImageChunk(image=image))

158 +

159 + image_processed = torch.tensor(image_inputs.image)

160 + image_tokens = torch.tensor(image_inputs.tokens)

161 +

162 + images_processed.append(image_processed)

163 + images_tokens.append(image_tokens)

164 + images_embed_is_patch.append(image_tokens == image_token_id)

165 + images_num_patches.append(len(image_tokens))

166 +

167 + return {

168 + "input_ids": torch.cat(images_tokens)[None].expand(len(text), -1),

169 + "images": images_processed,

170 + "embed_is_patch": images_embed_is_patch,

171 + "num_patches": torch.tensor(images_num_patches),

172 + }

173 +

174 +

175 +class PixtralProcessingInfo(BaseProcessingInfo):

176 +

177 + def get_tokenizer(self) -> MistralTokenizer:

178 + tokenizer = cached_tokenizer_from_config(self.ctx.model_config)

179 + if not isinstance(tokenizer, MistralTokenizer):

180 + raise ValueError("This model requires `--tokenizer-mode mistral`")

181 +

182 + return tokenizer

183 +

184 + def get_hf_processor(self) -> PixtralProcessorAdapter:

185 + return PixtralProcessorAdapter(self.get_tokenizer())

186 +

187 + def get_supported_mm_limits(self) -> Mapping[str, Optional[int]]:

188 + return {"image": None}

189 +

190 + def get_mm_max_tokens_per_item(

191 + self,

192 + seq_len: int,

193 + mm_counts: Mapping[str, int],

194 + ) -> Mapping[str, int]:

195 + return {"image": self.get_max_image_tokens()}

196 +

197 + def get_vision_config(

198 + self,

199 + processor: Optional[PixtralProcessorAdapter] = None,

200 + ):

201 + if processor is None:

202 + processor = self.get_hf_processor()

203 +

204 + return PixtralVisionConfig(

205 + image_size=processor.image_size,

206 + patch_size=processor.patch_size,

207 + )

208 +

209 + def get_num_image_tokens(

210 + self,

211 + *,

212 + image_width: int,

213 + image_height: int,

214 + processor: Optional[PixtralProcessorAdapter] = None,

215 + ) -> int:

216 + if processor is None:

217 + processor = self.get_hf_processor()

218 +

219 + ncols, nrows = processor.image_processor._image_to_num_tokens(

220 + Image.new("RGB", (image_width, image_height)))

221 +

222 + return (ncols + 1) * nrows

223 +

224 + def get_image_size_with_most_features(self) -> ImageSize:

225 + image_processor = self.get_hf_processor().image_processor

226 + max_image_size = image_processor.mm_config.max_image_size

227 +

228 + return ImageSize(width=max_image_size, height=max_image_size)

229 +

230 + def get_max_image_tokens(self) -> int:

231 + target_width, target_height = self.get_image_size_with_most_features()

232 +

233 + return self.get_num_image_tokens(

234 + image_width=target_width,

235 + image_height=target_height,

236 + )

237 +

238 +

239 +class PixtralDummyInputsBuilder(BaseDummyInputsBuilder[PixtralProcessingInfo]):

240 +

241 + def get_dummy_processor_inputs(

242 + self,

243 + seq_len: int,

244 + mm_counts: Mapping[str, int],

245 + ) -> ProcessorInputs:

246 + num_images = mm_counts.get("image", 0)

247 +

248 + target_width, target_height = \

249 + self.info.get_image_size_with_most_features()

250 +

251 + mm_data = {

252 + "image":

253 + self._get_dummy_images(width=target_width,

254 + height=target_height,

255 + num_images=num_images)

256 + }

257 +

258 + return ProcessorInputs(

259 + prompt_text="",

260 + mm_data=mm_data,

261 + )

262 +

263 +

264 +class PixtralMultiModalProcessor(BaseMultiModalProcessor[PixtralProcessingInfo]

265 + ):

266

267 + def _get_mm_fields_config(

268 + self,

269 + hf_inputs: Mapping[str, NestedTensors],

270 + hf_processor_mm_kwargs: Mapping[str, object],

271 + ) -> Mapping[str, MultiModalFieldConfig]:

272 + return dict(

273 + images=MultiModalFieldConfig.batched("image"),

274 + embed_is_patch=MultiModalFieldConfig.batched("image"),

275 + num_patches=MultiModalFieldConfig.batched("image"),

276 + )

277 +

278 + def _get_prompt_updates(

279 + self,

280 + mm_items: MultiModalDataItems,

281 + hf_processor_mm_kwargs: Mapping[str, object],

282 + out_mm_kwargs: MultiModalKwargs,

283 + ) -> Sequence[PromptUpdate]:

284 + processor = self.info.get_hf_processor(**hf_processor_mm_kwargs)

285 +

286 + image_break_id = processor.image_break_id

287 + image_token_id = processor.image_token_id

288 + image_end_id = processor.image_end_id

289 +

290 + def get_replacement(item_idx: int):

291 + images = mm_items.get_items("image", ImageProcessorItems)

292 + image_size = images.get_image_size(item_idx)

293 +

294 + ncols, nrows = processor.image_processor._image_to_num_tokens(

295 + Image.new("RGB", (image_size.width, image_size.height)))

296 +

297 + tokens = ([image_token_id] * ncols + [image_break_id]) * nrows

298 + tokens[-1] = image_end_id

299 +

300 + return tokens

301 +

302 + return [

303 + PromptReplacement(

304 + modality="image",

305 + target="", # Never match the prompt (see below note)

306 + replacement=get_replacement,

307 + ),

308 + ]

309 +

310 + def _cached_apply_hf_processor(

311 + self,

312 + prompt: Union[str, list[int]],

313 + mm_data_items: MultiModalDataItems,

314 + hf_processor_mm_kwargs: Mapping[str, object],

315 + ) -> tuple[list[int], MultiModalKwargs, bool]:

316 + prompt_ids, mm_kwargs, _ = super()._cached_apply_hf_processor(

317 + prompt=prompt,

318 + mm_data_items=mm_data_items,

319 + hf_processor_mm_kwargs=hf_processor_mm_kwargs,

320 + )

321 +

322 + # NOTE: The tokens are already inserted by the chat template

323 + return prompt_ids, mm_kwargs, True

324

173 -@MULTIMODAL_REGISTRY.register_image_input_mapper(input_mapper_for_pixtral)

174 -@MULTIMODAL_REGISTRY.register_max_image_tokens(get_max_pixtral_image_tokens)

175 -@INPUT_REGISTRY.register_dummy_data(dummy_data_for_pixtral)

176 -@INPUT_REGISTRY.register_input_processor(input_processor_for_pixtral)

325 +

326 +@MULTIMODAL_REGISTRY.register_processor(PixtralMultiModalProcessor,

327 + info=PixtralProcessingInfo,

328 + dummy_inputs=PixtralDummyInputsBuilder)

329 class PixtralForConditionalGeneration(nn.Module, SupportsMultiModal,

330 SupportsPP):

331

@@ -191,13 +343,6 @@ class PixtralForConditionalGeneration(nn.Module, SupportsMultiModal,

343 if key in dataclass_fields

344 }

345

194 - if not ("image_break_token_id" in vision_args

195 - and "image_end_token_id" in vision_args):

196 - raise ValueError(

197 - "'image_break_token_id' and 'image_end_token_id' not found "

198 - "in the vision_encoder arguments. Please download the latest "

199 - "version of 'params.json' from the model repository.")

200 -

346 self.vision_args = VisionEncoderArgs(**vision_args)

347

348 # init MistralForCausalLM

@@ -221,36 +366,92 @@ class PixtralForConditionalGeneration(nn.Module, SupportsMultiModal,

366

367 return get_sampler()

368

369 + def _parse_and_validate_image_input(

370 + self, **kwargs: object) -> Optional[PixtralImagePixelInputs]:

371 + images = kwargs.pop("images", None)

372 + if images is None:

373 + return None

374 +

375 + if not isinstance(images, (torch.Tensor, list)):

376 + raise ValueError("Incorrect type of images. "

377 + f"Got type: {type(images)}")

378 +

379 + embed_is_patch = kwargs.pop("embed_is_patch")

380 + if not isinstance(embed_is_patch, (torch.Tensor, list)):

381 + raise ValueError("Incorrect type of embed_is_patch. "

382 + f"Got type: {type(embed_is_patch)}")

383 +

384 + num_patches = kwargs.pop("num_patches")

385 + if not isinstance(num_patches, (torch.Tensor, list)):

386 + raise ValueError("Incorrect type of num_patches. "

387 + f"Got type: {type(num_patches)}")

388 +

389 + return PixtralImagePixelInputs(

390 + type="pixel_values",

391 + images=flatten_bn(images),

392 + embed_is_patch=embed_is_patch,

393 + num_patches=num_patches,

394 + )

395 +

396 + def _process_image_input(

397 + self,

398 + image_input: PixtralImagePixelInputs,

399 + ) -> tuple[torch.Tensor, ...]:

400 + images = image_input["images"]

401 +

402 + image_features = self.vision_encoder(images)

403 + feature_sizes = [

404 + image_feature.shape[0] for image_feature in image_features

405 + ]

406 +

407 + image_embeds = self.vision_language_adapter(torch.cat(image_features))

408 + image_embeds = torch.split(image_embeds, feature_sizes)

409 + return image_embeds

410 +

411 + def _get_mm_embeds(

412 + self,

413 + features: torch.Tensor, # Shape: (num_patch, d)

414 + num_patches: torch.Tensor, # Shape: (num_images,)

415 + embed_is_patch: torch.Tensor, # Shape: (num_images, num_embeds)

416 + ) -> tuple[torch.Tensor, ...]:

417 + """Scatter the patch features into a contiguous tensor that corresponds

418 + to the embedding tokens defined by the multimodal processor.

419 +

420 + Mostly copied from `Molmo._get_mm_embeds`. See following fixme comment.

421 + """

422 + # Insert columns of nan values according to `embed_is_patch`. This work

423 + # ideally should be done in `_process_image_input`, but

424 + # `_process_image_input` is used in both V0 and V1 path. It's safer to

425 + # put the logic here.

426 + # FIXME: Move this logic to `_process_image_input` when v0 is

427 + # deprecated. Merge this function with `Molmo._get_mm_embeds`.

428 + num_patches_per_image: list[int] = num_patches.tolist()

429 +

430 + embeds_flat = features.new_full(

431 + (sum(num_patches_per_image), *features.shape[1:]),

432 + fill_value=torch.nan,

433 + )

434 + embeds_flat[embed_is_patch.view(-1)] = features

435 +

436 + return embeds_flat.split(num_patches_per_image)

437 +

438 def get_multimodal_embeddings(

439 self, **kwargs: object) -> Optional[MultiModalEmbeddings]:

226 - image_input, image_tokens = self._parse_and_validate_image_input(

227 - **kwargs)

440 + image_input = self._parse_and_validate_image_input(**kwargs)

441 if image_input is None:

442 return None

443

231 - vision_embeddings = self._process_image_input(image_input)

232 -

233 - # NOTE: We patch the outputs of the vision encoder with embeddings

234 - # from `[IMG_BREAK]` and `[IMG_END]` tokens.

235 - image_embeds = self.language_model.get_input_embeddings(image_tokens)

236 - image_token_mask = image_tokens == self.vision_args.image_token_id

237 - image_embeds[image_token_mask] = vision_embeddings

444 + image_features = self._process_image_input(image_input)

445

239 - # NOTE: Image embeddings are split into separate tensors for each image

240 - # by the indices of `[IMG_END]` token.

241 - image_end_mask = image_tokens == self.vision_args.image_end_token_id

242 - split_indices = torch.where(image_end_mask)[0] + 1

243 - if len(split_indices) <= 1:

244 - # Do not split, return as tensor of shape [1, fs, hs]

245 - return image_embeds.unsqueeze(0)

446 + if kwargs.get("v0_path", False):

447 + return image_features

448

247 - # If the last split index is the last index in image_tokens, we

248 - # ignore it to avoid empty split tensor

249 - if split_indices[-1] == len(image_tokens):

250 - split_indices = split_indices[:-1]

251 -

252 - image_embeds = image_embeds.tensor_split(split_indices.cpu())

253 - return image_embeds

449 + return flatten_2d_lists(

450 + self._get_mm_embeds(*args) for args in zip(

451 + image_features,

452 + image_input["num_patches"],

453 + image_input["embed_is_patch"],

454 + ))

455

456 def get_input_embeddings(

457 self,

@@ -259,12 +460,17 @@ class PixtralForConditionalGeneration(nn.Module, SupportsMultiModal,

460 ) -> torch.Tensor:

461 inputs_embeds = self.language_model.get_input_embeddings(input_ids)

462 if multimodal_embeddings is not None:

463 + # Extract the patch tokens

464 + patch_embeddings = json_map_leaves(

465 + lambda x: x[~x.isnan()].view(-1, *x.shape[1:]),

466 + cast(JSONTree[torch.Tensor], multimodal_embeddings),

467 + )

468 inputs_embeds = merge_multimodal_embeddings(

263 - input_ids, inputs_embeds, multimodal_embeddings, [

264 - self.vision_args.image_token_id,

265 - self.vision_args.image_break_token_id,

266 - self.vision_args.image_end_token_id,

267 - ])

469 + input_ids,

470 + inputs_embeds,

471 + cast(NestedTensors, patch_embeddings),

472 + self.vision_args.image_token_id,

473 + )

474 return inputs_embeds

475

476 def forward(

@@ -275,14 +481,14 @@ class PixtralForConditionalGeneration(nn.Module, SupportsMultiModal,

481 inputs_embeds: Optional[torch.Tensor] = None,

482 **kwargs: object,

483 ) -> Union[torch.Tensor, IntermediateTensors]:

278 - """Run forward pass for pixtral.

279 - """

484 + """Run forward pass for pixtral."""

485 if intermediate_tensors is not None:

486 inputs_embeds = None

487

488 # NOTE: In v1, inputs_embeds is always generated at model runner, this

489 # condition is for v0 compatibility.

490 elif inputs_embeds is None:

491 + kwargs.update({"v0_path": True})

492 vision_embeddings = self.get_multimodal_embeddings(**kwargs)

493 inputs_embeds = self.get_input_embeddings(input_ids,

494 vision_embeddings)

@@ -295,47 +501,6 @@ class PixtralForConditionalGeneration(nn.Module, SupportsMultiModal,

501

502 return hidden_states

503

298 - def _parse_and_validate_image_input(

299 - self,

300 - images: Optional[Union[List[List[torch.Tensor]], List[torch.Tensor],

301 - torch.Tensor]] = None,

302 - image_tokens: Optional[torch.Tensor] = None,

303 - ) -> Tuple[Optional[List[torch.Tensor]], Optional[torch.Tensor]]:

304 - if images is None:

305 - return None, None

306 -

307 - if isinstance(images, torch.Tensor):

308 - # if passed as batch take all images

309 - N, B, C, W, H = images.shape

310 - images = images.reshape(N * B, C, W, H)

311 - images = [images[i] for i in range(images.size(0))]

312 - elif isinstance(images, list):

313 - # if passed as list flatten lists of tensors

314 - flatten_images = []

315 - for imgs_per_req in images:

316 - imgs_per_req = [

317 - imgs_per_req[i] for i in range(imgs_per_req.size(0))

318 - ] if isinstance(imgs_per_req, torch.Tensor) else imgs_per_req

319 -

320 - flatten_images.extend(imgs_per_req)

321 -

322 - images = flatten_images

323 -

324 - if isinstance(image_tokens, torch.Tensor):

325 - # image_tokens are batched

326 - image_tokens = image_tokens.flatten()

327 - elif isinstance(image_tokens, list):

328 - # image_tokens are of different lengths thus passed as a list

329 - image_tokens = torch.cat(image_tokens)

330 -

331 - assert image_tokens.dim() == 1

332 -

333 - return images, image_tokens

334 -

335 - def _process_image_input(self,

336 - image_input: List[torch.Tensor]) -> torch.Tensor:

337 - return self.vision_language_adapter(self.vision_encoder(image_input))

338 -

504 def compute_logits(

505 self,

506 hidden_states: torch.Tensor,

@@ -400,8 +565,6 @@ class VisionEncoderArgs:

565 num_attention_heads: int

566 rope_theta: float # for rope-2D

567 image_token_id: int

403 - image_break_token_id: int

404 - image_end_token_id: int

568 adapter_bias: bool = True

569

570

@@ -637,9 +800,13 @@ class VisionTransformer(nn.Module):

800 self.patch_conv(img.unsqueeze(0).to(self.dtype)) for img in images

801 ]

802

803 + patch_embeds = [

804 + p.flatten(2).permute(0, 2, 1) for p in patch_embeds_list

805 + ]

806 + embed_sizes = [p.shape[1] for p in patch_embeds]

807 +

808 # flatten to a single sequence

641 - patch_embeds = torch.cat(

642 - [p.flatten(2).permute(0, 2, 1) for p in patch_embeds_list], dim=1)

809 + patch_embeds = torch.cat(patch_embeds, dim=1)

810 patch_embeds = self.ln_pre(patch_embeds)

811

812 # positional embeddings

@@ -655,8 +822,8 @@ class VisionTransformer(nn.Module):

822 "with the Mistral format")

823 out = self.transformer(patch_embeds, mask=mask, freqs_cis=freqs_cis)

824

658 - # remove batch dimension of the single sequence

659 - return out.squeeze(0)

825 + # squeeze dim 0 and split into separate tensors for each image

826 + return torch.split(out.squeeze(0), embed_sizes)

827

828

829 class VisionLanguageAdapter(nn.Module):

@@ -978,9 +1145,9 @@ class PixtralHFVisionModel(nn.Module):

1145

1146 def forward(

1147 self,

981 - pixel_values: List[torch.Tensor],

1148 + pixel_values: list[torch.Tensor],

1149 feature_sample_layers: Optional[list[int]] = None,

983 - ) -> torch.Tensor:

1150 + ) -> tuple[torch.Tensor, ...]:

1151 """

1152 Args:

1153 pixel_values: Each image to be processed will be a separate tensor

@@ -1039,8 +1206,7 @@ class PixtralHFVisionModel(nn.Module):

1206 self.config.num_hidden_layers)

1207

1208 # squeeze dim 0 and split into separate tensors for each image

1042 - out = torch.split(torch.squeeze(out), embed_sizes)

1043 - return out

1209 + return torch.split(out.squeeze(0), embed_sizes)

1210

1211 # (TODO) Add prefix argument for filtering out weights to be loaded

1212 # ref: https://github.com/vllm-project/vllm/pull/7186#discussion_r1734163986

vllm/multimodal/processing.py

+9 -5

@@ -77,7 +77,9 @@ class PromptIndexTargets:

77 else:

78 if isinstance(prefix, str):

79 # Make both `list[int]`

80 - prefix = encode_tokens(tokenizer, prefix)

80 + prefix = encode_tokens(tokenizer,

81 + prefix,

82 + add_special_tokens=False)

84 match_idx = len(prefix)

85 return match_idx if prompt[:match_idx] == prefix else None

@@ -318,7 +320,7 @@ def _cached_encode(

320 tokenizer: AnyTokenizer,

321 text: str,

322 *,

321 - add_special_tokens: bool = False,

323 + add_special_tokens: Optional[bool] = None,

324 ) -> list[int]:

325 return encode_tokens(tokenizer,

326 text,

@@ -330,7 +332,7 @@ def _cached_decode(

332 tokenizer: AnyTokenizer,

333 token_ids: tuple[int, ...],

334 *,

333 - skip_special_tokens: bool = False,

335 + skip_special_tokens: Optional[bool] = None,

336 ) -> str:

337 return decode_tokens(tokenizer,

338 list(token_ids),

@@ -395,7 +397,9 @@ class _BoundPromptSequence:

397 def token_ids(self) -> list[int]:

398 if self._token_ids is None:

399 assert self._text is not None

398 - self._token_ids = _cached_encode(self.tokenizer, self._text)

400 + self._token_ids = _cached_encode(self.tokenizer,

401 + self._text,

402 + add_special_tokens=False)

403

404 return self._token_ids

405

@@ -1046,7 +1050,7 @@ class BaseMultiModalProcessor(ABC, Generic[_I]):

1050 mm_items: MultiModalDataItems,

1051 hf_processor_mm_kwargs: Mapping[str, object],

1052 out_mm_kwargs: MultiModalKwargs,

1049 - ) -> list[PromptUpdate]:

1053 + ) -> Sequence[PromptUpdate]:

1054 """

1055 Given the original multi-modal items for this modality

1056 and HF-processed data, output the updates to perform.

vllm/transformers_utils/tokenizer.py

+15 -4

@@ -34,13 +34,20 @@ def decode_tokens(

34 tokenizer: AnyTokenizer,

35 token_ids: list[int],

36 *,

37 - skip_special_tokens: bool = False,

37 + skip_special_tokens: Optional[bool] = None,

38 ) -> str:

39 """

40 Backend-agnostic equivalent of HF's

41 - :code:`tokenizer.decode(token_ids, skip_special_tokens=...)`.

41 + :code:`tokenizer.decode(token_ids, ...)`.

42 +

43 + :code:`skip_special_tokens=None` means to use the backend's default

44 + settings.

45 """

43 - return tokenizer.decode(token_ids, skip_special_tokens=skip_special_tokens)

46 + if skip_special_tokens is not None:

47 + return tokenizer.decode(token_ids,

48 + skip_special_tokens=skip_special_tokens)

49 +

50 + return tokenizer.decode(token_ids)

53 def encode_tokens(

@@ -51,10 +58,14 @@ def encode_tokens(

58 ) -> list[int]:

59 """

60 Backend-agnostic equivalent of HF's

54 - :code:`tokenizer.encode(text, add_special_tokens=...)`.

61 + :code:`tokenizer.encode(text, ...)`.

62 +

63 + :code:`add_special_tokens=None` means to use the backend's default

64 + settings.

65 """

66 if add_special_tokens is not None:

67 return tokenizer.encode(text, add_special_tokens=add_special_tokens)

68 +

69 return tokenizer.encode(text)

vllm/utils.py

+1 -1

@@ -845,7 +845,7 @@ def is_list_of(

845 assert_never(check)

846

847

848 -def flatten_2d_lists(lists: list[list[T]]) -> list[T]:

848 +def flatten_2d_lists(lists: Iterable[Iterable[T]]) -> list[T]:

849 """Flatten a list of lists to a single list."""

850 return [item for sublist in lists for item in sublist]

851