@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 7f1bcd18ff
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

[3/N] Initialize MM components in context managers (I-L) (#32650)

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

Cyrus Leung committed Jan 20, 2026 at 18:21 UTC 7f1bcd18ffd9313a72e1d72c55a3d44f4fbb6157
10 files changed +202 -388
vllm/model_executor/models/interns1.py
+13 -18
@@ -547,20 +547,20 @@ class InternS1ForConditionalGeneration(
547 )
548 self.downsample_ratio = config.downsample_ratio
549
550 - self.llm_arch_name = config.text_config.architectures[0]
551 - self.vision_tower = self._init_vision_model(
552 - config,
553 - quant_config=quant_config,
554 - prefix=maybe_prefix(prefix, "vision_tower"),
555 - )
556 -
557 - self.language_model = init_vllm_registered_model(
558 - vllm_config=vllm_config,
559 - hf_config=config.text_config,
560 - prefix=maybe_prefix(prefix, "language_model"),
561 - )
550 + with self._mark_tower_model(vllm_config, {"image", "video"}):
551 + self.vision_tower = self._init_vision_model(
552 + config,
553 + quant_config=quant_config,
554 + prefix=maybe_prefix(prefix, "vision_tower"),
555 + )
556 + self.multi_modal_projector = self._init_mlp1(config)
557
563 - self.multi_modal_projector = self._init_mlp1(config)
558 + with self._mark_language_model(vllm_config):
559 + self.language_model = init_vllm_registered_model(
560 + vllm_config=vllm_config,
561 + hf_config=config.text_config,
562 + prefix=maybe_prefix(prefix, "language_model"),
563 + )
564
565 self.img_context_token_id = None
566 self.video_context_token_id = None
@@ -699,8 +699,6 @@ class InternS1ForConditionalGeneration(
699 ):
700 return image_input["data"]
701
702 - assert self.vision_tower is not None
703 -
702 image_embeds = self.extract_feature(image_input["pixel_values"])
703
704 num_patches = image_input["num_patches"]
@@ -737,9 +735,6 @@ class InternS1ForConditionalGeneration(
735 def _set_visual_token_mask(self, input_ids: torch.Tensor) -> None:
736 self.visual_token_mask = None
737
740 - def get_language_model(self) -> torch.nn.Module:
741 - return self.language_model
742 -
738 def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:
739 modalities = self._parse_and_validate_multimodal_inputs(**kwargs)
740 if not modalities:
vllm/model_executor/models/internvl.py
+16 -19
@@ -1092,22 +1092,24 @@ class InternVLChatModel(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA)
1092 self.downsample_ratio = config.downsample_ratio
1093 self.ps_version = config.ps_version
1094
1095 - self.llm_arch_name = config.text_config.architectures[0]
1096 - self.is_mono = self.llm_arch_name == "InternLM2VEForCausalLM"
1097 - self.vision_model = self._init_vision_model(
1098 - config,
1099 - quant_config=quant_config,
1100 - is_mono=self.is_mono,
1101 - prefix=maybe_prefix(prefix, "vision_model"),
1102 - )
1095 + llm_arch_name = config.text_config.architectures[0]
1096 + self.is_mono = llm_arch_name == "InternLM2VEForCausalLM"
1097
1104 - self.language_model = init_vllm_registered_model(
1105 - vllm_config=vllm_config,
1106 - hf_config=config.text_config,
1107 - prefix=maybe_prefix(prefix, "language_model"),
1108 - )
1098 + with self._mark_tower_model(vllm_config, {"image", "video"}):
1099 + self.vision_model = self._init_vision_model(
1100 + config,
1101 + quant_config=quant_config,
1102 + is_mono=self.is_mono,
1103 + prefix=maybe_prefix(prefix, "vision_model"),
1104 + )
1105 + self.mlp1 = self._init_mlp1(config)
1106
1110 - self.mlp1 = self._init_mlp1(config)
1107 + with self._mark_language_model(vllm_config):
1108 + self.language_model = init_vllm_registered_model(
1109 + vllm_config=vllm_config,
1110 + hf_config=config.text_config,
1111 + prefix=maybe_prefix(prefix, "language_model"),
1112 + )
1113
1114 self.img_context_token_id = None
1115 self.video_context_token_id = None
@@ -1281,8 +1283,6 @@ class InternVLChatModel(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA)
1283 ):
1284 return image_input["data"]
1285
1284 - assert self.vision_model is not None
1285 -
1286 image_embeds = self.extract_feature(image_input["pixel_values_flat"])
1287
1288 num_patches = image_input["num_patches"]
@@ -1325,9 +1325,6 @@ class InternVLChatModel(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA)
1325 else:
1326 self.visual_token_mask = None
1327
1328 - def get_language_model(self) -> torch.nn.Module:
1329 - return self.language_model
1330 -
1328 def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:
1329 modalities = self._parse_and_validate_multimodal_inputs(**kwargs)
1330 if not modalities:
vllm/model_executor/models/isaac.py
+18 -25
@@ -1342,11 +1342,14 @@ class IsaacForConditionalGeneration(
1342 "mrope_interleaved", rope_scaling["mrope_interleaved"]
1343 )
1344 target_cfg.rope_parameters = rope_parameters
1345 - self.language_model = init_vllm_registered_model(
1346 - vllm_config=vllm_config,
1347 - architectures=["Qwen3ForCausalLM"],
1348 - prefix=maybe_prefix(prefix, "language_model"),
1349 - )
1345 +
1346 + with self._mark_language_model(vllm_config):
1347 + self.language_model = init_vllm_registered_model(
1348 + vllm_config=vllm_config,
1349 + architectures=["Qwen3ForCausalLM"],
1350 + prefix=maybe_prefix(prefix, "language_model"),
1351 + )
1352 +
1353 self.make_empty_intermediate_tensors = (
1354 self.language_model.make_empty_intermediate_tensors
1355 )
@@ -1363,14 +1366,16 @@ class IsaacForConditionalGeneration(
1366 vision_cfg._attn_implementation = attn_impl
1367
1368 hidden_dim = vision_cfg.hidden_size * (vision_cfg.pixel_shuffle_scale_factor**2)
1366 - self.vision_embedding = IsaacVisionEmbedding(
1367 - vision_cfg=vision_cfg,
1368 - hidden_dim=hidden_dim,
1369 - output_dim=config.hidden_size,
1370 - quant_config=quant_config,
1371 - multimodal_config=self.multimodal_config,
1372 - prefix=maybe_prefix(prefix, "vision_embedding"),
1373 - )
1369 +
1370 + with self._mark_tower_model(vllm_config, "image"):
1371 + self.vision_embedding = IsaacVisionEmbedding(
1372 + vision_cfg=vision_cfg,
1373 + hidden_dim=hidden_dim,
1374 + output_dim=config.hidden_size,
1375 + quant_config=quant_config,
1376 + multimodal_config=self.multimodal_config,
1377 + prefix=maybe_prefix(prefix, "vision_embedding"),
1378 + )
1379
1380 def iter_mm_grid_hw(
1381 self, input_tokens: list[int], mm_features: list[MultiModalFeatureSpec]
@@ -1457,18 +1462,6 @@ class IsaacForConditionalGeneration(
1462 return ()
1463 return self._process_image_input(image_input)
1464
1460 - def get_multimodal_embeddings(
1461 - self, **kwargs: object
1462 - ) -> MultiModalEmbeddings | None:
1463 - # Backward compatibility for older runners.
1464 - embeddings = self.embed_multimodal(**kwargs)
1465 - if not embeddings:
1466 - return []
1467 - return embeddings
1468 -
1469 - def get_language_model(self) -> torch.nn.Module:
1470 - return self.language_model
1471 -
1465 def forward(
1466 self,
1467 input_ids: torch.Tensor,
vllm/model_executor/models/kanana_v.py
+15 -13
@@ -586,16 +586,21 @@ class KananaVForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP)
586 config = vllm_config.model_config.hf_config
587 self.config = config
588
589 - self.vision_model = CustomQwen2VLVE._from_config(config.vision_config)
590 - self.abstractor = DynamicCAbstractor(
591 - config.projector_config, num_input_tokens=self.vision_model.get_num_tokens()
592 - )
593 - self.language_model = init_vllm_registered_model(
594 - vllm_config=vllm_config,
595 - hf_config=config.text_config,
596 - prefix=maybe_prefix(prefix, "model"),
597 - architectures=["LlamaForCausalLM"],
598 - )
589 + with self._mark_tower_model(vllm_config, "image"):
590 + self.vision_model = CustomQwen2VLVE._from_config(config.vision_config)
591 + self.abstractor = DynamicCAbstractor(
592 + config.projector_config,
593 + num_input_tokens=self.vision_model.get_num_tokens(),
594 + )
595 +
596 + with self._mark_language_model(vllm_config):
597 + self.language_model = init_vllm_registered_model(
598 + vllm_config=vllm_config,
599 + hf_config=config.text_config,
600 + prefix=maybe_prefix(prefix, "model"),
601 + architectures=["LlamaForCausalLM"],
602 + )
603 +
604 self.make_empty_intermediate_tensors = (
605 self.language_model.make_empty_intermediate_tensors
606 )
@@ -718,9 +723,6 @@ class KananaVForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP)
723 visual_embeds = self.forward_projector(visual_features, image_metas=image_metas)
724 return visual_embeds
725
721 - def get_language_model(self) -> torch.nn.Module:
722 - return self.language_model
723 -
726 def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:
727 image_input = self._parse_and_validate_image_input(**kwargs)
728 if image_input is None:
vllm/model_executor/models/keye.py
+21 -23
@@ -1242,7 +1242,7 @@ class KeyeMultiModalProcessor(BaseMultiModalProcessor[KeyeProcessingInfo]):
1242 return _keye_field_config(hf_inputs)
1243
1244
1245 -class BaseKeyeModule(nn.Module):
1245 +class BaseKeyeModule(nn.Module, SupportsMultiModal):
1246 packed_modules_mapping = {
1247 "qkv_proj": [
1248 "q_proj",
@@ -1280,25 +1280,26 @@ class BaseKeyeModule(nn.Module):
1280 self.config = config
1281 self.multimodal_config = multimodal_config
1282
1283 - self.visual = KeyeSiglipVisionModel(
1284 - config.vision_config,
1285 - quant_config=quant_config,
1286 - multimodal_config=multimodal_config,
1287 - prefix=maybe_prefix(prefix, "visual"),
1288 - )
1289 -
1290 - self.mlp_AR = self._build_projector(
1291 - config,
1292 - config.vision_config,
1293 - quant_config=quant_config,
1294 - prefix=maybe_prefix(prefix, "mlp_AR"),
1295 - )
1283 + with self._mark_tower_model(vllm_config, {"image", "video"}):
1284 + self.visual = KeyeSiglipVisionModel(
1285 + config.vision_config,
1286 + quant_config=quant_config,
1287 + multimodal_config=multimodal_config,
1288 + prefix=maybe_prefix(prefix, "visual"),
1289 + )
1290 + self.mlp_AR = self._build_projector(
1291 + config,
1292 + config.vision_config,
1293 + quant_config=quant_config,
1294 + prefix=maybe_prefix(prefix, "mlp_AR"),
1295 + )
1296
1297 - self.language_model = init_vllm_registered_model(
1298 - vllm_config=vllm_config,
1299 - prefix=maybe_prefix(prefix, "language_model"),
1300 - architectures=["Qwen3ForCausalLM"],
1301 - )
1297 + with self._mark_language_model(vllm_config):
1298 + self.language_model = init_vllm_registered_model(
1299 + vllm_config=vllm_config,
1300 + prefix=maybe_prefix(prefix, "language_model"),
1301 + architectures=["Qwen3ForCausalLM"],
1302 + )
1303
1304 self.make_empty_intermediate_tensors = (
1305 self.language_model.make_empty_intermediate_tensors
@@ -1312,7 +1313,7 @@ class BaseKeyeModule(nn.Module):
1313 quant_config: QuantizationConfig | None = None,
1314 prefix: str = "",
1315 ) -> nn.Module:
1315 - raise ValueError("Need projector")
1316 + raise NotImplementedError("Need projector")
1317
1318 def _process_image_input(self, image_input: Any) -> tuple[torch.Tensor, ...]:
1319 siglip_position_ids = list()
@@ -1429,9 +1430,6 @@ class BaseKeyeModule(nn.Module):
1430
1431 return modalities
1432
1432 - def get_language_model(self) -> torch.nn.Module:
1433 - return self.language_model
1434 -
1433 def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings | None:
1434 modalities = self._parse_and_validate_multimodal_inputs(**kwargs)
1435 if not modalities:
vllm/model_executor/models/kimi_vl.py
+27 -188
@@ -42,7 +42,6 @@
42 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
43 # SOFTWARE.
44
45 -import copy
45 import math
46 from collections.abc import Iterable, Mapping, Sequence
47 from dataclasses import dataclass
@@ -50,23 +49,12 @@ from typing import Annotated, Any, Literal
49
50 import torch
51 from torch import nn
53 -from transformers import BatchFeature, DeepseekV2Config
52 +from transformers import BatchFeature
53 from transformers.activations import GELUActivation
54
55 from vllm.config import VllmConfig
56 from vllm.config.multimodal import BaseDummyOptions
58 -from vllm.distributed import get_pp_group
59 -from vllm.model_executor.layers.fused_moe import FusedMoE
57 from vllm.model_executor.layers.linear import ReplicatedLinear
61 -from vllm.model_executor.layers.logits_processor import LogitsProcessor
62 -from vllm.model_executor.layers.vocab_parallel_embedding import (
63 - ParallelLMHead,
64 -)
65 -from vllm.model_executor.model_loader.weight_utils import (
66 - default_weight_loader,
67 - maybe_remap_kv_scale_name,
68 -)
69 -from vllm.model_executor.models.deepseek_v2 import DeepseekV2Model
58 from vllm.model_executor.models.interfaces import SupportsMultiModal, SupportsPP
59 from vllm.model_executor.models.moonvit import MoonVitPretrainedModel
60 from vllm.multimodal import MULTIMODAL_REGISTRY
@@ -92,7 +80,7 @@ from vllm.sequence import IntermediateTensors
80 from vllm.transformers_utils.configs import KimiVLConfig, MoonViTConfig
81 from vllm.utils.tensor_schema import TensorSchema, TensorShape
82
95 -from .utils import PPMissingLayer, is_pp_missing_parameter, maybe_prefix
83 +from .utils import AutoWeightsLoader, init_vllm_registered_model, maybe_prefix
84 from .vision import run_dp_sharded_mrope_vision_model
85
86
@@ -315,48 +303,41 @@ class KimiVLForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP):
303 super().__init__()
304 model_config = vllm_config.model_config
305 config: KimiVLConfig = model_config.hf_config
318 - self.config = config
306 quant_config = vllm_config.quant_config
307
308 + self.config = config
309 + self.quant_config = quant_config
310 +
311 assert isinstance(config.vision_config, MoonViTConfig)
312 self.use_data_parallel = (
313 model_config.multimodal_config.mm_encoder_tp_mode == "data"
314 )
315 self.hidden_size = config.text_config.hidden_size
326 - self.vision_tower = MoonVitPretrainedModel(
327 - config.vision_config,
328 - multimodal_config=model_config.multimodal_config,
329 - prefix=maybe_prefix(prefix, "vision_tower"),
330 - )
316
332 - self.multi_modal_projector = KimiVLMultiModalProjector(
333 - config=config,
334 - use_data_parallel=self.use_data_parallel,
335 - prefix=maybe_prefix(prefix, "multi_modal_projector"),
336 - )
317 + with self._mark_tower_model(vllm_config, "image"):
318 + self.vision_tower = MoonVitPretrainedModel(
319 + config.vision_config,
320 + multimodal_config=model_config.multimodal_config,
321 + prefix=maybe_prefix(prefix, "vision_tower"),
322 + )
323 + self.multi_modal_projector = KimiVLMultiModalProjector(
324 + config=config,
325 + use_data_parallel=self.use_data_parallel,
326 + prefix=maybe_prefix(prefix, "multi_modal_projector"),
327 + )
328
338 - self.quant_config = quant_config
339 - sub_vllm_config = copy.deepcopy(vllm_config)
340 - sub_vllm_config.model_config.hf_config = (
341 - sub_vllm_config.model_config.hf_config.text_config
342 - )
343 - self.language_model = DeepseekV2Model(
344 - vllm_config=sub_vllm_config,
345 - prefix=maybe_prefix(prefix, "language_model"),
346 - )
347 - if get_pp_group().is_last_rank:
348 - self.lm_head = ParallelLMHead(
349 - config.vocab_size,
350 - config.text_config.hidden_size,
351 - prefix=maybe_prefix(prefix, "lm_head"),
329 + with self._mark_language_model(vllm_config):
330 + self.language_model = init_vllm_registered_model(
331 + vllm_config=vllm_config,
332 + hf_config=config.text_config,
333 + prefix=maybe_prefix(prefix, "language_model"),
334 + architectures=["DeepseekV2ForCausalLM"],
335 )
353 - else:
354 - self.lm_head = PPMissingLayer()
336 +
337 self.make_empty_intermediate_tensors = (
338 self.language_model.make_empty_intermediate_tensors
339 )
358 - logit_scale = getattr(config, "logit_scale", 1.0)
359 - self.logits_processor = LogitsProcessor(config.vocab_size, scale=logit_scale)
340 +
341 self.media_placeholder: int = self.config.media_placeholder_token_id
342
343 def _parse_and_validate_image_input(
@@ -378,8 +359,6 @@ class KimiVLForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP):
359 # perform vt on processored pixel_values
360 @torch.inference_mode()
361 def _process_image_pixels(self, inputs: KimiVLImagePixelInputs) -> torch.Tensor:
381 - assert self.vision_tower is not None
382 -
362 pixel_values = inputs["pixel_values"]
363 image_grid_hws = inputs["image_grid_hws"]
364 if self.use_data_parallel:
@@ -399,9 +378,6 @@ class KimiVLForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP):
378 lengths = [x.shape[0] for x in image_features]
379 return self.multi_modal_projector(torch.cat(image_features)).split(lengths)
380
402 - def get_language_model(self) -> torch.nn.Module:
403 - return self.language_model
404 -
381 def embed_multimodal(self, **kwargs: object) -> NestedTensors | None:
382 # Validate the multimodal input keyword arguments
383 image_input = self._parse_and_validate_image_input(**kwargs)
@@ -433,145 +409,8 @@ class KimiVLForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP):
409 return hidden_states
410
411 def compute_logits(self, hidden_states: torch.Tensor, **kwargs) -> torch.Tensor:
436 - logits = self.logits_processor(self.lm_head, hidden_states, **kwargs)
437 - return logits
412 + return self.language_model.compute_logits(hidden_states)
413
414 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
440 - config = self.config.text_config
441 - _KEYS_TO_MODIFY_MAPPING = {
442 - "language_model.lm_head": "lm_head",
443 - "language_model.model": "language_model",
444 - }
445 - # only doing this for language model part for now.
446 - stacked_params_mapping = [
447 - # (param_name, shard_name, shard_id)
448 - (".gate_up_proj", ".gate_proj", 0),
449 - (".gate_up_proj", ".up_proj", 1),
450 - ]
451 - use_mha = (
452 - config.model_type == "deepseek"
453 - or config.qk_nope_head_dim + config.qk_rope_head_dim == 0
454 - )
455 - if use_mha:
456 - stacked_params_mapping += [
457 - (".qkv_proj", ".q_proj", "q"),
458 - (".qkv_proj", ".k_proj", "k"),
459 - (".qkv_proj", ".v_proj", "v"),
460 - ]
461 - if getattr(config, "n_routed_experts", None):
462 - # Params for weights, fp8 weight scales, fp8 activation scales
463 - # (param_name, weight_name, expert_id, shard_id)
464 - expert_params_mapping = FusedMoE.make_expert_params_mapping(
465 - self,
466 - ckpt_gate_proj_name="gate_proj",
467 - ckpt_down_proj_name="down_proj",
468 - ckpt_up_proj_name="up_proj",
469 - num_experts=config.n_routed_experts,
470 - )
471 - else:
472 - expert_params_mapping = []
473 -
474 - params_dict = dict(self.named_parameters())
475 -
476 - for args in weights:
477 - name, loaded_weight = args[:2]
478 - kwargs = args[2] if len(args) > 2 else {}
479 - if "rotary_emb.inv_freq" in name:
480 - continue
481 -
482 - spec_layer = get_spec_layer_idx_from_weight_name(config, name)
483 - if spec_layer is not None:
484 - continue # skip spec decode layers for main model
485 -
486 - if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name:
487 - # Models trained using ColossalAI may include these tensors in
488 - # the checkpoint. Skip them.
489 - continue
490 - for key_to_modify, new_key in _KEYS_TO_MODIFY_MAPPING.items():
491 - if key_to_modify in name:
492 - name = name.replace(key_to_modify, new_key)
493 - use_default_weight_loading = False
494 - if "vision" in name:
495 - if self.vision_tower is not None:
496 - # We only do sharding for language model and
497 - # not vision model for now.
498 - use_default_weight_loading = True
499 - else:
500 - for param_name, weight_name, shard_id in stacked_params_mapping:
501 - if weight_name not in name:
502 - continue
503 - # We have mlp.experts[0].gate_proj in the checkpoint.
504 - # Since we handle the experts below in expert_params_mapping,
505 - # we need to skip here BEFORE we update the name, otherwise
506 - # name will be updated to mlp.experts[0].gate_up_proj, which
507 - # will then be updated below in expert_params_mapping
508 - # for mlp.experts[0].gate_gate_up_proj, which breaks load.
509 - if ("mlp.experts." in name) and name not in params_dict:
510 - continue
511 - name = name.replace(weight_name, param_name)
512 - # Skip loading extra bias for GPTQ models.
513 - if name.endswith(".bias") and name not in params_dict:
514 - continue
515 -
516 - if is_pp_missing_parameter(name, self):
517 - continue
518 -
519 - param = params_dict[name]
520 - weight_loader = param.weight_loader
521 - weight_loader(param, loaded_weight, shard_id, **kwargs)
522 - break
523 - else:
524 - for idx, (
525 - param_name,
526 - weight_name,
527 - expert_id,
528 - shard_id,
529 - ) in enumerate(expert_params_mapping):
530 - if weight_name not in name:
531 - continue
532 - name = name.replace(weight_name, param_name)
533 -
534 - if is_pp_missing_parameter(name, self):
535 - continue
536 -
537 - param = params_dict[name]
538 - weight_loader = param.weight_loader
539 - weight_loader(
540 - param,
541 - loaded_weight,
542 - name,
543 - expert_id=expert_id,
544 - shard_id=shard_id,
545 - **kwargs,
546 - )
547 - break
548 - else:
549 - use_default_weight_loading = True
550 - if use_default_weight_loading:
551 - # Skip loading extra bias for GPTQ models.
552 - if name.endswith(".bias") and name not in params_dict:
553 - continue
554 - # Remapping the name of FP8 kv-scale.
555 - name = maybe_remap_kv_scale_name(name, params_dict)
556 - if name is None:
557 - continue
558 -
559 - if is_pp_missing_parameter(name, self):
560 - continue
561 -
562 - param = params_dict[name]
563 - weight_loader = getattr(param, "weight_loader", default_weight_loader)
564 - weight_loader(param, loaded_weight, **kwargs)
565 -
566 -
567 -def get_spec_layer_idx_from_weight_name(
568 - config: DeepseekV2Config, weight_name: str
569 -) -> int | None:
570 - if hasattr(config, "num_nextn_predict_layers") and (
571 - config.num_nextn_predict_layers > 0
572 - ):
573 - layer_idx = config.num_hidden_layers
574 - for i in range(config.num_nextn_predict_layers):
575 - if weight_name.startswith(f"model.layers.{layer_idx + i}."):
576 - return layer_idx + i
577 - return None
415 + loader = AutoWeightsLoader(self)
416 + return loader.load_weights(weights)
vllm/model_executor/models/lfm2_vl.py
+25 -27
@@ -546,38 +546,37 @@ class Lfm2VLForConditionalGeneration(
546 self.multimodal_config = multimodal_config
547 self.use_data_parallel = multimodal_config.mm_encoder_tp_mode == "data"
548
549 - if vision_config.model_type == "siglip2_vision_model":
550 - self.vision_tower = Siglip2Model(
551 - config=vision_config,
552 - quant_config=quant_config,
553 - multimodal_config=multimodal_config,
554 - prefix=maybe_prefix(prefix, "vision_tower"),
555 - )
556 - else:
557 - raise ValueError(
558 - f"Unsupported visual tokenizer model_type: {vision_config.model_type}"
559 - )
549 + with self._mark_tower_model(vllm_config, "image"):
550 + if vision_config.model_type == "siglip2_vision_model":
551 + self.vision_tower = Siglip2Model(
552 + config=vision_config,
553 + quant_config=quant_config,
554 + multimodal_config=multimodal_config,
555 + prefix=maybe_prefix(prefix, "vision_tower"),
556 + )
557 + else:
558 + raise ValueError(
559 + f"Unsupported visual tokenizer type: {vision_config.model_type}"
560 + )
561
561 - self.multi_modal_projector = Lfm2VLMultiModalProjector(
562 - config=config,
563 - use_data_parallel=self.use_data_parallel,
564 - prefix=f"{prefix}.multi_modal_projector",
565 - )
562 + self.multi_modal_projector = Lfm2VLMultiModalProjector(
563 + config=config,
564 + use_data_parallel=self.use_data_parallel,
565 + prefix=maybe_prefix(prefix, "multi_modal_projector"),
566 + )
567
567 - self.language_model = init_vllm_registered_model(
568 - vllm_config=vllm_config,
569 - hf_config=config.text_config,
570 - prefix=maybe_prefix(prefix, "language"),
571 - architectures=config.text_config.architectures,
572 - )
568 + with self._mark_language_model(vllm_config):
569 + self.language_model = init_vllm_registered_model(
570 + vllm_config=vllm_config,
571 + hf_config=config.text_config,
572 + prefix=maybe_prefix(prefix, "language"),
573 + architectures=config.text_config.architectures,
574 + )
575
576 self.make_empty_intermediate_tensors = (
577 self.language_model.make_empty_intermediate_tensors
578 )
579
578 - def get_language_model(self) -> torch.nn.Module:
579 - return self.language_model
580 -
580 def _parse_and_validate_image_input(
581 self, **kwargs: object
582 ) -> LFM2VLImageInputs | None:
@@ -714,8 +713,7 @@ class Lfm2VLForConditionalGeneration(
713 self,
714 hidden_states: torch.Tensor,
715 ) -> torch.Tensor | None:
717 - logits = self.language_model.compute_logits(hidden_states)
718 - return logits
716 + return self.language_model.compute_logits(hidden_states)
717
718 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
719 loader = AutoWeightsLoader(self)
vllm/model_executor/models/llava_next.py
+23 -25
@@ -268,27 +268,30 @@ class LlavaNextForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsP
268 self.config = config
269 self.multimodal_config = multimodal_config
270
271 - # TODO: Optionally initializes this for supporting embeddings.
272 - self.vision_tower = init_vision_tower_for_llava(
273 - config,
274 - quant_config=quant_config,
275 - multimodal_config=multimodal_config,
276 - require_post_norm=False,
277 - prefix=maybe_prefix(prefix, "vision_tower"),
278 - )
279 - self.image_newline = nn.Parameter(torch.empty(config.text_config.hidden_size))
280 - self.multi_modal_projector = LlavaMultiModalProjector(
281 - vision_hidden_size=vision_hidden_size,
282 - text_hidden_size=config.text_config.hidden_size,
283 - projector_hidden_act=config.projector_hidden_act,
284 - multimodal_projector_bias=config.multimodal_projector_bias,
285 - )
271 + with self._mark_tower_model(vllm_config, "image"):
272 + self.vision_tower = init_vision_tower_for_llava(
273 + config,
274 + quant_config=quant_config,
275 + multimodal_config=multimodal_config,
276 + require_post_norm=False,
277 + prefix=maybe_prefix(prefix, "vision_tower"),
278 + )
279 + self.image_newline = nn.Parameter(
280 + torch.empty(config.text_config.hidden_size)
281 + )
282 + self.multi_modal_projector = LlavaMultiModalProjector(
283 + vision_hidden_size=vision_hidden_size,
284 + text_hidden_size=config.text_config.hidden_size,
285 + projector_hidden_act=config.projector_hidden_act,
286 + multimodal_projector_bias=config.multimodal_projector_bias,
287 + )
288
287 - self.language_model = init_vllm_registered_model(
288 - vllm_config=vllm_config,
289 - hf_config=config.text_config,
290 - prefix=maybe_prefix(prefix, "language_model"),
291 - )
289 + with self._mark_language_model(vllm_config):
290 + self.language_model = init_vllm_registered_model(
291 + vllm_config=vllm_config,
292 + hf_config=config.text_config,
293 + prefix=maybe_prefix(prefix, "language_model"),
294 + )
295
296 self.make_empty_intermediate_tensors = (
297 self.language_model.make_empty_intermediate_tensors
@@ -427,8 +430,6 @@ class LlavaNextForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsP
430 self,
431 inputs: LlavaNextImagePixelInputs,
432 ) -> torch.Tensor | tuple[torch.Tensor, ...]:
430 - assert self.vision_tower is not None
431 -
433 pixel_values = inputs["pixel_values"]
434
435 if isinstance(pixel_values, torch.Tensor):
@@ -480,9 +481,6 @@ class LlavaNextForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsP
481 for i, patch_features_batch in enumerate(patch_embeddings)
482 ]
483
483 - def get_language_model(self) -> torch.nn.Module:
484 - return self.language_model
485 -
484 def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:
485 image_input = self._parse_and_validate_image_input(**kwargs)
486 if image_input is None:
vllm/model_executor/models/llava_next_video.py
+24 -28
@@ -312,12 +312,10 @@ class LlavaNextVideoForConditionalGeneration(nn.Module, SupportsMultiModal, Supp
312
313 @classmethod
314 def get_placeholder_str(cls, modality: str, i: int) -> str | None:
315 - if modality.startswith("image"):
316 - return "<image>"
315 if modality.startswith("video"):
316 return "<video>"
317
320 - raise ValueError("Only image or video modality is supported")
318 + raise ValueError("Only video modality is supported")
319
320 def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
321 super().__init__()
@@ -329,26 +327,29 @@ class LlavaNextVideoForConditionalGeneration(nn.Module, SupportsMultiModal, Supp
327 self.config = config
328 self.multimodal_config = multimodal_config
329
332 - # Initialize the vision tower only up to the required feature layer
333 - self.vision_tower = init_vision_tower_for_llava(
334 - config,
335 - quant_config=quant_config,
336 - multimodal_config=multimodal_config,
337 - require_post_norm=False,
338 - prefix=maybe_prefix(prefix, "vision_tower"),
339 - )
340 - self.vision_resampler = LlavaNextVideoPooler(config)
341 - self.multi_modal_projector = LlavaNextMultiModalProjector(
342 - vision_hidden_size=config.vision_config.hidden_size,
343 - text_hidden_size=config.text_config.hidden_size,
344 - projector_hidden_act=config.projector_hidden_act,
345 - multimodal_projector_bias=config.multimodal_projector_bias,
346 - )
347 - self.language_model = init_vllm_registered_model(
348 - vllm_config=vllm_config,
349 - hf_config=config.text_config,
350 - prefix=maybe_prefix(prefix, "language_model"),
351 - )
330 + with self._mark_tower_model(vllm_config, "video"):
331 + # Initialize the vision tower only up to the required feature layer
332 + self.vision_tower = init_vision_tower_for_llava(
333 + config,
334 + quant_config=quant_config,
335 + multimodal_config=multimodal_config,
336 + require_post_norm=False,
337 + prefix=maybe_prefix(prefix, "vision_tower"),
338 + )
339 + self.vision_resampler = LlavaNextVideoPooler(config)
340 + self.multi_modal_projector = LlavaNextMultiModalProjector(
341 + vision_hidden_size=config.vision_config.hidden_size,
342 + text_hidden_size=config.text_config.hidden_size,
343 + projector_hidden_act=config.projector_hidden_act,
344 + multimodal_projector_bias=config.multimodal_projector_bias,
345 + )
346 +
347 + with self._mark_language_model(vllm_config):
348 + self.language_model = init_vllm_registered_model(
349 + vllm_config=vllm_config,
350 + hf_config=config.text_config,
351 + prefix=maybe_prefix(prefix, "language_model"),
352 + )
353
354 self.make_empty_intermediate_tensors = (
355 self.language_model.model.make_empty_intermediate_tensors
@@ -395,8 +396,6 @@ class LlavaNextVideoForConditionalGeneration(nn.Module, SupportsMultiModal, Supp
396 return image_features
397
398 def _process_video_pixels(self, inputs: LlavaNextVideoPixelInputs):
398 - assert self.vision_tower is not None
399 -
399 video_pixels = inputs["pixel_values_videos"]
400
401 if isinstance(video_pixels, torch.Tensor):
@@ -419,9 +418,6 @@ class LlavaNextVideoForConditionalGeneration(nn.Module, SupportsMultiModal, Supp
418
419 return [e.flatten(0, 1) for e in embeds]
420
422 - def get_language_model(self) -> torch.nn.Module:
423 - return self.language_model
424 -
421 def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:
422 video_input = self._parse_and_validate_video_input(**kwargs)
423 if video_input is None:
vllm/model_executor/models/llava_onevision.py
+20 -22
@@ -508,21 +508,26 @@ class LlavaOnevisionForConditionalGeneration(nn.Module, SupportsMultiModal, Supp
508 self.config = config
509 self.multimodal_config = multimodal_config
510
511 - # Initialize the vision tower only up to the required feature layer
512 - self.vision_tower = init_vision_tower_for_llava(
513 - config,
514 - quant_config=quant_config,
515 - multimodal_config=multimodal_config,
516 - require_post_norm=False,
517 - prefix=maybe_prefix(prefix, "vision_tower"),
518 - )
519 - self.multi_modal_projector = LlavaOnevisionMultiModalProjector(config)
520 - self.language_model = init_vllm_registered_model(
521 - vllm_config=vllm_config,
522 - hf_config=config.text_config,
523 - prefix=maybe_prefix(prefix, "language_model"),
524 - )
525 - self.image_newline = nn.Parameter(torch.empty(config.text_config.hidden_size))
511 + with self._mark_tower_model(vllm_config, {"image", "video"}):
512 + # Initialize the vision tower only up to the required feature layer
513 + self.vision_tower = init_vision_tower_for_llava(
514 + config,
515 + quant_config=quant_config,
516 + multimodal_config=multimodal_config,
517 + require_post_norm=False,
518 + prefix=maybe_prefix(prefix, "vision_tower"),
519 + )
520 + self.image_newline = nn.Parameter(
521 + torch.empty(config.text_config.hidden_size)
522 + )
523 + self.multi_modal_projector = LlavaOnevisionMultiModalProjector(config)
524 +
525 + with self._mark_language_model(vllm_config):
526 + self.language_model = init_vllm_registered_model(
527 + vllm_config=vllm_config,
528 + hf_config=config.text_config,
529 + prefix=maybe_prefix(prefix, "language_model"),
530 + )
531
532 self.make_empty_intermediate_tensors = (
533 self.language_model.model.make_empty_intermediate_tensors
@@ -726,8 +731,6 @@ class LlavaOnevisionForConditionalGeneration(nn.Module, SupportsMultiModal, Supp
731 self,
732 inputs: LlavaOnevisionImagePixelInputs,
733 ) -> torch.Tensor | list[torch.Tensor]:
729 - assert self.vision_tower is not None
730 -
734 pixel_values = inputs["pixel_values"]
735
736 if isinstance(pixel_values, torch.Tensor):
@@ -801,8 +804,6 @@ class LlavaOnevisionForConditionalGeneration(nn.Module, SupportsMultiModal, Supp
804 return video_features
805
806 def _process_video_pixels(self, inputs: LlavaOnevisionVideoPixelInputs):
804 - assert self.vision_tower is not None
805 -
807 video_pixels = inputs["pixel_values_videos"]
808
809 if isinstance(video_pixels, torch.Tensor):
@@ -862,9 +863,6 @@ class LlavaOnevisionForConditionalGeneration(nn.Module, SupportsMultiModal, Supp
863 image_feature = image_feature.view(batch_frames, -1, dim)
864 return image_feature
865
865 - def get_language_model(self) -> torch.nn.Module:
866 - return self.language_model
867 -
866 def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:
867 mm_input_by_modality = self._parse_and_validate_multimodal_inputs(**kwargs)
868 if not mm_input_by_modality: