@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / b75e85dede
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 9, 23:00 UTC

[1/N] Initialize MM components in context managers (A-D) (#32632)

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

Cyrus Leung committed Jan 20, 2026 at 14:12 UTC b75e85dedea304a0da7b6ab696a159ab6a202741
11 files changed +239 -267
vllm/model_executor/models/aria.py
+21 -32
@@ -15,9 +15,7 @@ from vllm.distributed import get_tensor_model_parallel_rank
15 from vllm.model_executor.layers.activation import get_act_fn
16 from vllm.model_executor.layers.fused_moe import SharedFusedMoE
17 from vllm.model_executor.layers.linear import ColumnParallelLinear, RowParallelLinear
18 -from vllm.model_executor.layers.logits_processor import LogitsProcessor
18 from vllm.model_executor.layers.quantization import QuantizationConfig
20 -from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
19 from vllm.model_executor.model_loader.weight_utils import (
20 default_weight_loader,
21 maybe_remap_kv_scale_name,
@@ -539,30 +537,22 @@ class AriaForConditionalGeneration(nn.Module, SupportsMultiModal):
537 quant_config = vllm_config.quant_config
538
539 self.config = config
542 - self.vision_tower = AriaVisionTransformer(
543 - config.vision_config,
544 - quant_config=quant_config,
545 - prefix=f"{prefix}.vision_tower",
546 - )
547 - self.multi_modal_projector = AriaProjector(
548 - config, prefix=maybe_prefix(prefix, "multi_modal_projector")
549 - )
550 - self.vocab_size = config.text_config.vocab_size
551 - self.language_model = AriaTextModel(
552 - vllm_config=vllm_config.with_hf_config(config.text_config),
553 - prefix=maybe_prefix(prefix, "language_model.model"),
554 - )
555 - self.pad_token_id = (
556 - self.config.pad_token_id if self.config.pad_token_id is not None else -1
557 - )
558 - self.lm_head = ParallelLMHead(
559 - self.vocab_size,
560 - config.text_config.hidden_size,
561 - quant_config=quant_config,
562 - prefix=maybe_prefix(prefix, "lm_head"),
563 - )
564 - logit_scale = getattr(config, "logit_scale", 1.0)
565 - self.logits_processor = LogitsProcessor(self.vocab_size, scale=logit_scale)
540 +
541 + with self._mark_tower_model(vllm_config, "image"):
542 + self.vision_tower = AriaVisionTransformer(
543 + config.vision_config,
544 + quant_config=quant_config,
545 + prefix=f"{prefix}.vision_tower",
546 + )
547 + self.multi_modal_projector = AriaProjector(
548 + config, prefix=maybe_prefix(prefix, "multi_modal_projector")
549 + )
550 +
551 + with self._mark_language_model(vllm_config):
552 + self.language_model = AriaTextModel(
553 + vllm_config=vllm_config.with_hf_config(config.text_config),
554 + prefix=maybe_prefix(prefix, "language_model.model"),
555 + )
556
557 def _parse_and_validate_image_input(
558 self, **kwargs: object
@@ -618,9 +608,6 @@ class AriaForConditionalGeneration(nn.Module, SupportsMultiModal):
608
609 return self.multi_modal_projector(image_outputs, image_attn_mask)
610
621 - def get_language_model(self) -> torch.nn.Module:
622 - return self.language_model
623 -
611 def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:
612 image_input = self._parse_and_validate_image_input(**kwargs)
613 if image_input is None:
@@ -654,9 +641,11 @@ class AriaForConditionalGeneration(nn.Module, SupportsMultiModal):
641
642 return hidden_states
643
657 - def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor:
658 - logits = self.logits_processor(self.lm_head, hidden_states)
659 - return logits
644 + def compute_logits(
645 + self,
646 + hidden_states: torch.Tensor,
647 + ) -> torch.Tensor | None:
648 + return self.language_model.compute_logits(hidden_states)
649
650 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
651 loader = AutoWeightsLoader(self)
vllm/model_executor/models/audioflamingo3.py
+13 -15
@@ -460,20 +460,21 @@ class AudioFlamingo3ForConditionalGeneration(
460 multimodal_config = vllm_config.model_config.multimodal_config
461 self.config = config
462 self.multimodal_config = multimodal_config
463 -
464 - self.audio_tower = AudioFlamingo3Encoder(
465 - config.audio_config,
466 - )
467 - self.multi_modal_projector = AudioFlamingo3MultiModalProjector(config)
468 -
463 self.quant_config = quant_config
464
471 - self.language_model = init_vllm_registered_model(
472 - vllm_config=vllm_config,
473 - hf_config=config.text_config,
474 - prefix=maybe_prefix(prefix, "language_model"),
475 - architectures=["Qwen2ForCausalLM"],
476 - )
465 + with self._mark_tower_model(vllm_config, "audio"):
466 + self.audio_tower = AudioFlamingo3Encoder(
467 + config.audio_config,
468 + )
469 + self.multi_modal_projector = AudioFlamingo3MultiModalProjector(config)
470 +
471 + with self._mark_language_model(vllm_config):
472 + self.language_model = init_vllm_registered_model(
473 + vllm_config=vllm_config,
474 + hf_config=config.text_config,
475 + prefix=maybe_prefix(prefix, "language_model"),
476 + architectures=["Qwen2ForCausalLM"],
477 + )
478
479 self.make_empty_intermediate_tensors = (
480 self.language_model.make_empty_intermediate_tensors
@@ -599,9 +600,6 @@ class AudioFlamingo3ForConditionalGeneration(
600 current_idx += count
601 return tuple(grouped_embeddings)
602
602 - def get_language_model(self) -> torch.nn.Module:
603 - return self.language_model
604 -
603 def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:
604 audio_input = self._parse_and_validate_audio_input(**kwargs)
605 if audio_input is None:
vllm/model_executor/models/aya_vision.py
+17 -18
@@ -343,21 +343,23 @@ class AyaVisionForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsP
343 self.quant_config = quant_config
344 self.multimodal_config = multimodal_config
345
346 - self.vision_tower = SiglipVisionModel(
347 - config.vision_config,
348 - quant_config,
349 - num_hidden_layers_override=num_hidden_layers,
350 - prefix=maybe_prefix(prefix, "vision_model"),
351 - )
352 - self.vocab_size = config.text_config.vocab_size
353 - self.multi_modal_projector = AyaVisionMultiModalProjector(config)
354 - self.language_model = init_vllm_registered_model(
355 - vllm_config=vllm_config,
356 - hf_config=config.text_config,
357 - prefix=maybe_prefix(prefix, "model"),
358 - # Cohere2ForCausalLM and CohereForCausalLM are the same on vllm
359 - architectures=["Cohere2ForCausalLM"],
360 - )
346 + with self._mark_tower_model(vllm_config, "image"):
347 + self.vision_tower = SiglipVisionModel(
348 + config.vision_config,
349 + quant_config,
350 + num_hidden_layers_override=num_hidden_layers,
351 + prefix=maybe_prefix(prefix, "vision_model"),
352 + )
353 + self.multi_modal_projector = AyaVisionMultiModalProjector(config)
354 +
355 + with self._mark_language_model(vllm_config):
356 + self.language_model = init_vllm_registered_model(
357 + vllm_config=vllm_config,
358 + hf_config=config.text_config,
359 + prefix=maybe_prefix(prefix, "model"),
360 + # Cohere2ForCausalLM and CohereForCausalLM are the same on vllm
361 + architectures=["Cohere2ForCausalLM"],
362 + )
363
364 @property
365 def dtype(self):
@@ -410,9 +412,6 @@ class AyaVisionForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsP
412 },
413 )
414
413 - def get_language_model(self) -> torch.nn.Module:
414 - return self.language_model
415 -
415 def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:
416 image_input = self._parse_and_validate_image_input(**kwargs)
417 if image_input is None:
vllm/model_executor/models/bagel.py
+36 -43
@@ -44,6 +44,7 @@ from .interfaces import (
44 SupportsLoRA,
45 SupportsMultiModal,
46 SupportsPP,
47 + TowerMissingLayer,
48 )
49 from .siglip import SiglipVisionModel
50 from .utils import (
@@ -373,12 +374,13 @@ class BagelForConditionalGeneration(
374
375 # Initialize language model (Qwen2)
376 # Pass the llm_config from BagelConfig to initialize Qwen2 properly
376 - self.language_model = init_vllm_registered_model(
377 - vllm_config=vllm_config,
378 - hf_config=config.llm_config,
379 - prefix=maybe_prefix(prefix, "language_model"),
380 - architectures=["Qwen2ForCausalLM"],
381 - )
377 + with self._mark_language_model(vllm_config):
378 + self.language_model = init_vllm_registered_model(
379 + vllm_config=vllm_config,
380 + hf_config=config.llm_config,
381 + prefix=maybe_prefix(prefix, "language_model"),
382 + architectures=["Qwen2ForCausalLM"],
383 + )
384
385 # Initialize vision model (SigLIP) if visual understanding is enabled
386 if config.visual_und:
@@ -398,34 +400,35 @@ class BagelForConditionalGeneration(
400 )
401 vit_config.vision_use_head = False
402
401 - self.vit_model = SiglipVisionModel(
402 - config=vit_config,
403 - quant_config=quant_config,
404 - prefix=maybe_prefix(prefix, "vit_model"),
405 - )
403 + with self._mark_tower_model(vllm_config, "image"):
404 + self.vit_model = SiglipVisionModel(
405 + config=vit_config,
406 + quant_config=quant_config,
407 + prefix=maybe_prefix(prefix, "vit_model"),
408 + )
409
407 - # Initialize connector (MLP)
408 - vit_hidden_size = config.vit_config.hidden_size
409 - llm_hidden_size = config.llm_config.hidden_size
410 -
411 - self.connector = BagelVisionMLP(
412 - in_features=vit_hidden_size,
413 - hidden_features=llm_hidden_size,
414 - out_features=llm_hidden_size,
415 - act_layer=config.connector_act,
416 - quant_config=quant_config,
417 - prefix=maybe_prefix(prefix, "connector"),
418 - )
410 + # Initialize connector (MLP)
411 + vit_hidden_size = config.vit_config.hidden_size
412 + llm_hidden_size = config.llm_config.hidden_size
413 +
414 + self.connector = BagelVisionMLP(
415 + in_features=vit_hidden_size,
416 + hidden_features=llm_hidden_size,
417 + out_features=llm_hidden_size,
418 + act_layer=config.connector_act,
419 + quant_config=quant_config,
420 + prefix=maybe_prefix(prefix, "connector"),
421 + )
422
420 - # Position embedding for vision tokens
421 - self.vit_pos_embed = PositionEmbedding(
422 - max_num_patch_per_side=config.vit_max_num_patch_per_side,
423 - hidden_size=llm_hidden_size,
424 - )
423 + # Position embedding for vision tokens
424 + self.vit_pos_embed = PositionEmbedding(
425 + max_num_patch_per_side=config.vit_max_num_patch_per_side,
426 + hidden_size=llm_hidden_size,
427 + )
428 else:
426 - self.vit_model = None
427 - self.connector = None
428 - self.vit_pos_embed = None
429 + self.vit_model = TowerMissingLayer("image")
430 + self.connector = TowerMissingLayer("image")
431 + self.vit_pos_embed = TowerMissingLayer("image")
432
433 self.make_empty_intermediate_tensors = (
434 self.language_model.make_empty_intermediate_tensors
@@ -502,9 +505,6 @@ class BagelForConditionalGeneration(
505
506 return self._process_image_input(image_input)
507
505 - def get_language_model(self) -> nn.Module:
506 - return self.language_model
507 -
508 def forward(
509 self,
510 input_ids: torch.Tensor,
@@ -540,14 +540,6 @@ class BagelForConditionalGeneration(
540
541 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
542 """Load weights from checkpoint."""
543 - skip_prefixes = []
544 - # Skip vit_pos_embed.pos_embed as it's handled by PositionEmbedding module
545 - skip_prefixes.append("vit_pos_embed.pos_embed")
546 -
547 - # If visual understanding is disabled, skip vision-related weights
548 - if self.vit_model is None:
549 - skip_prefixes.extend(["vit_model.", "connector.", "vit_pos_embed"])
550 -
543 # Skip generation-related weights since we only support text2text and image2text
544 # Filter out all image generation components:
545 # - 'moe_gen': MoE generation weights
@@ -587,5 +579,6 @@ class BagelForConditionalGeneration(
579
580 filtered_weights.append((name, tensor))
581
590 - loader = AutoWeightsLoader(self, skip_prefixes=skip_prefixes)
582 + # Skip vit_pos_embed.pos_embed as it's handled by PositionEmbedding module
583 + loader = AutoWeightsLoader(self, skip_prefixes=["vit_pos_embed.pos_embed"])
584 return loader.load_weights(filtered_weights, mapper=self.hf_to_vllm_mapper)
vllm/model_executor/models/blip2.py
+24 -30
@@ -549,31 +549,31 @@ class Blip2ForConditionalGeneration(
549 + 1 # include class token
550 )
551
552 - # TODO: Optionally initializes this for supporting embeddings.
553 - self.vision_model = BlipVisionModel(vision_config, quant_config)
554 -
555 - self.query_tokens = nn.Parameter(
556 - torch.zeros(1, config.num_query_tokens, config.qformer_config.hidden_size)
557 - )
558 -
559 - self.qformer = Blip2QFormerModel(
560 - config.qformer_config,
561 - cache_config=cache_config,
562 - quant_config=quant_config,
563 - prefix=f"{prefix}.qformer",
564 - )
565 -
566 - self.language_projection = nn.Linear(
567 - config.qformer_config.hidden_size,
568 - config.text_config.hidden_size,
569 - bias=True,
570 - )
552 + with self._mark_tower_model(vllm_config, "image"):
553 + self.vision_model = BlipVisionModel(vision_config, quant_config)
554 + self.query_tokens = nn.Parameter(
555 + torch.zeros(
556 + 1, config.num_query_tokens, config.qformer_config.hidden_size
557 + )
558 + )
559 + self.qformer = Blip2QFormerModel(
560 + config.qformer_config,
561 + cache_config=cache_config,
562 + quant_config=quant_config,
563 + prefix=f"{prefix}.qformer",
564 + )
565 + self.language_projection = nn.Linear(
566 + config.qformer_config.hidden_size,
567 + config.text_config.hidden_size,
568 + bias=True,
569 + )
570
572 - self.language_model = init_vllm_registered_model(
573 - vllm_config=vllm_config,
574 - hf_config=config.text_config,
575 - prefix=maybe_prefix(prefix, "language_model"),
576 - )
571 + with self._mark_language_model(vllm_config):
572 + self.language_model = init_vllm_registered_model(
573 + vllm_config=vllm_config,
574 + hf_config=config.text_config,
575 + prefix=maybe_prefix(prefix, "language_model"),
576 + )
577
578 self.make_empty_intermediate_tensors = (
579 self.language_model.make_empty_intermediate_tensors
@@ -614,8 +614,6 @@ class Blip2ForConditionalGeneration(
614 return image_features
615
616 def _process_image_pixels(self, inputs: Blip2ImagePixelInputs) -> torch.Tensor:
617 - assert self.vision_model is not None
618 -
617 pixel_values = inputs["data"]
618
619 return self._image_pixels_to_features(self.vision_model, pixel_values)
@@ -624,7 +622,6 @@ class Blip2ForConditionalGeneration(
622 if image_input["type"] == "image_embeds":
623 return image_input["data"]
624
627 - assert self.vision_model is not None
625 image_features = self._process_image_pixels(image_input)
626
627 query_tokens = self.query_tokens.expand(image_features.shape[0], -1, -1)
@@ -635,9 +632,6 @@ class Blip2ForConditionalGeneration(
632
633 return self.language_projection(query_output)
634
638 - def get_language_model(self) -> torch.nn.Module:
639 - return self.language_model
640 -
635 def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:
636 image_input = self._parse_and_validate_image_input(**kwargs)
637 if image_input is None:
vllm/model_executor/models/clip.py
+23 -24
@@ -853,28 +853,30 @@ class CLIPEmbeddingModel(nn.Module, SupportsMultiModal, SupportsQuant):
853 self.text_embed_dim = text_config.hidden_size
854 self.vision_embed_dim = vision_config.hidden_size
855
856 - self.text_model = CLIPTextTransformer(
857 - text_config,
858 - quant_config=quant_config,
859 - prefix=maybe_prefix(prefix, "text_model"),
860 - )
861 - self.vision_model = CLIPVisionTransformer(
862 - vision_config,
863 - quant_config=quant_config,
864 - multimodal_config=multimodal_config,
865 - prefix=maybe_prefix(prefix, "vision_model"),
866 - )
856 + with self._mark_language_model(vllm_config):
857 + self.text_model = CLIPTextTransformer(
858 + text_config,
859 + quant_config=quant_config,
860 + prefix=maybe_prefix(prefix, "text_model"),
861 + )
862 + self.text_projection = nn.Linear(
863 + self.text_embed_dim,
864 + self.projection_dim,
865 + bias=False,
866 + )
867
868 - self.visual_projection = nn.Linear(
869 - self.vision_embed_dim,
870 - self.projection_dim,
871 - bias=False,
872 - )
873 - self.text_projection = nn.Linear(
874 - self.text_embed_dim,
875 - self.projection_dim,
876 - bias=False,
877 - )
868 + with self._mark_tower_model(vllm_config, "image"):
869 + self.vision_model = CLIPVisionTransformer(
870 + vision_config,
871 + quant_config=quant_config,
872 + multimodal_config=multimodal_config,
873 + prefix=maybe_prefix(prefix, "vision_model"),
874 + )
875 + self.visual_projection = nn.Linear(
876 + self.vision_embed_dim,
877 + self.projection_dim,
878 + bias=False,
879 + )
880
881 pooler_config = vllm_config.model_config.pooler_config
882 assert pooler_config is not None
@@ -940,9 +942,6 @@ class CLIPEmbeddingModel(nn.Module, SupportsMultiModal, SupportsQuant):
942
943 return self.get_image_features(pixel_values)
944
943 - def get_language_model(self) -> torch.nn.Module:
944 - return self.text_model
945 -
945 def _embed_text_input_ids(
946 self,
947 input_ids: torch.Tensor,
vllm/model_executor/models/cohere2_vision.py
+17 -18
@@ -353,21 +353,23 @@ class Cohere2VisionForConditionalGeneration(nn.Module, SupportsMultiModal, Suppo
353 self.multimodal_config = multimodal_config
354 self._patch_quant_config(config, quant_config)
355
356 - self.vision_tower = SiglipVisionModel(
357 - config.vision_config,
358 - quant_config,
359 - prefix=maybe_prefix(prefix, "vision_tower"),
360 - )
361 - self.vocab_size = config.text_config.vocab_size
362 - self.multi_modal_projector = Cohere2VisionMultiModalProjector(
363 - config, prefix=maybe_prefix(prefix, "multi_modal_projector")
364 - )
365 - self.language_model = init_vllm_registered_model(
366 - vllm_config=vllm_config,
367 - hf_config=config.text_config,
368 - prefix=maybe_prefix(prefix, "language_model"),
369 - architectures=config.text_config.architectures,
370 - )
356 + with self._mark_tower_model(vllm_config, "image"):
357 + self.vision_tower = SiglipVisionModel(
358 + config.vision_config,
359 + quant_config,
360 + prefix=maybe_prefix(prefix, "vision_tower"),
361 + )
362 + self.multi_modal_projector = Cohere2VisionMultiModalProjector(
363 + config, prefix=maybe_prefix(prefix, "multi_modal_projector")
364 + )
365 +
366 + with self._mark_language_model(vllm_config):
367 + self.language_model = init_vllm_registered_model(
368 + vllm_config=vllm_config,
369 + hf_config=config.text_config,
370 + prefix=maybe_prefix(prefix, "language_model"),
371 + architectures=config.text_config.architectures,
372 + )
373
374 @property
375 def dtype(self):
@@ -437,9 +439,6 @@ class Cohere2VisionForConditionalGeneration(nn.Module, SupportsMultiModal, Suppo
439 ):
440 quant_config.modules_to_not_convert.append("vision_tower")
441
440 - def get_language_model(self) -> torch.nn.Module:
441 - return self.language_model
442 -
442 def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:
443 image_input = self._parse_and_validate_image_input(**kwargs)
444 if image_input is None:
vllm/model_executor/models/deepseek_ocr.py
+40 -41
@@ -383,46 +383,48 @@ class DeepseekOCRForCausalLM(nn.Module, SupportsMultiModal, SupportsPP, Supports
383 tokenizer = cached_tokenizer_from_config(model_config)
384 self.image_token_id = tokenizer.vocab[_IMAGE_TOKEN]
385
386 - self.sam_model = build_sam_vit_b()
387 - clip_vision_config = CLIPVisionConfig(
388 - hidden_size=1024,
389 - intermediate_size=4096,
390 - num_attention_heads=16,
391 - num_hidden_layers=24,
392 - image_size=224,
393 - patch_size=14,
394 - projection_dim=512,
395 - layer_norm_eps=1e-5,
396 - )
397 - self.vision_model = DeepCLIPVisionTransformer(
398 - config=clip_vision_config,
399 - quant_config=quant_config,
400 - multimodal_config=multimodal_config,
401 - prefix=maybe_prefix(prefix, "vision_model"),
402 - )
403 -
404 - self.projector = MlpProjector(self.projector_config)
405 - self.tile_tag = config.tile_tag
406 - self.global_view_pos = config.global_view_pos
407 -
408 - # special token for image token sequence format
409 - n_embed = self.projector_config.n_embed
410 - embed_std = 1 / torch.sqrt(torch.tensor(n_embed, dtype=torch.float32))
411 - if self.tile_tag == "2D":
412 - # <|view_separator|>, <|\n|>
413 - self.image_newline = nn.Parameter(torch.randn(n_embed) * embed_std)
414 - # This is a typo in original implementation
415 - self.view_seperator = nn.Parameter(torch.randn(n_embed) * embed_std)
416 - else:
417 - raise ValueError(
418 - f"Only 2D tile_tag is supported currently, got: {self.tile_tag}"
386 + with self._mark_tower_model(vllm_config, "image"):
387 + self.sam_model = build_sam_vit_b()
388 + clip_vision_config = CLIPVisionConfig(
389 + hidden_size=1024,
390 + intermediate_size=4096,
391 + num_attention_heads=16,
392 + num_hidden_layers=24,
393 + image_size=224,
394 + patch_size=14,
395 + projection_dim=512,
396 + layer_norm_eps=1e-5,
397 + )
398 + self.vision_model = DeepCLIPVisionTransformer(
399 + config=clip_vision_config,
400 + quant_config=quant_config,
401 + multimodal_config=multimodal_config,
402 + prefix=maybe_prefix(prefix, "vision_model"),
403 )
404
421 - self.language_model = init_vllm_registered_model(
422 - vllm_config=vllm_config,
423 - hf_config=self.text_config,
424 - prefix=maybe_prefix(prefix, "language_model"),
425 - )
405 + self.projector = MlpProjector(self.projector_config)
406 + self.tile_tag = config.tile_tag
407 + self.global_view_pos = config.global_view_pos
408 +
409 + # special token for image token sequence format
410 + n_embed = self.projector_config.n_embed
411 + embed_std = 1 / torch.sqrt(torch.tensor(n_embed, dtype=torch.float32))
412 + if self.tile_tag == "2D":
413 + # <|view_separator|>, <|\n|>
414 + self.image_newline = nn.Parameter(torch.randn(n_embed) * embed_std)
415 + # This is a typo in original implementation
416 + self.view_seperator = nn.Parameter(torch.randn(n_embed) * embed_std)
417 + else:
418 + raise ValueError(
419 + f"Only 2D tile_tag is supported currently, got: {self.tile_tag}"
420 + )
421 +
422 + with self._mark_language_model(vllm_config):
423 + self.language_model = init_vllm_registered_model(
424 + vllm_config=vllm_config,
425 + hf_config=self.text_config,
426 + prefix=maybe_prefix(prefix, "language_model"),
427 + )
428
429 self.make_empty_intermediate_tensors = (
430 self.language_model.make_empty_intermediate_tensors
@@ -552,9 +554,6 @@ class DeepseekOCRForCausalLM(nn.Module, SupportsMultiModal, SupportsPP, Supports
554
555 return vision_features
556
555 - def get_language_model(self) -> torch.nn.Module:
556 - return self.language_model
557 -
557 def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings | None:
558 image_input = self._parse_and_validate_image_input(**kwargs)
559 if image_input is None:
vllm/model_executor/models/deepseek_vl2.py
+29 -30
@@ -374,37 +374,39 @@ class DeepseekVLV2ForCausalLM(nn.Module, SupportsMultiModal, SupportsPP):
374 tokenizer = cached_tokenizer_from_config(model_config)
375 self.image_token_id: int = tokenizer.vocab[_IMAGE_TOKEN]
376
377 - self.vision = self._init_vision_module(
378 - self.vision_config, quant_config, maybe_prefix(prefix, "vision")
379 - )
377 + with self._mark_tower_model(vllm_config, "image"):
378 + self.vision = self._init_vision_module(
379 + self.vision_config, quant_config, maybe_prefix(prefix, "vision")
380 + )
381
381 - self.projector = MlpProjector(self.projector_config)
382 - self.tile_tag = config.tile_tag
383 - self.global_view_pos = config.global_view_pos
382 + self.projector = MlpProjector(self.projector_config)
383 + self.tile_tag = config.tile_tag
384 + self.global_view_pos = config.global_view_pos
385
385 - # special token for image token sequence format
386 - embed_std = 1 / torch.sqrt(
387 - torch.tensor(self.projector_config.n_embed, dtype=torch.float32)
388 - )
389 - if self.tile_tag == "2D":
390 - # <|view_seperator|>, <|\n|>
391 - self.image_newline = nn.Parameter(
392 - torch.randn(self.projector_config.n_embed) * embed_std
393 - )
394 - # This is a typo in original implementation
395 - self.view_seperator = nn.Parameter(
396 - torch.randn(self.projector_config.n_embed) * embed_std
397 - )
398 - else:
399 - raise ValueError(
400 - f"Only 2D tile_tag is supported currently, got: {self.tile_tag}"
386 + # special token for image token sequence format
387 + embed_std = 1 / torch.sqrt(
388 + torch.tensor(self.projector_config.n_embed, dtype=torch.float32)
389 )
390 + if self.tile_tag == "2D":
391 + # <|view_seperator|>, <|\n|>
392 + self.image_newline = nn.Parameter(
393 + torch.randn(self.projector_config.n_embed) * embed_std
394 + )
395 + # This is a typo in original implementation
396 + self.view_seperator = nn.Parameter(
397 + torch.randn(self.projector_config.n_embed) * embed_std
398 + )
399 + else:
400 + raise ValueError(
401 + f"Only 2D tile_tag is supported currently, got: {self.tile_tag}"
402 + )
403
403 - self.language_model = init_vllm_registered_model(
404 - vllm_config=vllm_config,
405 - hf_config=self.text_config,
406 - prefix=maybe_prefix(prefix, "language"),
407 - )
404 + with self._mark_language_model(vllm_config):
405 + self.language_model = init_vllm_registered_model(
406 + vllm_config=vllm_config,
407 + hf_config=self.text_config,
408 + prefix=maybe_prefix(prefix, "language"),
409 + )
410
411 self.make_empty_intermediate_tensors = (
412 self.language_model.make_empty_intermediate_tensors
@@ -603,9 +605,6 @@ class DeepseekVLV2ForCausalLM(nn.Module, SupportsMultiModal, SupportsPP):
605 pixel_values=pixel_values, images_spatial_crop=images_spatial_crop
606 )
607
606 - def get_language_model(self) -> torch.nn.Module:
607 - return self.language_model
608 -
608 def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:
609 image_input = self._parse_and_validate_image_input(**kwargs)
610 if image_input is None:
vllm/model_executor/models/dots_ocr.py
+15 -15
@@ -689,18 +689,21 @@ class DotsOCRForCausalLM(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA
689 else:
690 vision_config = self.config.vision_config
691
692 - self.vision_tower = DotsVisionTransformer(
693 - vision_config,
694 - quant_config=self.quant_config,
695 - multimodal_config=multimodal_config,
696 - prefix=maybe_prefix(prefix, "vision_tower"),
697 - )
698 - self.language_model: Qwen2ForCausalLM = init_vllm_registered_model(
699 - vllm_config=vllm_config,
700 - hf_config=self.config,
701 - prefix=maybe_prefix(prefix, "language_model"),
702 - architectures=["Qwen2ForCausalLM"],
703 - )
692 + with self._mark_tower_model(vllm_config, "image"):
693 + self.vision_tower = DotsVisionTransformer(
694 + vision_config,
695 + quant_config=self.quant_config,
696 + multimodal_config=multimodal_config,
697 + prefix=maybe_prefix(prefix, "vision_tower"),
698 + )
699 +
700 + with self._mark_language_model(vllm_config):
701 + self.language_model: Qwen2ForCausalLM = init_vllm_registered_model(
702 + vllm_config=vllm_config,
703 + hf_config=self.config,
704 + prefix=maybe_prefix(prefix, "language_model"),
705 + architectures=["Qwen2ForCausalLM"],
706 + )
707
708 self.make_empty_intermediate_tensors = (
709 self.language_model.make_empty_intermediate_tensors
@@ -763,9 +766,6 @@ class DotsOCRForCausalLM(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA
766
767 return image_embeds.split(sizes)
768
766 - def get_language_model(self) -> torch.nn.Module:
767 - return self.language_model
768 -
769 def get_num_mm_encoder_tokens(self, num_image_tokens: int) -> int:
770 merge_size = self.vision_tower.spatial_merge_size
771 return num_image_tokens * (merge_size**2)
vllm/model_executor/models/interfaces.py
+4 -1
@@ -83,7 +83,10 @@ class LMMissingLayer(nn.Module):
83 class TowerMissingLayer(nn.Module):
84 packed_modules_mapping: dict[str, list[str]] = {}
85
86 - def __init__(self, modalities: set[str]) -> None:
86 + def __init__(self, modalities: set[str] | str) -> None:
87 + if isinstance(modalities, str):
88 + modalities = {modalities}
89 +
90 super().__init__()
91
92 self.modalities = modalities