[Mypy] Fix mypy typing for N/O models (#54142)

Signed-off-by: Taneem Ibrahim <taneem.ibrahim@gmail.com>
This commit is contained in:
Taneem Ibrahim
2026-09-18 12:50:01 +01:00
committed by GitHub
parent e562003adb
commit 2b8a5495a8
19 changed files with 321 additions and 143 deletions
-2
View File
@@ -113,8 +113,6 @@ SEPARATE_GROUPS = [
EXCLUDE = [
r"vllm/model_executor/models/[kK]",
r"vllm/model_executor/models/[nN]",
r"vllm/model_executor/models/[oO]",
r"vllm/model_executor/models/[qQ]",
r"vllm/model_executor/models/[uU]",
r"vllm/model_executor/models/[vV]",
@@ -222,12 +222,18 @@ class MambaStateShapeCalculator:
@classmethod
def append_replayssm_ring(
cls,
base_shapes: tuple[tuple[int, ...], ...],
base_shapes: tuple[tuple[int, int], tuple[int, int, int]],
n_groups: int,
tp_world_size: int,
logical_window: int,
backend: MambaBackendEnum,
) -> tuple[tuple[int, ...], ...]:
) -> tuple[
tuple[int, int],
tuple[int, int, int],
tuple[int, int, int],
tuple[int, int],
tuple[int, int, int],
]:
"""Append the physical ReplaySSM ring shapes.
``base_shapes[1]`` is ``(nheads // tp, head_dim, state_size)``;
+14 -5
View File
@@ -26,7 +26,10 @@ from vllm.transformers_utils.configs.ovis import AIMv2Config
class AIMv2SwiGLUFFN(nn.Module):
def __init__(
self, config: AIMv2Config, quant_config: QuantizationConfig, prefix: str
self,
config: AIMv2Config,
quant_config: QuantizationConfig | None,
prefix: str,
):
super().__init__()
hidden_features = config.intermediate_size
@@ -91,7 +94,10 @@ class AIMv2ViTPreprocessor(nn.Module):
class AIMv2Attention(nn.Module):
def __init__(
self, config: AIMv2Config, quant_config: QuantizationConfig, prefix: str
self,
config: AIMv2Config,
quant_config: QuantizationConfig | None,
prefix: str,
):
super().__init__()
self.config = config
@@ -144,7 +150,10 @@ class AIMv2Attention(nn.Module):
class AIMv2Block(nn.Module):
def __init__(
self, config: AIMv2Config, quant_config: QuantizationConfig, prefix: str
self,
config: AIMv2Config,
quant_config: QuantizationConfig | None,
prefix: str,
):
super().__init__()
self.attn = AIMv2Attention(
@@ -166,7 +175,7 @@ class AIMv2Transformer(nn.Module):
def __init__(
self,
config: AIMv2Config,
quant_config: QuantizationConfig,
quant_config: QuantizationConfig | None,
*,
require_post_norm: bool | None = None,
prefix: str = "",
@@ -206,7 +215,7 @@ class AIMv2Model(torch.nn.Module):
def __init__(
self,
config: AIMv2Config,
quant_config: QuantizationConfig,
quant_config: QuantizationConfig | None,
*,
require_post_norm: bool | None = None,
prefix: str = "",
+7
View File
@@ -91,6 +91,13 @@ MambaStateShapes: TypeAlias = (
tuple[int, int, int],
tuple[int, int, int],
]
| tuple[
tuple[int, int],
tuple[int, int, int],
tuple[int, int, int],
tuple[int, int],
tuple[int, int, int],
]
)
+85 -27
View File
@@ -12,7 +12,7 @@ import warnings
from collections.abc import Iterable, Mapping, Sequence
from functools import cached_property
from io import BytesIO
from typing import Annotated, Literal, TypeAlias
from typing import Annotated, Literal, TypeAlias, TypedDict, cast
import torch
import torch.nn as nn
@@ -20,7 +20,11 @@ from transformers import BatchFeature, PretrainedConfig
from vllm import envs
from vllm.config import VllmConfig
from vllm.config.multimodal import BaseDummyOptions, VideoDummyOptions
from vllm.config.multimodal import (
AudioDummyOptions,
BaseDummyOptions,
VideoDummyOptions,
)
from vllm.inputs import MultiModalDataDict, MultiModalInput
from vllm.logger import init_logger
from vllm.lora.layers.base import BaseLayerWithLoRA
@@ -114,7 +118,7 @@ class NanoNemotronVLAudioFeatureInputs(TensorSchema):
type: Literal["audio_features"] = "audio_features"
input_audio_features: Annotated[torch.Tensor, TensorShape("c", "t", "f")]
feature_attention_mask: Annotated[torch.Tensor, TensorShape("c", "t")]
audio_num_clips: list[int]
audio_num_clips: list[int] | torch.Tensor
class NanoNemotronVLImagePixelInputs(TensorSchema):
@@ -195,6 +199,12 @@ NanoNemotronVLVideoInputs: TypeAlias = (
)
class NanoNemotronVLMultiModalInputs(TypedDict, total=False):
images: NanoNemotronVLImageInputs | None
videos: NanoNemotronVLVideoInputs | None
audios: NanoNemotronVLAudioFeatureInputs | None
class NanoNemotronVLProcessingInfo(BaseProcessingInfo):
def get_hf_processor(self, **kwargs: object) -> NanoNemotronVLProcessor:
# Not accepted by NanoNemotronVLProcessor.__init__
@@ -465,8 +475,11 @@ class NanoNemotronVLMultiModalProcessor(
if isinstance(images, ImageEmbeddingItems):
feature_size = images.get_feature_size(item_idx)
elif self.info.is_dynamic_tiler:
feature_size = out_mm_data["num_tokens_per_image"][item_idx]
num_tokens_per_image = out_mm_data["num_tokens_per_image"]
assert isinstance(num_tokens_per_image, list)
feature_size = num_tokens_per_image[item_idx]
else:
assert isinstance(images, ImageProcessorItems)
image_size = images.get_image_size(item_idx)
max_num_tiles = hf_processor.max_num_tiles
feature_size = hf_processor.get_num_image_tokens(
@@ -530,10 +543,12 @@ class NanoNemotronVLMultiModalProcessor(
assert isinstance(num_patches, int)
T = hf_processor.video_temporal_patch_size
num_tubelets: int | None
if T > 1 and num_patches is not None:
num_tubelets = math.ceil(num_patches / T)
else:
num_tubelets = num_patches
assert num_tubelets is not None
video_pruning_rate = self.info.ctx.get_mm_config().video_pruning_rate
if video_pruning_rate is not None and video_pruning_rate > 0.0:
@@ -635,6 +650,7 @@ class NanoNemotronVLMultiModalProcessor(
audio_items: list[AudioItem] = []
has_audio: list[bool] = []
for idx, metadata in enumerate(metadata_list):
assert metadata is not None
video_bytes = metadata.get("original_video_bytes")
if video_bytes is None or len(video_bytes) == 0:
raise ValueError(
@@ -661,21 +677,25 @@ class NanoNemotronVLMultiModalProcessor(
# Create a new VideoProcessorItems with metadata that does not contain
# the large video bytes, to avoid modifying the input `mm_items`.
new_metadata_list = [
{k: v for k, v in meta.items() if k != "original_video_bytes"}
for meta in metadata_list
]
new_metadata_list = metadata_list.copy()
for idx, metadata in enumerate(metadata_list):
assert metadata is not None
new_metadata_list[idx] = {
key: value
for key, value in metadata.items()
if key != "original_video_bytes"
}
new_videos = VideoProcessorItems(data=videos.data, metadata=new_metadata_list)
audio_parsed = {}
audio_parsed = MultiModalDataItems()
if audio_items:
audio_parsed = self.data_parser.parse_mm_data({"audio": audio_items})
# Create a new MultiModalDataItems with the new video and audio items.
new_mm_items_dict = {**mm_items, **audio_parsed, "video": new_videos}
mm_items = MultiModalDataItems(new_mm_items_dict)
new_mm_items = MultiModalDataItems(new_mm_items_dict)
return mm_items, audio_items, has_audio
return new_mm_items, audio_items, has_audio
def apply(
self,
@@ -800,7 +820,7 @@ class NanoNemotronVLDummyInputsBuilder(
)
videos = [v.copy() for v in videos]
video_items = []
video_items: list[VideoItem] = []
for video in videos:
video_num_frames = video.shape[0]
video_metadata = {
@@ -869,6 +889,9 @@ class NanoNemotronVLDummyInputsBuilder(
target_num_frames = num_frames
num_videos = mm_counts.get("video", 0)
video_overrides = mm_options.get("video")
assert video_overrides is None or isinstance(
video_overrides, VideoDummyOptions
)
dummy_video = {
"video": self._get_dummy_videos(
width=video_width,
@@ -884,6 +907,9 @@ class NanoNemotronVLDummyInputsBuilder(
if sound_config := self.info.sound_config:
num_audios = mm_counts.get("audio", 0)
audio_overrides = mm_options.get("audio") if mm_options else None
assert audio_overrides is None or isinstance(
audio_overrides, AudioDummyOptions
)
tokens_per_audio = max(1, seq_len // max(num_audios, 1))
max_audio_num_samples = MAX_AUDIO_LEN_S * sound_config.sampling_rate
calculated_max_audio_num_samples = ParakeetExtractor.audio_length(
@@ -946,6 +972,7 @@ class NemotronH_Nano_VL_V2(
model_config = vllm_config.model_config
config = model_config.hf_config
multimodal_config = model_config.multimodal_config
assert multimodal_config is not None
image_size = config.force_image_size
patch_size = config.patch_size
self.patch_size = patch_size
@@ -1138,17 +1165,26 @@ class NemotronH_Nano_VL_V2(
return None
if self.dynamic_resolution:
assert isinstance(pixel_values_flat, (list, torch.Tensor))
pixel_values_flat = DynamicResolutionImageTiler.stack(
pixel_values_flat, self.patch_size
)
imgs_sizes = kwargs.pop("imgs_sizes")
num_tokens_per_image = kwargs.pop("num_tokens_per_image")
assert isinstance(imgs_sizes, list)
assert isinstance(num_tokens_per_image, list)
return NanoNemotronVLImagePixelInputsDynamic(
pixel_values_flat=pixel_values_flat, **kwargs
pixel_values_flat=pixel_values_flat,
imgs_sizes=imgs_sizes,
num_tokens_per_image=num_tokens_per_image,
)
else:
assert isinstance(pixel_values_flat, torch.Tensor)
image_num_patches = kwargs.pop("image_num_patches")
assert isinstance(image_num_patches, torch.Tensor)
return NanoNemotronVLImagePixelInputs(
pixel_values_flat=pixel_values_flat,
num_patches=kwargs.pop("image_num_patches"),
**kwargs,
num_patches=image_num_patches,
)
def _process_image_input_dynamic(
@@ -1166,7 +1202,8 @@ class NemotronH_Nano_VL_V2(
return image_embeds.split(num_tokens_per_image)
def _process_image_input(
self, image_input: NanoNemotronVLImagePixelInputs
self,
image_input: NanoNemotronVLImagePixelInputs | NanoNemotronVLVideoPixelInputs,
) -> tuple[torch.Tensor, ...]:
image_embeds = self.extract_feature(image_input["pixel_values_flat"])
num_patches = image_input["num_patches"]
@@ -1362,7 +1399,8 @@ class NemotronH_Nano_VL_V2(
# of them than max_num_batched_tokens.
# Embed them with the base weights -
# the LoRA delta is undefined for tokens with no mapping entry.
embed_tokens = self.get_language_model().model.embed_tokens
language_model = cast(NemotronHForCausalLM, self.get_language_model())
embed_tokens = language_model.model.embed_tokens
if isinstance(embed_tokens, BaseLayerWithLoRA):
embed_tokens = embed_tokens.base_layer
text_embeddings = embed_tokens(repl_token_ids)
@@ -1388,26 +1426,29 @@ class NemotronH_Nano_VL_V2(
return None
if video_embeds is not None:
assert isinstance(video_embeds, (torch.Tensor, list))
return NanoNemotronVLVideoEmbeddingInputs(
type="video_embeds",
data=video_embeds,
)
if pixel_values_flat_video is not None:
if torch.is_tensor(frames_indices):
if isinstance(frames_indices, torch.Tensor):
frames_indices = frames_indices.flatten()
else:
assert isinstance(frames_indices, (list, tuple))
frames_indices = torch.cat([f.flatten() for f in frames_indices], dim=0)
if torch.is_tensor(frame_duration_ms):
if isinstance(frame_duration_ms, torch.Tensor):
frame_duration_ms = frame_duration_ms.flatten()
else:
assert isinstance(frame_duration_ms, (list, tuple))
frame_duration_ms = torch.cat(
[f.flatten() for f in frame_duration_ms], dim=0
)
if (
torch.is_tensor(pixel_values_flat_video)
isinstance(pixel_values_flat_video, torch.Tensor)
and pixel_values_flat_video.ndim == 5
):
# batched._reduce_data stacked same-shape videos into
@@ -1415,9 +1456,11 @@ class NemotronH_Nano_VL_V2(
# same-H,W cat path below handles it uniformly.
pixel_values_flat_video = list(pixel_values_flat_video)
if not torch.is_tensor(pixel_values_flat_video):
if not isinstance(pixel_values_flat_video, torch.Tensor):
assert isinstance(pixel_values_flat_video, (list, tuple))
pixel_values_flat_video = torch.cat(pixel_values_flat_video, dim=0)
assert isinstance(video_num_patches, torch.Tensor)
expected_h = pixel_values_flat_video.shape[-2]
expected_w = pixel_values_flat_video.shape[-1]
num_frames = video_num_patches[0].item()
@@ -1434,8 +1477,10 @@ class NemotronH_Nano_VL_V2(
raise AssertionError("This line should be unreachable.")
def _parse_and_validate_multimodal_inputs(self, **kwargs: object) -> dict:
modalities = {}
def _parse_and_validate_multimodal_inputs(
self, **kwargs: object
) -> NanoNemotronVLMultiModalInputs:
modalities: NanoNemotronVLMultiModalInputs = {}
# Preserve the order of modalities if there are multiple of them
# from the order of kwargs.
for input_key in kwargs:
@@ -1458,8 +1503,17 @@ class NemotronH_Nano_VL_V2(
)
and "audios" not in modalities
):
input_audio_features = kwargs.get("input_audio_features")
feature_attention_mask = kwargs.get("feature_attention_mask")
audio_num_clips = kwargs.get("audio_num_clips")
assert isinstance(input_audio_features, torch.Tensor)
assert isinstance(feature_attention_mask, torch.Tensor)
assert isinstance(audio_num_clips, (list, torch.Tensor))
modalities["audios"] = NanoNemotronVLAudioFeatureInputs(
**kwargs, validate=False
input_audio_features=input_audio_features,
feature_attention_mask=feature_attention_mask,
audio_num_clips=audio_num_clips,
validate=False,
)
return modalities
@@ -1479,23 +1533,26 @@ class NemotronH_Nano_VL_V2(
for modality in modalities:
if modality == "images":
image_input = modalities["images"]
if image_input["type"] == "image_embeds":
assert image_input is not None
if isinstance(image_input, NanoNemotronVLImageEmbeddingInputs):
image_embeddings = image_input["data"]
elif self.dynamic_resolution:
assert image_input["type"] == "pixel_values_dynamic"
elif isinstance(image_input, NanoNemotronVLImagePixelInputsDynamic):
image_embeddings = self._process_image_input_dynamic(image_input)
else:
image_embeddings = self._process_image_input(image_input)
multimodal_embeddings += tuple(image_embeddings)
if modality == "videos":
video_input = modalities["videos"]
assert video_input is not None
if video_input["type"] == "video_embeds":
video_embeddings = video_input["data"]
else:
assert isinstance(video_input, NanoNemotronVLVideoPixelInputs)
video_embeddings = self._process_video_input(video_input)
multimodal_embeddings += tuple(video_embeddings)
if modality == "audios":
audio_input = modalities["audios"]
assert audio_input is not None
audio_embeddings = self._process_audio_input(audio_input)
multimodal_embeddings += tuple(audio_embeddings)
@@ -1538,6 +1595,7 @@ class NemotronH_Nano_VL_V2(
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
mm_config = self.model_config.multimodal_config
assert mm_config is not None
load_multimodal_weights = not all(
mm_config.get_limit_per_prompt(modality) == 0
for modality in ("image", "video", "audio")
+7 -2
View File
@@ -68,6 +68,7 @@ from vllm.model_executor.models.interfaces import (
EagleModelMixin,
HasInnerState,
IsHybrid,
MambaStateShapes,
MixtureOfExperts,
SupportsEagle,
SupportsEagle3,
@@ -138,6 +139,7 @@ class NemotronHMoE(nn.Module):
prefix: str = "",
):
super().__init__()
assert parallel_config is not None
self.tp_size = get_tensor_model_parallel_world_size()
self.routed_scaling_factor = config.routed_scaling_factor
@@ -188,6 +190,8 @@ class NemotronHMoE(nn.Module):
prefix=f"{prefix}.shared_experts",
)
self.fc1_latent_proj: ReplicatedLinear | None
self.fc2_latent_proj: ReplicatedLinear | None
if self.use_latent_moe:
self.fc1_latent_proj = ReplicatedLinear(
input_size=config.hidden_size,
@@ -722,7 +726,7 @@ class NemotronHForCausalLM(
SupportsReplaySSM,
):
# Relevant only if self.has_moe is True
is_non_gated_moe: bool = True
is_non_gated_moe = True
hf_to_vllm_mapper = WeightsMapper(
orig_to_new_prefix={"backbone": "model", "mtp": None},
@@ -773,7 +777,7 @@ class NemotronHForCausalLM(
def get_mamba_state_shape_from_config(
cls,
vllm_config: "VllmConfig",
) -> tuple[tuple[int, ...], ...]:
) -> MambaStateShapes:
"""Calculate shapes for Mamba's convolutional and state caches.
Args:
@@ -858,6 +862,7 @@ class NemotronHForCausalLM(
self.moe_layers.append(layer.mixer.experts)
self.num_moe_layers = len(self.moe_layers)
assert example_moe is not None
self.num_logical_experts = example_moe.n_logical_experts
self.num_physical_experts = example_moe.n_physical_experts
self.num_local_physical_experts = example_moe.n_local_physical_experts # noqa: E501
+16 -8
View File
@@ -92,10 +92,12 @@ class NemotronHMTPAttentionDecoderLayer(NemotronHAttentionDecoderLayer):
def forward(
self,
inputs_embeds: torch.Tensor,
positions: torch.Tensor,
hidden_states: torch.Tensor,
residual: torch.Tensor | None = None,
*,
inputs_embeds: torch.Tensor | None = None,
**kwargs: object,
) -> tuple[torch.Tensor, torch.Tensor | None]:
# Start projections (Fusion)
if self.has_start_projections:
@@ -179,10 +181,11 @@ class NemotronHMTPMoEDecoderLayer(NemotronHMoEDecoderLayer):
def forward(
self,
inputs_embeds: torch.Tensor,
positions: torch.Tensor,
hidden_states: torch.Tensor,
residual: torch.Tensor | None = None,
*,
inputs_embeds: torch.Tensor | None = None,
**kwargs: object,
) -> tuple[torch.Tensor, torch.Tensor | None]:
# Start projections (Fusion)
if self.has_start_projections:
@@ -305,13 +308,14 @@ class NemotronHMultiTokenPredictor(nn.Module):
def forward(
self,
input_ids: torch.Tensor,
input_ids: torch.Tensor | None,
positions: torch.Tensor,
hidden_states: torch.Tensor,
intermediate_tensors: IntermediateTensors | None = None,
inputs_embeds: torch.Tensor | None = None,
) -> torch.Tensor | IntermediateTensors:
if inputs_embeds is None:
assert input_ids is not None
inputs_embeds = self.get_input_embeddings(input_ids)
residual = None
@@ -386,6 +390,7 @@ class NemotronHMTP(nn.Module, SupportsPP, SupportsQuant):
@staticmethod
def _find_quant_config(*args, **kwargs) -> QuantizationConfig | None:
vllm_config = kwargs.get("vllm_config")
assert isinstance(vllm_config, VllmConfig)
return get_draft_quant_config(vllm_config)
def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
@@ -393,14 +398,15 @@ class NemotronHMTP(nn.Module, SupportsPP, SupportsQuant):
def forward(
self,
input_ids: torch.Tensor,
input_ids: torch.Tensor | None,
positions: torch.Tensor,
hidden_states: torch.Tensor,
hidden_states: torch.Tensor | None = None,
intermediate_tensors: IntermediateTensors | None = None,
inputs_embeds: torch.Tensor | None = None,
**kwargs: object,
) -> torch.Tensor:
"""Forward - applies attention-based MTP."""
assert hidden_states is not None
hidden_states = self.model(
input_ids,
positions,
@@ -473,9 +479,10 @@ class NemotronHMTP(nn.Module, SupportsPP, SupportsQuant):
name = name.replace("backbone.", "model.")
if "scale" in name or "zero_point" in name:
name = maybe_remap_kv_scale_name(name, params_dict)
if name is None:
remapped_name = maybe_remap_kv_scale_name(name, params_dict)
if remapped_name is None:
continue
name = remapped_name
# Handle stacked parameters (qkv_proj) for attention layers
is_stacked = False
@@ -547,6 +554,7 @@ class NemotronHMTP(nn.Module, SupportsPP, SupportsQuant):
param = params_dict[name]
weight_loader = getattr(param, "weight_loader", default_weight_loader)
assert weight_loader is not None
weight_loader(param, loaded_weight)
loaded_params.add(name)
+14 -5
View File
@@ -138,7 +138,7 @@ class LlamaNemotronVLChatModel(nn.Module, SupportsMultiModal, SupportsPP, Suppor
prefix=maybe_prefix(prefix, "language_model"),
)
self.img_context_token_id = None
self.img_context_token_id: int | None = None
self.visual_token_mask = None
self.make_empty_intermediate_tensors = (
@@ -146,7 +146,7 @@ class LlamaNemotronVLChatModel(nn.Module, SupportsMultiModal, SupportsPP, Suppor
)
def _patch_quant_config(
self, config: PretrainedConfig, quant_config: QuantizationConfig
self, config: PretrainedConfig, quant_config: QuantizationConfig | None
):
# the awq models from OpenGVLab missing `modules_to_not_convert`
# patch the quant_config to add `modules_to_not_convert` back
@@ -518,12 +518,21 @@ class LlamaNemotronVLForEmbedding(LlamaNemotronVLChatModel, VllmModelForPooling)
use_head=False,
)
def _init_mlp1(self, config: PretrainedConfig) -> nn.Module:
def _init_mlp1(
self,
config: PretrainedConfig,
vit_hidden_size: int | None = None,
vision_projection_hidden_size: int | None = None,
) -> nn.Module:
"""Override to use different MLP structure for embedding model."""
if vit_hidden_size is None:
vit_hidden_size = config.vision_config.hidden_size
if vision_projection_hidden_size is None:
vision_projection_hidden_size = config.get_text_config().hidden_size
return super()._init_mlp1(
config,
vit_hidden_size=config.vision_config.hidden_size,
vision_projection_hidden_size=config.get_text_config().hidden_size,
vit_hidden_size=vit_hidden_size,
vision_projection_hidden_size=vision_projection_hidden_size,
)
def _call_vision_model(self, pixel_values: torch.Tensor) -> torch.Tensor:
+1
View File
@@ -128,6 +128,7 @@ class NVLMMultiModalProcessor(BaseInternVLMultiModalProcessor[NVLMProcessingInfo
if isinstance(images, ImageEmbeddingItems):
feature_size = images.get_feature_size(item_idx)
else:
assert isinstance(images, ImageProcessorItems)
image_size = images.get_image_size(item_idx)
feature_size = self.info.get_num_image_tokens(
image_width=image_size.width,
@@ -143,6 +143,7 @@ class OlmoHybridAttention(nn.Module):
rope_parameters["rope_theta"] is not None
)
self.rotary_emb: nn.Module | None
if self._use_rope:
self.rotary_emb = get_rope(
self.head_dim,
@@ -183,6 +184,7 @@ class OlmoHybridAttention(nn.Module):
q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
q, k = self._apply_qk_norm(q, k)
if self._use_rope:
assert self.rotary_emb is not None
q, k = self.rotary_emb(positions, q, k)
attn_output = self.attn(q, k, v)
output, _ = self.o_proj(attn_output)
+10 -2
View File
@@ -7,6 +7,7 @@
"""Inference-only OpenCUA-7B model compatible with HuggingFace weights."""
import typing
from collections.abc import Mapping, Sequence
from typing import Any
@@ -27,6 +28,7 @@ from vllm.multimodal.inputs import (
)
from vllm.multimodal.parse import MultiModalDataItems
from vllm.multimodal.processing import (
BaseDummyInputsBuilder,
BaseMultiModalProcessor,
PromptReplacement,
PromptUpdate,
@@ -178,7 +180,13 @@ class OpenCUAMultiModalProcessor(BaseMultiModalProcessor[OpenCUAProcessingInfo])
]
class OpenCUADummyInputsBuilder(Qwen2VLDummyInputsBuilder):
if typing.TYPE_CHECKING:
_OpenCUADummyInputsBuilderBase = BaseDummyInputsBuilder[OpenCUAProcessingInfo]
else:
_OpenCUADummyInputsBuilderBase = Qwen2VLDummyInputsBuilder
class OpenCUADummyInputsBuilder(_OpenCUADummyInputsBuilderBase):
def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
num_images = mm_counts.get("image", 0)
@@ -220,7 +228,7 @@ class OpenCUAForConditionalGeneration(Qwen2_5_VLForConditionalGeneration):
nn.Module.__init__(self)
config = vllm_config.model_config.hf_config
quant_config = vllm_config.quant_config
multimodal_config = vllm_config.model_config.multimodal_config
multimodal_config = vllm_config.model_config.get_multimodal_config()
self.use_data_parallel = multimodal_config.mm_encoder_tp_mode == "data"
self.config = config
+5 -1
View File
@@ -175,6 +175,7 @@ class OpenPanguMoE(nn.Module):
self.n_physical_experts = self.n_logical_experts + self.n_redundant_experts
self.n_local_physical_experts = self.n_physical_experts // self.ep_size
self.shared_experts: OpenPanguMLP | None
if config.n_shared_experts is not None:
intermediate_size = config.moe_intermediate_size * config.n_shared_experts
self.shared_experts = OpenPanguMLP(
@@ -283,7 +284,7 @@ class OpenPanguMLAAttention(nn.Module):
)
self.q_a_layernorm = RMSNorm(self.q_lora_rank, eps=config.rms_norm_eps)
self.q_b_proj = ColumnParallelLinear(
q_lora_rank,
self.q_lora_rank,
self.num_heads * self.qk_head_dim,
bias=False,
quant_config=quant_config,
@@ -568,6 +569,9 @@ class OpenPanguSinkAttention(nn.Module):
self.qk_nope_dim = getattr(config, "qk_nope_dim", None)
self.qk_rope_dim = getattr(config, "qk_rope_dim", None)
self.v_channels = getattr(config, "v_channels", None)
assert self.qk_nope_dim is not None
assert self.qk_rope_dim is not None
assert self.v_channels is not None
self.head_dim = self.qk_rope_dim + self.qk_nope_dim
self.q_size = self.num_heads * self.head_dim
self.k_size = self.num_kv_heads * self.head_dim
+12 -5
View File
@@ -52,7 +52,9 @@ class OpenPanguMultiTokenPredictorLayer(DeepSeekMultiTokenPredictorLayer):
def __init__(self, vllm_config: VllmConfig, prefix: str) -> None:
nn.Module.__init__(self)
config = vllm_config.speculative_config.draft_model_config.hf_config
speculative_config = vllm_config.speculative_config
assert speculative_config is not None
config = speculative_config.draft_model_config.hf_config
self.config = config
quant_config = vllm_config.quant_config
@@ -200,10 +202,15 @@ class OpenPanguMTP(nn.Module):
break
else:
for mapping in expert_params_mapping:
param_name, weight_name, expert_id, shard_id = mapping
if weight_name not in name:
(
expert_param_name,
expert_weight_name,
expert_id,
expert_shard_id,
) = mapping
if expert_weight_name not in name:
continue
name = name.replace(weight_name, param_name)
name = name.replace(expert_weight_name, expert_param_name)
param = params_dict[name]
weight_loader = param.weight_loader
@@ -211,7 +218,7 @@ class OpenPanguMTP(nn.Module):
param,
loaded_weight,
name,
shard_id=shard_id,
shard_id=expert_shard_id,
expert_id=expert_id,
)
break
+101 -74
View File
@@ -22,7 +22,7 @@
from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence
from functools import lru_cache, partial
from typing import Annotated, Literal, Optional
from typing import Annotated, Literal, Optional, TypeAlias
import torch
import torch.nn as nn
@@ -627,7 +627,7 @@ class OpenPanguVisionTransformer(nn.Module):
pad[-(dim + 1) * 2 + 1] = pad_len
return F.pad(w, pad, mode="constant", value=0)
stacked_params_mapping = [
stacked_params_mapping: list[tuple[str, str, str | int]] = [
("attn.qkv.", "attn.q.", "q"),
("attn.qkv.", "attn.k.", "k"),
("attn.qkv.", "attn.v.", "v"),
@@ -677,15 +677,11 @@ class OpenPanguVLProcessingInfo(Qwen2_5_VLProcessingInfo):
def get_hf_config(self):
return self.ctx.model_config.hf_config
def get_hf_processor(
self,
*,
min_pixels: int | None = None,
max_pixels: int | None = None,
size: dict[str, int] | None = None,
fps: float | list[float] | None = None,
**kwargs: object,
):
def get_hf_processor(self, **kwargs: object):
kwargs.pop("min_pixels", None)
kwargs.pop("max_pixels", None)
kwargs.pop("size", None)
fps = kwargs.pop("fps", None)
if fps is not None:
kwargs["fps"] = fps
@@ -747,11 +743,19 @@ class OpenPanguVLVideoEmbeddingInputs(TensorSchema):
]
OpenPanguVLImageInputs: TypeAlias = (
OpenPanguVLImagePixelInputs | OpenPanguVLImageEmbeddingInputs
)
OpenPanguVLVideoInputs: TypeAlias = (
OpenPanguVLVideoPixelInputs | OpenPanguVLVideoEmbeddingInputs
)
class OpenPanguVLMultiModalProcessor(Qwen2_5_VLMultiModalProcessor):
def _get_prompt_updates(
self,
mm_items: MultiModalDataItems,
hf_processor_mm_kwargs: Mapping[str, any],
hf_processor_mm_kwargs: Mapping[str, object],
out_mm_kwargs: MultiModalKwargsItems,
) -> Sequence[PromptUpdate]:
hf_processor = self.info.get_hf_processor(**hf_processor_mm_kwargs)
@@ -878,10 +882,12 @@ class OpenPanguVLForConditionalGeneration(
self.do_rescale = image_processor.do_rescale
self.rescale_factor = image_processor.rescale_factor
self.do_normalize = image_processor.do_normalize
self.image_mean = tuple(image_processor.image_mean)
self.image_std = tuple(image_processor.image_std)
self.image_mean: tuple[float, ...] = tuple(image_processor.image_mean)
self.image_std: tuple[float, ...] = tuple(image_processor.image_std)
def _maybe_ignore_quant_config(self, quant_config: QuantizationConfig):
def _maybe_ignore_quant_config(
self, quant_config: QuantizationConfig | None
) -> QuantizationConfig | None:
if isinstance(quant_config, AutoGPTQConfig):
return None
return quant_config
@@ -904,7 +910,9 @@ class OpenPanguVLForConditionalGeneration(
else:
return torch.concat(mm_input)
def _parse_and_validate_image_input(self, **kwargs: object):
def _parse_and_validate_image_input(
self, **kwargs: object
) -> OpenPanguVLImageInputs | None:
pixel_values = kwargs.pop("pixel_values", None)
image_embeds = kwargs.pop("image_embeds", None)
image_grid_thw = kwargs.pop("image_grid_thw", None)
@@ -932,26 +940,27 @@ class OpenPanguVLForConditionalGeneration(
image_grid_thw=image_grid_thw,
)
if image_embeds is not None:
image_embeds = self._validate_and_reshape_mm_tensor(
image_embeds, "image embeds"
)
image_grid_thw = self._validate_and_reshape_mm_tensor(
image_grid_thw, "image grid_thw"
)
assert image_embeds is not None
image_embeds = self._validate_and_reshape_mm_tensor(
image_embeds, "image embeds"
)
image_grid_thw = self._validate_and_reshape_mm_tensor(
image_grid_thw, "image grid_thw"
)
if not isinstance(image_embeds, torch.Tensor):
raise ValueError(
"Incorrect type of image embeddings. "
f"Got type: {type(image_embeds)}"
)
return OpenPanguVLImageEmbeddingInputs(
type="image_embeds",
image_embeds=image_embeds,
image_grid_thw=image_grid_thw,
if not isinstance(image_embeds, torch.Tensor):
raise ValueError(
f"Incorrect type of image embeddings. Got type: {type(image_embeds)}"
)
return OpenPanguVLImageEmbeddingInputs(
type="image_embeds",
image_embeds=image_embeds,
image_grid_thw=image_grid_thw,
)
def _parse_and_validate_video_input(self, **kwargs: object):
def _parse_and_validate_video_input(
self, **kwargs: object
) -> OpenPanguVLVideoInputs | None:
pixel_values_videos = kwargs.pop("pixel_values_videos", None)
video_embeds = kwargs.pop("video_embeds", None)
video_grid_thw = kwargs.pop("video_grid_thw", None)
@@ -973,27 +982,30 @@ class OpenPanguVLForConditionalGeneration(
video_grid_thw=video_grid_thw,
)
if video_embeds is not None:
video_embeds = self._validate_and_reshape_mm_tensor(
video_embeds, "video embeds"
)
video_grid_thw = self._validate_and_reshape_mm_tensor(
video_grid_thw, "video grid_thw"
)
assert video_embeds is not None
video_embeds = self._validate_and_reshape_mm_tensor(
video_embeds, "video embeds"
)
video_grid_thw = self._validate_and_reshape_mm_tensor(
video_grid_thw, "video grid_thw"
)
if not isinstance(video_embeds, torch.Tensor):
raise ValueError(
"Incorrect type of video embeddings. "
f"Got type: {type(video_embeds)}"
)
return OpenPanguVLVideoEmbeddingInputs(
type="video_embeds",
video_embeds=video_embeds,
video_grid_thw=video_grid_thw,
if not isinstance(video_embeds, torch.Tensor):
raise ValueError(
f"Incorrect type of video embeddings. Got type: {type(video_embeds)}"
)
return OpenPanguVLVideoEmbeddingInputs(
type="video_embeds",
video_embeds=video_embeds,
video_grid_thw=video_grid_thw,
)
def _parse_and_validate_multimodal_inputs(self, **kwargs: object) -> dict:
mm_input_by_modality = {}
def _parse_and_validate_multimodal_inputs(
self, **kwargs: object
) -> dict[str, OpenPanguVLImageInputs | OpenPanguVLVideoInputs | None]:
mm_input_by_modality: dict[
str, OpenPanguVLImageInputs | OpenPanguVLVideoInputs | None
] = {}
for input_key in kwargs:
if (
input_key in ("pixel_values", "image_embeds")
@@ -1019,16 +1031,19 @@ class OpenPanguVLForConditionalGeneration(
multimodal_embeddings: tuple[torch.Tensor, ...] = ()
for modality in mm_input_by_modality:
multimodal_input = mm_input_by_modality[modality]
if modality == "image":
vision_embeddings = self._process_image_input(multimodal_input)
image_input = mm_input_by_modality["image"]
assert isinstance(image_input, OpenPanguVLImageInputs)
vision_embeddings = self._process_image_input(image_input)
multimodal_embeddings = (
multimodal_embeddings
if not vision_embeddings
else (multimodal_embeddings + vision_embeddings)
)
if modality == "video":
video_embeddings = self._process_video_input(multimodal_input)
video_input = mm_input_by_modality["video"]
assert isinstance(video_input, OpenPanguVLVideoInputs)
video_embeddings = self._process_video_input(video_input)
multimodal_embeddings = (
multimodal_embeddings
if not video_embeddings
@@ -1041,20 +1056,19 @@ class OpenPanguVLForConditionalGeneration(
input_ids: torch.Tensor,
multimodal_embeddings: MultiModalEmbeddings | None = None,
) -> torch.Tensor:
if multimodal_embeddings is None:
return self.language_model.embed_input_ids(input_ids)
inputs_embeds = self.language_model.embed_input_ids(input_ids)
if multimodal_embeddings is not None:
inputs_embeds = self.embed_input_ids( # type: ignore[call-overload]
input_ids,
inputs_embeds,
multimodal_embeddings,
[self.config.image_token_id, self.config.video_token_id],
)
return inputs_embeds
mm_token_ids = input_ids.new_tensor(
[self.config.image_token_id, self.config.video_token_id]
)
is_multimodal = torch.isin(input_ids, mm_token_ids)
return self.embed_input_ids(
input_ids,
multimodal_embeddings,
is_multimodal=is_multimodal,
)
def _process_image_input(self, image_input) -> tuple[torch.Tensor, ...]:
def _process_image_input(
self, image_input: OpenPanguVLImageInputs
) -> tuple[torch.Tensor, ...]:
grid_thw = image_input["image_grid_thw"]
if grid_thw.ndim != 2:
raise ValueError(f"grid_thw.ndim must be 2, but it is {grid_thw.ndim}")
@@ -1085,7 +1099,9 @@ class OpenPanguVLForConditionalGeneration(
sizes = grid_thw.prod(-1) // merge_size // merge_size
return image_embeds.split(sizes.tolist())
def _process_video_input(self, video_input) -> torch.Tensor:
def _process_video_input(
self, video_input: OpenPanguVLVideoInputs
) -> tuple[torch.Tensor, ...]:
grid_thw = video_input["video_grid_thw"]
if grid_thw.ndim != 2:
raise ValueError(f"grid_thw.ndim must be 2, but it is {grid_thw.ndim}")
@@ -1158,8 +1174,13 @@ class OpenPanguVLForConditionalGeneration(
for mm_feature in sorted(mm_features, key=lambda f: f.mm_position.offset):
offset = mm_feature.mm_position.offset
modality = mm_feature.modality
feature_data = mm_feature.data
assert feature_data is not None
if modality == "image":
t, h, w = mm_feature.data["image_grid_thw"].data.tolist()
grid_thw_item = feature_data["image_grid_thw"]
grid_thw = grid_thw_item.data
assert isinstance(grid_thw, torch.Tensor)
t, h, w = grid_thw.tolist()
assert t == 1, f"Image must have 1 frame, got {t}"
yield (
modality,
@@ -1169,7 +1190,10 @@ class OpenPanguVLForConditionalGeneration(
w // spatial_merge_size,
)
elif modality == "video":
t, h, w = mm_feature.data["video_grid_thw"].data.tolist()
grid_thw_item = feature_data["video_grid_thw"]
grid_thw = grid_thw_item.data
assert isinstance(grid_thw, torch.Tensor)
t, h, w = grid_thw.tolist()
yield (
modality,
offset,
@@ -1281,14 +1305,17 @@ def normalize(image, mean, std):
@lru_cache(maxsize=10)
def _fuse_mean_std_and_rescale_factor(
do_normalize: bool | None = None,
image_mean: float | list[float] | None = None,
image_std: float | list[float] | None = None,
image_mean: float | tuple[float, ...] | None = None,
image_std: float | tuple[float, ...] | None = None,
do_rescale: bool | None = None,
rescale_factor: float | None = None,
device: Optional["torch.device"] = None,
) -> tuple:
if do_rescale and do_normalize:
# Fused rescale and normalize
assert rescale_factor is not None
assert image_mean is not None
assert image_std is not None
image_mean = torch.tensor(image_mean, device=device) * (1.0 / rescale_factor)
image_std = torch.tensor(image_std, device=device) * (1.0 / rescale_factor)
do_rescale = False
@@ -1300,8 +1327,8 @@ def rescale_and_normalize(
do_rescale: bool,
rescale_factor: float,
do_normalize: bool,
image_mean: float | list[float],
image_std: float | list[float],
image_mean: float | tuple[float, ...],
image_std: float | tuple[float, ...],
dtype: torch.dtype = torch.bfloat16,
) -> "torch.Tensor":
"""Rescale and normalize images."""
+1
View File
@@ -349,6 +349,7 @@ class OpenVLAMultiModalProcessor(BaseMultiModalProcessor[OpenVLAProcessingInfo])
if isinstance(images, ImageEmbeddingItems):
num_image_tokens = images.get_feature_size(item_idx)
else:
assert isinstance(images, ImageProcessorItems)
image_size = images.get_image_size(item_idx)
num_image_tokens = self.info.get_num_image_tokens(
image_width=image_size.width,
+2
View File
@@ -218,6 +218,7 @@ class OPTDecoder(nn.Module):
)
# Project out & in will be replicated if they exist.
self.project_out: ReplicatedLinear | None
if config.word_embed_proj_dim != config.hidden_size:
self.project_out = ReplicatedLinear(
config.hidden_size,
@@ -229,6 +230,7 @@ class OPTDecoder(nn.Module):
else:
self.project_out = None
self.project_in: ReplicatedLinear | None
if config.word_embed_proj_dim != config.hidden_size:
self.project_in = ReplicatedLinear(
config.word_embed_proj_dim,
+10 -3
View File
@@ -59,7 +59,12 @@ from vllm.sequence import IntermediateTensors
from vllm.transformers_utils.processors.ovis import OvisProcessor
from vllm.utils.tensor_schema import TensorSchema, TensorShape
from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP
from .interfaces import (
MultiModalEmbeddings,
SupportsMultiModal,
SupportsPP,
supports_pp,
)
# Cannot find the following number from hf config.
IMAGE_TOKEN = "<image>"
@@ -281,7 +286,7 @@ class OvisProcessingInfo(BaseProcessingInfo):
# minus 1 for presented image token
return (patch_grid_length // hidden_stride) ** 2 - 1
def get_image_pad_token(self) -> str:
def get_image_pad_token(self) -> str | None:
hf_text_config = self.get_hf_config().get_text_config()
text_model_type = hf_text_config.model_type
return IMAGE_PAD_TOKEN_MAP.get(text_model_type)
@@ -439,8 +444,10 @@ class Ovis(nn.Module, SupportsMultiModal, SupportsPP):
text_model_type = self.config.get_text_config().model_type
self.image_pad_token_id = IMAGE_PAD_TOKEN_ID_MAP[text_model_type]
language_model = self.get_language_model()
assert supports_pp(language_model)
self.make_empty_intermediate_tensors = (
self.get_language_model().make_empty_intermediate_tensors
language_model.make_empty_intermediate_tensors
)
def _parse_and_validate_image_input(
+23 -6
View File
@@ -4,14 +4,14 @@
from collections.abc import Iterable, Mapping
from functools import partial
from typing import Annotated, Literal
from typing import Annotated, Literal, TypedDict
import torch
import torch.nn as nn
from transformers import BaseImageProcessor, BatchFeature, PretrainedConfig
from vllm.config import VllmConfig
from vllm.config.multimodal import BaseDummyOptions
from vllm.config.multimodal import BaseDummyOptions, VideoDummyOptions
from vllm.inputs import MultiModalDataDict
from vllm.model_executor.layers.linear import ReplicatedLinear
from vllm.model_executor.layers.quantization import QuantizationConfig
@@ -40,7 +40,12 @@ from vllm.sequence import IntermediateTensors
from vllm.transformers_utils.processors.ovis2_5 import Ovis2_5Processor
from vllm.utils.tensor_schema import TensorSchema, TensorShape
from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP
from .interfaces import (
MultiModalEmbeddings,
SupportsMultiModal,
SupportsPP,
supports_pp,
)
IMAGE_TOKEN = "<image>"
VIDEO_TOKEN = "<video>"
@@ -80,6 +85,11 @@ class Ovis2_5VideoPatchInputs(TensorSchema):
# This is used to restore the first two dimensions of `flat_data`.
class Ovis2_5MultiModalInputs(TypedDict, total=False):
images: Ovis2_5ImagePatchInputs | None
videos: Ovis2_5VideoPatchInputs | None
class VisualTokenizer(torch.nn.Module):
"""VIT."""
@@ -292,6 +302,7 @@ class Ovis2_5DummyInputsBuilder(BaseDummyInputsBuilder[Ovis2_5ProcessingInfo]):
image_overrides = mm_options.get("image")
video_overrides = mm_options.get("video")
assert video_overrides is None or isinstance(video_overrides, VideoDummyOptions)
mm_data = {
"image": self._get_dummy_images(
@@ -453,8 +464,10 @@ class Ovis2_5(nn.Module, SupportsMultiModal, SupportsPP):
self.image_pad_token_id: int = IMAGE_PAD_TOKEN_ID
language_model = self.get_language_model()
assert supports_pp(language_model)
self.make_empty_intermediate_tensors = (
self.get_language_model().make_empty_intermediate_tensors
language_model.make_empty_intermediate_tensors
)
def _parse_and_validate_image_input(
@@ -562,8 +575,10 @@ class Ovis2_5(nn.Module, SupportsMultiModal, SupportsPP):
vision_embeddings.append(torch.cat(vision_embeddings_per_image, dim=0))
return tuple(vision_embeddings)
def _parse_and_validate_multimodal_inputs(self, **kwargs: object) -> dict:
modalities = {}
def _parse_and_validate_multimodal_inputs(
self, **kwargs: object
) -> Ovis2_5MultiModalInputs:
modalities: Ovis2_5MultiModalInputs = {}
# Preserve the order of modalities if there are multiple of them
# from the order of kwargs.
@@ -593,10 +608,12 @@ class Ovis2_5(nn.Module, SupportsMultiModal, SupportsPP):
for modality in modalities:
if modality == "images":
image_input = modalities["images"]
assert image_input is not None
image_embeddings = self._process_visual_input(image_input)
multimodal_embeddings += tuple(image_embeddings)
if modality == "videos":
video_input = modalities["videos"]
assert video_input is not None
video_embeddings = self._process_visual_input(video_input)
multimodal_embeddings += tuple(video_embeddings)
@@ -551,7 +551,9 @@ class DynamicResolutionImageTiler:
)
@staticmethod
def stack(images: list[torch.Tensor], patch_size: int) -> torch.Tensor:
def stack(
images: list[torch.Tensor] | torch.Tensor, patch_size: int
) -> torch.Tensor:
assert len(images) > 0, "No images to stack"
def rearrange_img(x):