@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 5d802522a7

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, 21:00 UTC

[V1][VLM][Pixtral-HF] Support Pixtral-HF on V1 (#14275)

Signed-off-by: Linkun Chen <github@lkchen.net>

lkchen committed Mar 6, 2025 at 00:58 UTC 5d802522a783ba31a150df80499a71acabd8009c

4 files changed +175 -16

docs/source/models/supported_models.md

+1 -5

@@ -866,7 +866,7 @@ See [this page](#generative-models) for more information on how to use generativ

866 - * `PixtralForConditionalGeneration`

867 * Pixtral

868 * T + I<sup>+</sup>

869 - * `mistralai/Pixtral-12B-2409`, `mistral-community/pixtral-12b` (see note), etc.

869 + * `mistralai/Pixtral-12B-2409`, `mistral-community/pixtral-12b`, etc.

870 *

871 * ✅︎

872 * ✅︎

@@ -930,10 +930,6 @@ For more details, please see: <gh-pr:4087#issuecomment-2250397630>

930 Currently the PaliGemma model series is implemented without PrefixLM attention mask. This model series may be deprecated in a future release.

931 :::

932

933 -:::{note}

934 -`mistral-community/pixtral-12b` does not support V1 yet.

935 -:::

936 -

933 :::{note}

934 To use Qwen2.5-VL series models, you have to install Hugging Face Transformers library from source via `pip install git+https://github.com/huggingface/transformers`.

935 :::

vllm/model_executor/models/llava.py

+164 -7

@@ -4,7 +4,7 @@ from abc import abstractmethod

4 from collections.abc import Iterable, Mapping, Sequence

5 from functools import cached_property

6 from typing import (Final, List, Literal, Optional, Protocol, Set, Tuple,

7 - TypedDict, TypeVar, Union)

7 + TypedDict, TypeVar, Union, cast)

9 import torch

10 import torch.nn as nn

@@ -35,6 +35,7 @@ from vllm.multimodal.processing import (BaseMultiModalProcessor,

35 PromptReplacement, PromptUpdate)

36 from vllm.multimodal.profiling import BaseDummyInputsBuilder, ProcessorInputs

37 from vllm.sequence import IntermediateTensors

38 +from vllm.utils import JSONTree, flatten_2d_lists, json_map_leaves

40 from .clip import CLIPVisionModel

41 from .interfaces import SupportsMultiModal, SupportsPP

@@ -56,6 +57,25 @@ class LlavaImagePixelInputs(TypedDict):

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

58 """

60 + feat_is_patch: Union[torch.Tensor, List[torch.Tensor]]

61 + """

62 + A boolean mask indicating which image features correspond

63 + to patch tokens.

64 +

65 + Shape: `(batch_size, num_crops, num_patch)`

66 + """

67 +

68 + embed_is_patch: Union[torch.Tensor, List[torch.Tensor]]

69 + """

70 + A boolean mask indicating which image embeddings correspond

71 + to patch tokens.

72 +

73 + Shape: `(batch_size, num_embeds)`

74 + """

75 +

76 + num_crops: torch.Tensor

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

78 +

80 class LlavaImageEmbeddingInputs(TypedDict):

81 type: Literal["image_embeds"]

@@ -65,6 +85,25 @@ class LlavaImageEmbeddingInputs(TypedDict):

85 `hidden_size` must match the hidden size of language model backbone.

86 """

88 + feat_is_patch: Union[torch.Tensor, List[torch.Tensor]]

89 + """

90 + A boolean mask indicating which image features correspond

91 + to patch tokens.

92 +

93 + Shape: `(batch_size, num_crops, num_patch)`

94 + """

95 +

96 + embed_is_patch: Union[torch.Tensor, List[torch.Tensor]]

97 + """

98 + A boolean mask indicating which image embeddings correspond

99 + to patch tokens.

100 +

101 + Shape: `(batch_size, num_embeds)`

102 + """

103 +

104 + num_crops: torch.Tensor

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

106 +

107

108 LlavaImageInputs = Union[LlavaImagePixelInputs, LlavaImageEmbeddingInputs]

109

@@ -317,6 +356,26 @@ class PixtralHFMultiModalProcessor(

356 for p, (h, w) in zip(pixel_values, image_sizes)

357 ]

358

359 + hf_config = self.info.get_hf_config()

360 +

361 + tile_sizes = [

362 + get_pixtral_hf_image_feature_grid_size(

363 + hf_config.vision_config,

364 + image_width=pixel_value.shape[-1],

365 + image_height=pixel_value.shape[-2])

366 + for pixel_value in processed_outputs["pixel_values"]

367 + ]

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

369 + for ncols, nrows in tile_sizes])

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

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

