[Model] Enable LoRA support for tower and connector in LlavaNextForConditionalGeneration (#49788)
Signed-off-by: gangula-karthik <gkarthik923@gmail.com>
karthik committed
Aug 18, 2026 at 23:53 UTC
241ff8c443175ae3e5e2b42b866e4c9275bd7744
2 files changed
+53
-5
docs/models/supported_models.md
+1
-1
@@ -571,7 +571,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
571
| `Llama4ForConditionalGeneration` | Llama 4 | T + I<sup>+</sup> | `meta-llama/Llama-4-Scout-17B-16E-Instruct`, `meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8`, `meta-llama/Llama-4-Maverick-17B-128E-Instruct`, etc. | ✅︎ | ✅︎ |
572
| `Llama_Nemotron_Nano_VL` | Llama Nemotron Nano VL | T + I<sup>E+</sup> | `nvidia/Llama-3.1-Nemotron-Nano-VL-8B-V1` | ✅︎ | ✅︎ |
573
| `LlavaForConditionalGeneration` | LLaVA-1.5, Pixtral (HF Transformers) | T + I<sup>E+</sup> | `llava-hf/llava-1.5-7b-hf`, `mistral-community/pixtral-12b`, etc. | ✅︎ | ✅︎ |
574
-| `LlavaNextForConditionalGeneration` | LLaVA-NeXT, Granite Vision | T + I<sup>E+</sup> | `llava-hf/llava-v1.6-mistral-7b-hf`, `llava-hf/llava-v1.6-vicuna-7b-hf`, `ibm-granite/granite-vision-3.3-2b`, etc. | | ✅︎ |
574
+| `LlavaNextForConditionalGeneration` | LLaVA-NeXT, Granite Vision | T + I<sup>E+</sup> | `llava-hf/llava-v1.6-mistral-7b-hf`, `llava-hf/llava-v1.6-vicuna-7b-hf`, `ibm-granite/granite-vision-3.3-2b`, etc. | ✅︎ | ✅︎ |
575
| `LlavaNextVideoForConditionalGeneration` | LLaVA-NeXT-Video | T + V | `llava-hf/LLaVA-NeXT-Video-7B-hf`, etc. | ✅︎ | ✅︎ |
576
| `LlavaOnevision2ForConditionalGeneration` | LLaVA-OneVision-2 | T + I<sup>+</sup> + V<sup>+</sup> | `lmms-lab-encoder/LLaVA-OneVision-2-8B-Instruct` | | |
577
| `LlavaOnevisionForConditionalGeneration` | LLaVA-Onevision | T + I<sup>+</sup> + V<sup>+</sup> | `llava-hf/llava-onevision-qwen2-7b-ov-hf`, `llava-hf/llava-onevision-qwen2-0.5b-ov-hf`, etc. | | ✅︎ |
vllm/model_executor/models/llava_next.py
+52
-4
@@ -15,13 +15,18 @@ from transformers.models.llava_next.modeling_llava_next import (
15
16
from vllm.config import VllmConfig
17
from vllm.multimodal import MULTIMODAL_REGISTRY
18
-from vllm.multimodal.inputs import MultiModalFieldConfig
18
+from vllm.multimodal.inputs import MultiModalFieldConfig, MultiModalKwargsItem
19
from vllm.multimodal.parse import ImageSize
20
from vllm.sequence import IntermediateTensors
21
from vllm.utils.tensor_schema import TensorSchema, TensorShape
22
23
from .clip import CLIPVisionModel
24
-from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP
24
+from .interfaces import (
25
+ MultiModalEmbeddings,
26
+ SupportsLoRA,
27
+ SupportsMultiModal,
28
+ SupportsPP,
29
+)
30
from .llava import (
31
BaseLlavaMultiModalProcessor,
32
BaseLlavaProcessingInfo,
@@ -30,6 +35,7 @@ from .llava import (
35
LlavaMultiModalProjector,
36
init_vision_tower_for_llava,
37
)
38
+from .module_mapping import MultiModelKeys
39
from .siglip import SiglipVisionModel
40
from .utils import (
41
AutoWeightsLoader,
@@ -37,7 +43,7 @@ from .utils import (
43
init_vllm_registered_model,
44
maybe_prefix,
45
)
40
-from .vision import get_num_selected_vision_tokens
46
+from .vision import get_num_selected_vision_tokens, get_vision_encoder_info
47
48
49
class LlavaNextImagePixelInputs(TensorSchema):
@@ -222,7 +228,14 @@ class LlavaNextMultiModalProcessor(BaseLlavaNextMultiModalProcessor[_I]):
228
info=LlavaNextProcessingInfo,
229
dummy_inputs=LlavaDummyInputsBuilder,
230
)
225
-class LlavaNextForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP):
231
+class LlavaNextForConditionalGeneration(
232
+ nn.Module, SupportsLoRA, SupportsMultiModal, SupportsPP
233
+):
234
+ packed_modules_mapping = {
235
+ "qkv_proj": ["q_proj", "k_proj", "v_proj"],
236
+ "gate_up_proj": ["gate_proj", "up_proj"],
237
+ }
238
+
239
hf_to_vllm_mapper = WeightsMapper(
240
orig_to_new_prefix={
241
# mapping for new names in checkpoint saved after transformers v4.52
@@ -582,3 +595,38 @@ model_executor.models.llava_next.LlavaNextProcessingInfo.get_num_image_tokens].
595
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
596
loader = AutoWeightsLoader(self)
597
return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
598
+
599
+ def get_mm_mapping(self) -> MultiModelKeys:
600
+ """
601
+ Get the module prefix in multimodal models
602
+ """
603
+ return MultiModelKeys.from_string_field(
604
+ language_model="language_model",
605
+ connector="multi_modal_projector",
606
+ tower_model="vision_tower",
607
+ )
608
+
609
+ def get_mm_lora_token_counts(
610
+ self,
611
+ *,
612
+ modality: str,
613
+ mm_kwargs: MultiModalKwargsItem | None,
614
+ num_mm_embeds: int,
615
+ ) -> tuple[int, int | None]:
616
+ del modality
617
+
618
+ pixel_values = mm_kwargs.get("pixel_values") if mm_kwargs else None
619
+ if pixel_values is None or not isinstance(pixel_values.data, torch.Tensor):
620
+ return num_mm_embeds, num_mm_embeds
621
+
622
+ # Unpad runs after the connector, so `num_mm_embeds` is not invertible.
623
+ num_tiles = pixel_values.data.shape[0]
624
+ encoder_info = get_vision_encoder_info(self.config)
625
+ tile_size = encoder_info.get_image_size()
626
+ tokens_per_tile = encoder_info.get_num_image_tokens(
627
+ image_width=tile_size, image_height=tile_size
628
+ )
629
+ selected_per_tile = get_num_selected_vision_tokens(
630
+ tokens_per_tile, self.config.vision_feature_select_strategy
631
+ )
632
+ return num_tiles * tokens_per_tile, num_tiles * selected_per_tile