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

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

[Model] Move `vision_feature_select_strategy` into `resolve_visual_encoder_outputs` (#25938)

Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>

Cyrus Leung committed Sep 30, 2025 at 19:24 UTC d7e34b4210fd7ce7d8f49441c50f2b99c93913ae

12 files changed +155 -179

tests/models/test_vision.py

+6 -5

@@ -18,7 +18,7 @@ from vllm.utils import get_open_port, update_environment_variables

20 @pytest.mark.parametrize(

21 - ("feature_sample_layers", "num_layers_loaded", "max_possible_layers",

21 + ("select_layers", "num_layers_loaded", "max_possible_layers",

22 "expected_features"),

23 [

24 # All layers loaded

@@ -28,8 +28,8 @@ from vllm.utils import get_open_port, update_environment_variables

28 ([1, 10], 10, 20, [1, 10]),

29 ([-20, -11], 10, 20, [1, 10]),

30 ])

31 -def test_resolve_visual_encoder_outputs(feature_sample_layers,

32 - num_layers_loaded, max_possible_layers,

31 +def test_resolve_visual_encoder_outputs(select_layers, num_layers_loaded,

32 + max_possible_layers,

33 expected_features):

34 """

35 Test that offsets are correctly handled for vision feature layers.

@@ -39,9 +39,10 @@ def test_resolve_visual_encoder_outputs(feature_sample_layers,

39 ]

40 output_tensor = resolve_visual_encoder_outputs(

41 encoder_outputs=encoder_outputs,

42 - feature_sample_layers=feature_sample_layers,

42 post_layer_norm=None,

44 - max_possible_layers=max_possible_layers)

43 + select_layers=select_layers,

44 + max_possible_layers=max_possible_layers,

45 + )

46 assert torch.equal(torch.tensor(expected_features), output_tensor)

vllm/model_executor/models/aya_vision.py

+4 -23

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

27 PromptUpdateDetails)

28 from vllm.multimodal.profiling import BaseDummyInputsBuilder

29 from vllm.sequence import IntermediateTensors

30 -from vllm.utils.jsontree import json_map_leaves

30 from vllm.utils.tensor_schema import TensorSchema, TensorShape

32 from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP

@@ -350,29 +349,11 @@ class AyaVisionForConditionalGeneration(nn.Module, SupportsMultiModal,

349 self,

350 vision_tower: SiglipVisionModel,

351 pixel_values: torch.Tensor,

353 - **kwargs,

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

355 - target_dtype: torch.dtype = \

356 - vision_tower.get_input_embeddings().weight.dtype

357 - image_features: Union[torch.Tensor, tuple[torch.Tensor, ...]] = \

358 - vision_tower(pixel_values.to(dtype=target_dtype), **kwargs)

359 -

360 - def select_features(leaf: torch.Tensor):

361 - return self._select_image_features(

362 - leaf,

363 - strategy=self.config.vision_feature_select_strategy,

364 - )

365 -

366 - return json_map_leaves(select_features, image_features)

367 -

368 - def _select_image_features(self, image_features: torch.Tensor, *,

369 - strategy: str) -> torch.Tensor:

370 - if strategy == "default":

371 - return image_features[:, 1:]

372 - elif strategy == "full":

373 - return image_features

374 -

375 - raise ValueError(f"Unexpected select feature strategy: {strategy}")

353 + return vision_tower(

354 + pixel_values.to(dtype=vision_tower.dtype),

355 + feature_select_strategy=self.config.vision_feature_select_strategy,

356 + )

357

358 def _process_image_input(self, image_input: AyaVisionImagePixelInputs,

359 **kwargs) -> list[torch.Tensor]:

vllm/model_executor/models/clip.py

+21 -10

@@ -19,7 +19,8 @@ from vllm.model_executor.layers.quantization import QuantizationConfig

19 from vllm.model_executor.model_loader.weight_utils import default_weight_loader

20 from vllm.model_executor.models.interfaces import SupportsQuant

22 -from .vision import VisionEncoderInfo, resolve_visual_encoder_outputs

22 +from .vision import (VisionEncoderInfo, VisionFeatureSelectStrategy,

23 + resolve_visual_encoder_outputs)

26 class CLIPEncoderInfo(VisionEncoderInfo[CLIPVisionConfig]):

@@ -308,24 +309,29 @@ class CLIPVisionTransformer(nn.Module):

309 def forward(

310 self,

311 pixel_values: torch.Tensor,

311 - feature_sample_layers: Optional[list[int]] = None,

312 + *,

313 + select_layers: Optional[list[int]] = None,

314 + feature_select_strategy: Optional[VisionFeatureSelectStrategy] = None,

315 ) -> torch.Tensor:

316

317 hidden_states = self.embeddings(pixel_values)

318 hidden_states = self.pre_layrnorm(hidden_states)

319

317 - return_all_hidden_states = feature_sample_layers is not None

318 -

320 # Produces either the last layer output or all of the hidden states,

320 - # depending on if we have feature_sample_layers or not

321 + # depending on if we have select_layers or not

322 encoder_outputs = self.encoder(

323 inputs_embeds=hidden_states,

323 - return_all_hidden_states=return_all_hidden_states)

324 + return_all_hidden_states=select_layers is not None,

325 + )

326

327 # Handle post-norm (if applicable) and stacks feature layers if needed

328 encoder_outputs = resolve_visual_encoder_outputs(

327 - encoder_outputs, feature_sample_layers, self.post_layernorm,

328 - self.config.num_hidden_layers)

329 + encoder_outputs,

330 + self.post_layernorm,

331 + select_layers=select_layers,

332 + max_possible_layers=self.config.num_hidden_layers,

333 + feature_select_strategy=feature_select_strategy,

334 + )

335

336 return encoder_outputs

337

@@ -355,9 +361,14 @@ class CLIPVisionModel(nn.Module, SupportsQuant):

361 def forward(

362 self,

363 pixel_values: torch.Tensor,

358 - feature_sample_layers: Optional[list[int]] = None,

364 + select_layers: Optional[list[int]] = None,

365 + feature_select_strategy: Optional[VisionFeatureSelectStrategy] = None,

366 ) -> torch.Tensor:

360 - return self.vision_model(pixel_values, feature_sample_layers)

367 + return self.vision_model(

368 + pixel_values,

369 + select_layers=select_layers,

370 + feature_select_strategy=feature_select_strategy,

371 + )

372

373 @property

374 def device(self):

vllm/model_executor/models/llava.py

+4 -21

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

33 PromptUpdateDetails)

34 from vllm.multimodal.profiling import BaseDummyInputsBuilder

35 from vllm.sequence import IntermediateTensors

36 -from vllm.utils.jsontree import json_map_leaves

36 from vllm.utils.tensor_schema import TensorSchema, TensorShape

38 from .clip import CLIPVisionModel

@@ -604,16 +603,6 @@ class LlavaForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP):

603

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

605

607 - def _select_image_features(self, image_features: torch.Tensor, *,

608 - strategy: str) -> torch.Tensor:

609 - # Copied from https://github.com/huggingface/transformers/blob/39c3c0a72af6fbda5614dde02ff236069bb79827/src/transformers/models/llava/modeling_llava.py#L421 # noqa

610 - if strategy == "default":

611 - return image_features[:, 1:]

612 - elif strategy == "full":

613 - return image_features

614 -

615 - raise ValueError(f"Unexpected select feature strategy: {strategy}")

616 -

606 def _image_pixels_to_features(

607 self,

608 vision_tower: Union[CLIPVisionModel, SiglipVisionModel,

@@ -622,16 +611,10 @@ class LlavaForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP):

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

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

613 # this is already done inside the vision tower

625 - image_features: Union[torch.Tensor, tuple[torch.Tensor, ...]] = \

626 - vision_tower(pixel_values)

627 -

628 - def select_features(leaf: torch.Tensor):

629 - return self._select_image_features(

630 - leaf,

631 - strategy=self.config.vision_feature_select_strategy,

632 - )

633 -

634 - return json_map_leaves(select_features, image_features)

614 + return vision_tower(

615 + pixel_values,

616 + feature_select_strategy=self.config.vision_feature_select_strategy,

617 + )

618

619 def _process_image_pixels(

620 self,

vllm/model_executor/models/llava_next.py

+6 -19

@@ -235,12 +235,12 @@ class LlavaNextForConditionalGeneration(nn.Module, SupportsMultiModal,

235 # Determine the layer up to which we will initialize the vision tower

236 if isinstance(vision_feature_layer, int):

237 vision_hidden_size = config.vision_config.hidden_size

238 - self.feature_sample_layers = None

238 + self.select_layers = None

239 # Used for multimodal granite models to control encoder outputs

240 elif isinstance(vision_feature_layer, (list, tuple)):

241 vision_hidden_size = config.vision_config.hidden_size * len(

242 vision_feature_layer)

243 - self.feature_sample_layers = vision_feature_layer

243 + self.select_layers = vision_feature_layer

244 else:

245 raise TypeError(

246 f"vision_layer_feature type: {type(vision_feature_layer)}"

@@ -312,30 +312,17 @@ class LlavaNextForConditionalGeneration(nn.Module, SupportsMultiModal,

312

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

314

315 - def _select_image_features(self, image_features: torch.Tensor, *,

316 - strategy: str) -> torch.Tensor:

317 - # Copied from https://github.com/huggingface/transformers/blob/39c3c0a72af6fbda5614dde02ff236069bb79827/src/transformers/models/llava/modeling_llava.py#L421 # noqa

318 - if strategy == "default":

319 - return image_features[:, 1:]

320 - elif strategy == "full":

321 - return image_features

322 -

323 - raise ValueError(f"Unexpected select feature strategy: {strategy}")

324 -

315 def _image_pixels_to_features(

316 self,

317 vision_tower: Union[CLIPVisionModel, SiglipVisionModel],

318 pixel_values: torch.Tensor,

319 ) -> torch.Tensor:

330 -

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

321 # this is already done inside the vision tower

333 - image_features = vision_tower(

334 - pixel_values, feature_sample_layers=self.feature_sample_layers)

335 -

336 - return self._select_image_features(

337 - image_features,

338 - strategy=self.config.vision_feature_select_strategy,

322 + return vision_tower(

323 + pixel_values,

324 + select_layers=self.select_layers,

325 + feature_select_strategy=self.config.vision_feature_select_strategy,

326 )

327

328 # Based on: https://github.com/haotian-liu/LLaVA/blob/main/llava/model/llava_arch.py

vllm/model_executor/models/llava_next_video.py

+3 -14

@@ -349,27 +349,16 @@ class LlavaNextVideoForConditionalGeneration(nn.Module, SupportsMultiModal,

349 "w": expected_w,

350 })

351

352 - def _select_image_features(self, image_features: torch.Tensor, *,

353 - strategy: str) -> torch.Tensor:

354 - if strategy == "default":

355 - return image_features[:, 1:]

356 - elif strategy == "full":

357 - return image_features

358 -

359 - raise ValueError(f"Unexpected select feature strategy: {strategy}")

360 -

352 def _video_pixels_to_features(

353 self,

354 vision_tower: Union[CLIPVisionModel, SiglipVisionModel],

355 pixel_values: torch.Tensor,

356 ) -> torch.Tensor:

366 -

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

358 # this is already done inside the vision tower

369 - image_features = vision_tower(pixel_values)

370 - image_features = self._select_image_features(

371 - image_features,

372 - strategy=self.config.vision_feature_select_strategy,

359 + image_features = vision_tower(

360 + pixel_values,

361 + feature_select_strategy=self.config.vision_feature_select_strategy,

362 )

363 image_features = self.vision_resampler(image_features)

364 image_features = self.multi_modal_projector(image_features)

vllm/model_executor/models/llava_onevision.py

+6 -19

@@ -577,27 +577,16 @@ class LlavaOnevisionForConditionalGeneration(nn.Module, SupportsMultiModal,

577

578 return mm_input_by_modality

579

580 - def _select_image_features(self, image_features: torch.Tensor, *,

581 - strategy: str) -> torch.Tensor:

582 - if strategy == "default":

583 - return image_features[:, 1:]

584 - elif strategy == "full":

585 - return image_features

586 -

587 - raise ValueError(f"Unexpected select feature strategy: {strategy}")

588 -

580 def _image_pixels_to_features(

581 self,

582 vision_tower: Union[CLIPVisionModel, SiglipVisionModel],

583 pixel_values: torch.Tensor,

584 ) -> torch.Tensor:

594 -

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

586 # this is already done inside the vision tower

597 - image_features = vision_tower(pixel_values)

598 - return self._select_image_features(

599 - image_features,

600 - strategy=self.config.vision_feature_select_strategy,

587 + return vision_tower(

588 + pixel_values,

589 + feature_select_strategy=self.config.vision_feature_select_strategy,

590 )

591

592 # Based on: https://github.com/haotian-liu/LLaVA/blob/main/llava/model/llava_arch.py

@@ -750,13 +739,11 @@ class LlavaOnevisionForConditionalGeneration(nn.Module, SupportsMultiModal,

739 vision_tower: Union[CLIPVisionModel, SiglipVisionModel],

740 pixel_values: torch.Tensor,

741 ) -> torch.Tensor:

753 -

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

743 # this is already done inside the vision tower

756 - video_features = vision_tower(pixel_values)

757 - video_features = self._select_image_features(

758 - video_features,

759 - strategy=self.config.vision_feature_select_strategy,

744 + video_features = vision_tower(

745 + pixel_values,

746 + feature_select_strategy=self.config.vision_feature_select_strategy,

747 )

748 video_features = self.multi_modal_projector(video_features)

749 video_features = self.apply_pooling(video_features)

vllm/model_executor/models/minimax_vl_01.py

+4 -20

@@ -17,7 +17,6 @@ from vllm.model_executor.layers.quantization import QuantizationConfig

17 from vllm.multimodal import MULTIMODAL_REGISTRY

18 from vllm.multimodal.inputs import MultiModalFieldConfig

19 from vllm.sequence import IntermediateTensors

20 -from vllm.utils.jsontree import json_map_leaves

20 from vllm.utils.tensor_schema import TensorSchema, TensorShape

22 from .clip import CLIPVisionModel

@@ -221,15 +220,6 @@ class MiniMaxVL01ForConditionalGeneration(nn.Module, SupportsMultiModal,

220 def get_language_model(self) -> torch.nn.Module:

221 return self.language_model

222

224 - def _select_image_features(self, image_features: torch.Tensor, *,

225 - strategy: str) -> torch.Tensor:

226 - if strategy == "default":

227 - return image_features[:, 1:]

228 - elif strategy == "full":

229 - return image_features

230 -

231 - raise ValueError(f"Unexpected select feature strategy: {strategy}")

232 -

223 def _image_pixels_to_features(

224 self,

225 vision_tower: Union[CLIPVisionModel, SiglipVisionModel,

@@ -238,16 +228,10 @@ class MiniMaxVL01ForConditionalGeneration(nn.Module, SupportsMultiModal,

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

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

230 # this is already done inside the vision tower

241 - image_features: tuple[torch.Tensor, ...] = \

242 - tuple(vision_tower(p) for p in pixel_values)

243 -

244 - def select_features(leaf: torch.Tensor):

245 - return self._select_image_features(

246 - leaf,

247 - strategy=self.config.vision_feature_select_strategy,

248 - )

249 -

250 - return json_map_leaves(select_features, image_features)

231 + feature_select_strategy = self.config.vision_feature_select_strategy

232 + return tuple(

233 + vision_tower(p, feature_select_strategy=feature_select_strategy)

234 + for p in pixel_values)

235

236 # adapted from https://huggingface.co/MiniMaxAI/MiniMax-VL-01/blob/main/modeling_minimax_vl_01.py#L616-L631

237 def pack_image_features(self, image_features: list[torch.Tensor],

vllm/model_executor/models/pixtral.py

+15 -7

@@ -51,7 +51,8 @@ from vllm.utils.tensor_schema import TensorSchema, TensorShape

52 from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP

53 from .utils import flatten_bn, init_vllm_registered_model, maybe_prefix

54 -from .vision import VisionEncoderInfo, resolve_visual_encoder_outputs

54 +from .vision import (VisionEncoderInfo, VisionFeatureSelectStrategy,

55 + resolve_visual_encoder_outputs)

57 try:

58 from xformers import ops as xops

@@ -1218,7 +1219,9 @@ class PixtralHFVisionModel(nn.Module):

1219 def forward(

1220 self,

1221 pixel_values: list[torch.Tensor],

1221 - feature_sample_layers: Optional[list[int]] = None,

1222 + *,

1223 + select_layers: Optional[list[int]] = None,

1224 + feature_select_strategy: Optional[VisionFeatureSelectStrategy] = None,

1225 ) -> tuple[torch.Tensor, ...]:

1226 """

1227 Args:

@@ -1226,7 +1229,7 @@ class PixtralHFVisionModel(nn.Module):

1229 in pixel_values. This means it will be a list of tensors

1230 because multiple requests batched can have multiple images,

1231 each with their own shape potentially

1229 - feature_sample_layers: Layer indices whose features should be

1232 + select_layers: Layer indices whose features should be

1233 concatenated and used as the visual encoder output. If none

1234 are provided, the last layer is used.

1235

@@ -1267,15 +1270,20 @@ class PixtralHFVisionModel(nn.Module):

1270 [p.shape[-2] * p.shape[-1] for p in patch_embeds_list],

1271 patch_embeds)

1272

1270 - return_all_hidden_states = feature_sample_layers is not None

1273 out = self.transformer(

1274 patch_embeds,

1275 attention_mask,

1276 position_embedding,

1275 - return_all_hidden_states=return_all_hidden_states)

1277 + return_all_hidden_states=select_layers is not None,

1278 + )

1279

1277 - out = resolve_visual_encoder_outputs(out, feature_sample_layers, None,

1278 - self.config.num_hidden_layers)

1280 + out = resolve_visual_encoder_outputs(

1281 + out,

1282 + None,

1283 + select_layers=select_layers,

1284 + max_possible_layers=self.config.num_hidden_layers,

1285 + feature_select_strategy=feature_select_strategy,

1286 + )

1287

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

1289 return torch.split(out.squeeze(0), embed_sizes)

vllm/model_executor/models/siglip.py

+22 -12

@@ -23,7 +23,8 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (

23 from vllm.model_executor.model_loader.weight_utils import (

24 default_weight_loader, maybe_remap_kv_scale_name)

26 -from .vision import VisionEncoderInfo, resolve_visual_encoder_outputs

26 +from .vision import (VisionEncoderInfo, VisionFeatureSelectStrategy,

27 + resolve_visual_encoder_outputs)

30 class SiglipEncoderInfo(VisionEncoderInfo[SiglipVisionConfig]):

@@ -415,28 +416,31 @@ class SiglipVisionTransformer(nn.Module):

416 def forward(

417 self,

418 pixel_values: torch.Tensor,

418 - interpolate_pos_encoding: bool = True,

419 - feature_sample_layers: Optional[list[int]] = None,

419 + *,

420 + interpolate_pos_encoding: bool = False,

421 + select_layers: Optional[list[int]] = None,

422 + feature_select_strategy: Optional[VisionFeatureSelectStrategy] = None,

423 ) -> torch.Tensor:

421 -

424 hidden_states = self.embeddings(

425 pixel_values,

426 interpolate_pos_encoding=interpolate_pos_encoding,

427 )

428

427 - return_all_hidden_states = feature_sample_layers is not None

428 -

429 # Produces either the last layer output or all of the hidden states,

430 - # depending on if we have feature_sample_layers or not

430 + # depending on if we have select_layers or not

431 encoder_outputs = self.encoder(

432 inputs_embeds=hidden_states,

433 - return_all_hidden_states=return_all_hidden_states,

433 + return_all_hidden_states=select_layers is not None,

434 )

435

436 # Handle post-norm (if applicable) and stacks feature layers if needed

437 encoder_outputs = resolve_visual_encoder_outputs(

438 - encoder_outputs, feature_sample_layers, self.post_layernorm,

439 - self.config.num_hidden_layers)

438 + encoder_outputs,

439 + self.post_layernorm,

440 + select_layers=select_layers,

441 + max_possible_layers=self.config.num_hidden_layers,

442 + feature_select_strategy=feature_select_strategy,

443 + )

444

445 # TODO: add this back when pooled_output is used in inference.

446 # if self.use_head:

@@ -471,16 +475,22 @@ class SiglipVisionModel(nn.Module):

475 def get_input_embeddings(self) -> nn.Module:

476 return self.vision_model.embeddings.patch_embedding

477

478 + @property

479 + def dtype(self):

480 + return self.get_input_embeddings().weight.dtype

481 +

482 def forward(

483 self,

484 pixel_values: torch.Tensor,

485 interpolate_pos_encoding: bool = False,

478 - feature_sample_layers: Optional[list[int]] = None,

486 + select_layers: Optional[list[int]] = None,

487 + feature_select_strategy: Optional[VisionFeatureSelectStrategy] = None,

488 ) -> torch.Tensor:

489 return self.vision_model(

490 pixel_values=pixel_values,

491 interpolate_pos_encoding=interpolate_pos_encoding,

483 - feature_sample_layers=feature_sample_layers,

492 + select_layers=select_layers,

493 + feature_select_strategy=feature_select_strategy,

494 )

495

496 def load_weights(self, weights: Iterable[tuple[str,

vllm/model_executor/models/tarsier.py

+4 -19

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

33 PromptReplacement, PromptUpdate)

34 from vllm.multimodal.profiling import BaseDummyInputsBuilder

35 from vllm.sequence import IntermediateTensors

36 -from vllm.utils.jsontree import json_map_leaves

36 from vllm.utils.tensor_schema import TensorSchema, TensorShape

38 from .clip import CLIPVisionModel

@@ -476,30 +475,16 @@ class TarsierForConditionalGeneration(nn.Module, SupportsMultiModal,

475

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

477

479 - def _select_image_features(self, image_features: torch.Tensor, *,

480 - strategy: str) -> torch.Tensor:

481 - if strategy == "default":

482 - return image_features[:, 1:]

483 - elif strategy == "full":

484 - return image_features

485 - raise ValueError(f"Unexpected select feature strategy: {strategy}")

486 -

478 def _image_pixels_to_features(

479 self,

480 vision_tower: Union[CLIPVisionModel, SiglipVisionModel],

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

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

483 # From vLLM LLaVA, vision tower output handling

493 - image_hidden_states: Union[torch.Tensor, tuple[torch.Tensor, ...]] = \

494 - vision_tower(pixel_values)

495 -

496 - def select_features_fn(leaf: torch.Tensor):

497 - return self._select_image_features(

498 - leaf,

499 - strategy=self.config.vision_feature_select_strategy,

500 - )

501 -

502 - return json_map_leaves(select_features_fn, image_hidden_states)

484 + return vision_tower(

485 + pixel_values,

486 + feature_select_strategy=self.config.vision_feature_select_strategy,

487 + )

488

489 def _add_tarsier_split_tokens(

490 self, projected_image_features: torch.Tensor) -> torch.Tensor:

vllm/model_executor/models/vision.py

+60 -10

@@ -4,10 +4,12 @@

4 import itertools

5 import math

6 from abc import ABC, abstractmethod

7 -from typing import Final, Generic, Literal, Optional, Protocol, TypeVar, Union

7 +from typing import (Callable, Final, Generic, Literal, Optional, Protocol,

8 + TypeVar, Union)

10 import torch

11 from transformers import PretrainedConfig

12 +from typing_extensions import assert_never

14 from vllm.distributed import (get_tensor_model_parallel_rank,

15 get_tensor_model_parallel_world_size,

@@ -86,11 +88,39 @@ def get_vit_attn_backend(head_size: int, dtype: torch.dtype) -> _Backend:

88 return current_platform.get_vit_attn_backend(head_size, dtype)

91 +VisionFeatureSelectStrategy = Union[

92 + Literal["class", "default", "full"],

93 + Callable[[torch.Tensor], torch.Tensor],

94 +]

95 +

96 +

97 +def _get_vision_feature_selector(

98 + strategy: VisionFeatureSelectStrategy,

99 +) -> Callable[[torch.Tensor], torch.Tensor]:

100 + if callable(strategy):

101 + return strategy

102 +

103 + # https://github.com/huggingface/transformers/blob/cd74917ffc3e8f84e4a886052c5ab32b7ac623cc/src/transformers/models/clip/modeling_clip.py#L762

104 + if strategy == "class":

105 + return lambda feats: feats[:, 0, :]

106 +

107 + # https://github.com/huggingface/transformers/blob/4a02bc7004285bdb12cc033e87ad2578ce2fa900/src/transformers/models/llava/modeling_llava.py#L196

108 + if strategy == "default":

109 + return lambda feats: feats[:, 1:, :]

110 +

111 + if strategy == "full":

112 + return lambda feats: feats

113 +

114 + assert_never(strategy)

115 +

116 +

117 def resolve_visual_encoder_outputs(

118 encoder_outputs: Union[torch.Tensor, list[torch.Tensor]],

91 - feature_sample_layers: Optional[list[int]],

119 post_layer_norm: Optional[torch.nn.LayerNorm],

93 - max_possible_layers: int,

120 + *,

121 + select_layers: Optional[list[int]] = None,

122 + max_possible_layers: Optional[int] = None,

123 + feature_select_strategy: Optional[VisionFeatureSelectStrategy] = None,

124 ) -> torch.Tensor:

125 """Given the outputs a visual encoder module that may correspond to the

126 output of the last layer, or a list of hidden states to be stacked,

@@ -98,17 +128,32 @@ def resolve_visual_encoder_outputs(

128

129 Args:

130 encoder_outputs: Output of encoder's last layer or all hidden states.

101 - feature_sample_layers: Optional layer indices to grab from the encoder

102 - outputs; if provided, encoder outputs must be a list.

131 post_layer_norm: Post norm to apply to the output of the encoder.

132 + select_layers: Optional layer indices to grab from the encoder

133 + outputs; if provided, encoder outputs must be a list.

134 max_possible_layers: Total layers in the fully loaded visual encoder.

105 -

135 + feature_select_strategy: Defines how to select the hidden states

136 + from each layer.

137 """

107 - if feature_sample_layers is None:

138 + if select_layers is None:

139 + if not isinstance(encoder_outputs, torch.Tensor):

140 + raise ValueError("Expected only a single encoder output when "

141 + "`select_layers` is not provided")

142 +

143 + if feature_select_strategy is not None:

144 + select_features = _get_vision_feature_selector(

145 + feature_select_strategy)

146 + encoder_outputs = select_features(encoder_outputs)

147 +

148 if post_layer_norm is not None:

149 return post_layer_norm(encoder_outputs)

150 +

151 return encoder_outputs

152

153 + if max_possible_layers is None:

154 + raise ValueError("`max_possible_layers` must be provided "

155 + "alongside `select_layers`")

156 +

157 # Get the hidden states corresponding to the layer indices.

158 # Negative values are relative to the full visual encoder,

159 # so offset them depending on how many layers were loaded.

@@ -120,13 +165,18 @@ def resolve_visual_encoder_outputs(

165 hs_pool = [

166 encoder_outputs[layer_idx]

167 if layer_idx >= 0 else encoder_outputs[layer_idx + offset]

123 - for layer_idx in feature_sample_layers

168 + for layer_idx in select_layers

169 ]

170

171 + if feature_select_strategy is not None:

172 + select_features = _get_vision_feature_selector(feature_select_strategy)

173 + hs_pool = [select_features(hs) for hs in hs_pool]

174 +

175 # Apply post-norm on the final hidden state if we are using it

127 - uses_last_layer = feature_sample_layers[-1] in (len(hs_pool) - 1, -1)

176 + uses_last_layer = select_layers[-1] in (max_possible_layers - 1, -1)

177 if post_layer_norm is not None and uses_last_layer:

129 - hs_pool[-1] = post_layer_norm(encoder_outputs)

178 + hs_pool[-1] = post_layer_norm(hs_pool[-1])

179 +

180 return torch.cat(hs_pool, dim=-1)

181

182