372 + embed_is_patch = torch.tensor(

373 + flatten_2d_lists([([True] * ncols + [False]) * nrows

374 + for ncols, nrows in tile_sizes]))

375 + processed_outputs["num_crops"] = num_crops

376 + processed_outputs["embed_is_patch"] = embed_is_patch

377 + processed_outputs["feat_is_patch"] = embed_is_patch

378 +

379 return processed_outputs

380

381 def _get_mm_fields_config(

@@ -324,7 +383,13 @@ class PixtralHFMultiModalProcessor(

383 hf_inputs: BatchFeature,

384 hf_processor_mm_kwargs: Mapping[str, object],

385 ) -> Mapping[str, MultiModalFieldConfig]:

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

387 return dict(

388 + feat_is_patch=MultiModalFieldConfig.flat_from_sizes(

389 + "image", num_crops),

390 + embed_is_patch=MultiModalFieldConfig.flat_from_sizes(

391 + "image", num_crops),

392 + num_crops=MultiModalFieldConfig.batched("image"),

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

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

395 )

@@ -562,6 +627,23 @@ class LlavaForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP):

627 if pixel_values is None and image_embeds is None:

628 return None

629

630 + feat_is_patch = kwargs.pop("feat_is_patch", None)

631 + if feat_is_patch is not None and not isinstance(

632 + feat_is_patch, (torch.Tensor, list)):

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

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

635 +

636 + embed_is_patch = kwargs.pop("embed_is_patch", None)

637 + if embed_is_patch is not None and not isinstance(

638 + embed_is_patch, (torch.Tensor, list)):

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

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

641 +

642 + num_crops = kwargs.pop("num_crops", None)

643 + if num_crops is not None and not isinstance(num_crops, torch.Tensor):

644 + raise ValueError("Incorrect type of num_crops. "

645 + f"Got type: {type(num_crops)}")

646 +

647 if pixel_values is not None:

648 if not isinstance(pixel_values, (torch.Tensor, list)):

649 raise ValueError("Incorrect type of pixel values. "

@@ -571,12 +653,18 @@ class LlavaForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP):

653 return LlavaImagePixelInputs(

654 type="pixel_values",

655 data=flatten_bn(pixel_values),

656 + feat_is_patch=feat_is_patch,

657 + embed_is_patch=embed_is_patch,

658 + num_crops=num_crops,

659 )

660

661 return LlavaImagePixelInputs(

662 type="pixel_values",

663 data=self._validate_pixel_values(

664 flatten_bn(pixel_values, concat=True)),

665 + feat_is_patch=feat_is_patch,

666 + embed_is_patch=embed_is_patch,

667 + num_crops=num_crops,

668 )

669

670 if image_embeds is not None:

@@ -587,6 +675,9 @@ class LlavaForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP):

675 return LlavaImageEmbeddingInputs(

676 type="image_embeds",

677 data=flatten_bn(image_embeds, concat=True),

678 + feat_is_patch=feat_is_patch,

679 + embed_is_patch=embed_is_patch,

680 + num_crops=num_crops,

681 )

682

683 raise AssertionError("This line should be unreachable.")

@@ -633,16 +724,74 @@ class LlavaForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP):

724

725 assert self.vision_tower is not None

726 image_features = self._process_image_pixels(image_input)

636 - return self.multi_modal_projector(image_features)

727

638 - def get_multimodal_embeddings(

639 - self, **kwargs

640 - ) -> Union[list[torch.Tensor], torch.Tensor, tuple[torch.Tensor, ...]]:

728 + if isinstance(image_features, torch.Tensor):

729 + return self.multi_modal_projector(image_features)

730 +

731 + feature_sizes = [

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

733 + ]

734 +

735 + image_embeds = self.multi_modal_projector(torch.cat(image_features))

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

737 + return image_embeds

738 +

739 + def _get_mm_embeds(

740 + self,

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

742 + feat_is_patch: torch.Tensor, # Shape: (num_crop, num_patch)

743 + num_crops: torch.Tensor, # Shape: (num_images,)

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

745 + ) -> list[torch.Tensor]:

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

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

748 +

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

750 + """

751 +

752 + # Insert columns of nan values according to `feat_is_patch`. This work

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

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

755 + # put the logic here.

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

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

758 + feat_is_patch = feat_is_patch.view(-1)

759 + embed_is_patch = embed_is_patch.view(-1)

760 + expanded_embedding = torch.full(

761 + (sum(num_crops), *features.shape[1:]),

762 + torch.nan,

763 + dtype=features.dtype).to(features.device)

764 + expanded_embedding[feat_is_patch] = features

765 +

766 + num_crops_per_image = num_crops.tolist()

767 + feats_per_image = expanded_embedding.split(num_crops_per_image)

768 + f_is_patch_per_image = feat_is_patch.split(num_crops_per_image)

769 +

770 + embed_dim = expanded_embedding.shape[-1]

771 + num_embeds = embed_is_patch.shape[0]

772 +

773 + embeds_in_batch = list[torch.Tensor]()

774 + for feats, f_is_patch in zip(feats_per_image, f_is_patch_per_image):

775 + embeds = feats.new_full((num_embeds, embed_dim), torch.nan)

776 + embeds[embed_is_patch] = feats[f_is_patch]

777 + embeds_in_batch.append(embeds)

778 +

779 + return embeds_in_batch

780 +

781 + def get_multimodal_embeddings(self, **kwargs) -> Optional[NestedTensors]:

782 image_input = self._parse_and_validate_image_input(**kwargs)

783 if image_input is None:

784 return None

785 vision_embeddings = self._process_image_input(image_input)

645 - return vision_embeddings

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

787 + return vision_embeddings

788 + else:

789 + nested_emb = [

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

791 + vision_embeddings, image_input["feat_is_patch"],

792 + image_input["num_crops"], image_input["embed_is_patch"])

793 + ]

794 + return flatten_2d_lists(nested_emb)

795

796 def get_input_embeddings(

797 self,

@@ -651,8 +800,15 @@ class LlavaForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP):

800 ) -> torch.Tensor:

801 inputs_embeds = self.language_model.get_input_embeddings(input_ids)

802 if multimodal_embeddings is not None:

803 + # Extract the patch tokens

804 + patch_embeddings = json_map_leaves(

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

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

807 + )

808 +

809 inputs_embeds = merge_multimodal_embeddings(

655 - input_ids, inputs_embeds, multimodal_embeddings,

810 + input_ids, inputs_embeds, cast(NestedTensors,

811 + patch_embeddings),

812 self.config.image_token_index)

813 return inputs_embeds

814

@@ -705,6 +861,7 @@ class LlavaForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP):

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

862 # condition is for v0 compatibility.

863 elif inputs_embeds is None:

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

865 vision_embeddings = self.get_multimodal_embeddings(**kwargs)

866 inputs_embeds = self.get_input_embeddings(input_ids,

867 vision_embeddings)

vllm/model_executor/models/molmo.py

+2 -2

@@ -1484,8 +1484,8 @@ class MolmoForCausalLM(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA,

1484

1485 img_patch_id = kwargs.pop("img_patch_id", None)

1486 if not isinstance(img_patch_id, torch.Tensor):

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

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

1487 + raise ValueError("Incorrect type of img_patch_id. "

1488 + f"Got type: {type(img_patch_id)}")

1489 self.img_patch_id = img_patch_id.flatten().unique().item()

1490

1491 return MolmoImageInputs(

vllm/model_executor/models/pixtral.py

+8 -2

@@ -1042,9 +1042,13 @@ class PixtralHFVisionModel(nn.Module):

1042 for img in pixel_values

1043 ]

1044

1045 + patch_embeds = [

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

1047 + ]

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

1049 +

1050 # flatten to a single sequence

1046 - patch_embeds = torch.cat(

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

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

1052 patch_embeds = self.ln_pre(patch_embeds)

1053

1054 # positional embeddings

@@ -1075,6 +1079,8 @@ class PixtralHFVisionModel(nn.Module):

1079 out = resolve_visual_encoder_outputs(out, feature_sample_layers, None,

1080 self.config.num_hidden_layers)

1081

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

1083 + out = torch.split(torch.squeeze(out), embed_sizes)

1084 return out

1085

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