From f9083cb83af44debe290f3753ca2effc292ea1c1 Mon Sep 17 00:00:00 2001 From: Taneem Ibrahim Date: Thu, 10 Sep 2026 08:43:08 -0500 Subject: [PATCH] [Mypy] Fix typing for R/S models (#54157) Signed-off-by: Taneem Ibrahim --- tools/pre_commit/mypy.py | 2 - vllm/model_executor/models/radio.py | 34 ++++++++--- vllm/model_executor/models/registry.py | 7 ++- vllm/model_executor/models/roberta.py | 9 ++- vllm/model_executor/models/sarvam.py | 10 ++-- vllm/model_executor/models/siglip.py | 10 ++-- vllm/model_executor/models/siglip2navit.py | 1 + vllm/model_executor/models/skyworkr1v.py | 21 ++++--- vllm/model_executor/models/solar.py | 2 + vllm/model_executor/models/step3_text.py | 14 ++--- vllm/model_executor/models/step3_vl.py | 26 ++++++--- vllm/model_executor/models/step3p5.py | 68 +++++++++++++--------- vllm/model_executor/models/step3p5_mtp.py | 3 +- vllm/model_executor/models/step3p7.py | 9 ++- vllm/model_executor/models/step_vl.py | 15 +++-- 15 files changed, 146 insertions(+), 85 deletions(-) diff --git a/tools/pre_commit/mypy.py b/tools/pre_commit/mypy.py index c9154de048..9da39acc78 100755 --- a/tools/pre_commit/mypy.py +++ b/tools/pre_commit/mypy.py @@ -116,8 +116,6 @@ EXCLUDE = [ r"vllm/model_executor/models/[nN]", r"vllm/model_executor/models/[oO]", r"vllm/model_executor/models/[qQ]", - r"vllm/model_executor/models/[rR]", - r"vllm/model_executor/models/[sS]", r"vllm/model_executor/models/transformers", r"vllm/model_executor/models/[uU]", r"vllm/model_executor/models/[vV]", diff --git a/vllm/model_executor/models/radio.py b/vllm/model_executor/models/radio.py index 213083392d..47dc8f89d3 100644 --- a/vllm/model_executor/models/radio.py +++ b/vllm/model_executor/models/radio.py @@ -137,8 +137,9 @@ class ViTPatchGenerator(nn.Module): if isinstance(max_input_dims, int): max_input_dims = (max_input_dims, max_input_dims) - max_input_dims = tuple( - int(math.ceil(d / patch_size) * patch_size) for d in max_input_dims + max_input_dims = ( + int(math.ceil(max_input_dims[0] / patch_size) * patch_size), + int(math.ceil(max_input_dims[1] / patch_size) * patch_size), ) self.cpe_mode = max_input_dims != input_dims @@ -154,7 +155,10 @@ class ViTPatchGenerator(nn.Module): self.num_rows = max_input_dims[0] // patch_size self.num_cols = max_input_dims[1] // patch_size - self.input_dims = tuple(d // patch_size for d in input_dims) + self.input_dims = ( + input_dims[0] // patch_size, + input_dims[1] // patch_size, + ) self.num_patches = self.num_rows * self.num_cols self.max_input_dims = max_input_dims @@ -369,7 +373,10 @@ class ViTPatchGenerator(nn.Module): if input_size is None: input_dims = self.input_dims else: - input_dims = tuple(d // self.patch_size for d in input_size) + input_dims = ( + input_size[0] // self.patch_size, + input_size[1] // self.patch_size, + ) pos_embed = self._get_pos_embeddings(batch_size, input_dims) @@ -538,8 +545,9 @@ class RadioInternVisionModel(nn.Module): self.config = config self.img_size, self.grid_size, self.num_patches = self._init_img_size( - to_2tuple(config.patch_size), config.image_size + config.patch_size, config.image_size ) + assert self.img_size is not None max_img_size = int( round(config.cpe_max_size / config.patch_size) * config.patch_size ) @@ -565,11 +573,21 @@ class RadioInternVisionModel(nn.Module): prefix=f"{prefix}.encoder", ) - def _init_img_size(self, patch_size, img_size: int | tuple[int, int]): + def _init_img_size( + self, + patch_size: input_dim_t, + img_size: input_dim_t | None, + ) -> tuple[tuple[int, int] | None, tuple[int, int] | None, int | None]: if img_size is None: return None, None, None - img_size = to_2tuple(img_size) - grid_size = tuple([s // p for s, p in zip(img_size, patch_size)]) + if isinstance(img_size, int): + img_size = (img_size, img_size) + if isinstance(patch_size, int): + patch_size = (patch_size, patch_size) + grid_size = ( + img_size[0] // patch_size[0], + img_size[1] // patch_size[1], + ) num_patches = grid_size[0] * grid_size[1] return img_size, grid_size, num_patches diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index 3488d00272..a9fe7c2ec2 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -1264,9 +1264,11 @@ class _ModelRegistry: "'auto_map' (relevant if the model is custom)." ) + assert issubclass(model_module, transformers.PreTrainedModel) + transformers_model_cls: type[transformers.PreTrainedModel] = model_module if not ( - model_module.is_backend_compatible() - or model_module._can_set_attn_implementation() + transformers_model_cls.is_backend_compatible() + or transformers_model_cls._can_set_attn_implementation() ): if model_config.model_impl != "transformers": return None @@ -1323,6 +1325,7 @@ class _ModelRegistry: return (model_info, arch) elif model_config.model_impl == "terratorch": model_info = self._try_inspect_model_cls("Terratorch") + assert model_info is not None return (model_info, "Terratorch") # Fallback to transformers impl (after resolving convert_type) diff --git a/vllm/model_executor/models/roberta.py b/vllm/model_executor/models/roberta.py index aca929755f..59054b0667 100644 --- a/vllm/model_executor/models/roberta.py +++ b/vllm/model_executor/models/roberta.py @@ -139,11 +139,14 @@ class RobertaEmbeddingModel(BertEmbeddingModel): self, vllm_config: VllmConfig, prefix: str = "" ) -> BertModel | BertWithRope: hf_config = vllm_config.model_config.hf_config - kwargs = dict(vllm_config=vllm_config, prefix=prefix) if getattr(hf_config, "position_embedding_type", "absolute") == "absolute": - return BertModel(**kwargs, embedding_class=RobertaEmbedding) + return BertModel( + vllm_config=vllm_config, + prefix=prefix, + embedding_class=RobertaEmbedding, + ) else: - return JinaRobertaModel(**kwargs) + return JinaRobertaModel(vllm_config=vllm_config, prefix=prefix) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): weights_list = list(weights) diff --git a/vllm/model_executor/models/sarvam.py b/vllm/model_executor/models/sarvam.py index e790a27c31..14d68fec56 100644 --- a/vllm/model_executor/models/sarvam.py +++ b/vllm/model_executor/models/sarvam.py @@ -325,6 +325,7 @@ class SarvamMLAMoE(nn.Module): self.score_function = getattr(config, "score_function", "sigmoid") self.num_shared_experts = getattr(config, "num_shared_experts", 1) + self.shared_experts: SarvamMLAMLP | None if self.num_shared_experts > 0: if hasattr(config, "moe_shared_expert_intermediate_size"): shared_int = config.moe_shared_expert_intermediate_size @@ -590,6 +591,9 @@ class SarvamMLAModel(nn.Module, EagleModelMixin): class SarvamMixtureOfExperts(MixtureOfExperts): + moe_layers: list[MoERunner] + moe_mlp_layers: list[SarvamMLAMoE] + def extract_moe_parameters(self, example_moe: SarvamMLAMoE | None) -> None: if example_moe is None: raise RuntimeError("No SarvamMLAMoE layer found in model.layers.") @@ -626,12 +630,6 @@ class SarvamMixtureOfExperts(MixtureOfExperts): if hasattr(fused, "update_expert_map"): fused.update_expert_map() - def set_eplb_state(self, eplb_state) -> None: - self.eplb_state = eplb_state - for moe in self.moe_layers: - if hasattr(moe, "set_eplb_state"): - moe.set_eplb_state(eplb_state) - class SarvamMLAForCausalLM( nn.Module, SupportsPP, SupportsLoRA, SupportsEagle3, SarvamMixtureOfExperts diff --git a/vllm/model_executor/models/siglip.py b/vllm/model_executor/models/siglip.py index c5625b217f..6d1fd02962 100644 --- a/vllm/model_executor/models/siglip.py +++ b/vllm/model_executor/models/siglip.py @@ -127,13 +127,15 @@ class SiglipProcessingInfo(BaseProcessingInfo): pooler_config = self.ctx.model_config.pooler_config assert pooler_config is not None + pooling_type = pooler_config.seq_pooling_type + assert pooling_type is not None return get_num_selected_vision_tokens( vision_encoder_info.get_num_image_tokens( image_width=image_width, image_height=image_height, ), - _get_vision_feature_select_strategy(pooler_config.seq_pooling_type), + _get_vision_feature_select_strategy(pooling_type), ) def get_image_size_with_most_features(self) -> ImageSize: @@ -1021,9 +1023,9 @@ class SiglipEmbeddingModel(nn.Module, SupportsMultiModal, SupportsQuant): feature_select_strategy: VisionFeatureSelectStrategy | None = None, ) -> torch.Tensor: if feature_select_strategy is None: - feature_select_strategy = _get_vision_feature_select_strategy( - self.pooler_config.seq_pooling_type - ) + pooling_type = self.pooler_config.seq_pooling_type + assert pooling_type is not None + feature_select_strategy = _get_vision_feature_select_strategy(pooling_type) pooled_output = self.vision_model( pixel_values=pixel_values, diff --git a/vllm/model_executor/models/siglip2navit.py b/vllm/model_executor/models/siglip2navit.py index 5ee1c2c910..c16619612f 100644 --- a/vllm/model_executor/models/siglip2navit.py +++ b/vllm/model_executor/models/siglip2navit.py @@ -248,6 +248,7 @@ class Siglip2Attention(nn.Module): values = values.view(seq_length, self.num_heads_per_partition, self.head_dim) if self.use_rope: + assert position_embeddings is not None cos, sin = position_embeddings queries, keys = apply_rotary_pos_emb( queries.unsqueeze(0), diff --git a/vllm/model_executor/models/skyworkr1v.py b/vllm/model_executor/models/skyworkr1v.py index c44460aa0c..ef2626db1f 100644 --- a/vllm/model_executor/models/skyworkr1v.py +++ b/vllm/model_executor/models/skyworkr1v.py @@ -29,7 +29,12 @@ from vllm.transformers_utils.processors.internvl import ( ) from vllm.utils.tensor_schema import TensorSchema, TensorShape -from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP +from .interfaces import ( + MultiModalEmbeddings, + SupportsMultiModal, + SupportsPP, + supports_pp, +) from .internvl import ( BaseInternVLDummyInputsBuilder, BaseInternVLMultiModalProcessor, @@ -156,7 +161,7 @@ class SkyworkR1VChatModel(nn.Module, SupportsMultiModal, SupportsPP): 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.config = config self.multimodal_config = multimodal_config @@ -182,20 +187,22 @@ class SkyworkR1VChatModel(nn.Module, SupportsMultiModal, SupportsPP): ) with self._mark_language_model(vllm_config): - self.language_model = init_vllm_registered_model( + language_model = init_vllm_registered_model( vllm_config=vllm_config, hf_config=config.text_config, prefix=maybe_prefix(prefix, "language_model"), ) + self.language_model = language_model + assert supports_pp(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 = ( - self.language_model.make_empty_intermediate_tensors + language_model.make_empty_intermediate_tensors ) 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 @@ -232,7 +239,7 @@ class SkyworkR1VChatModel(nn.Module, SupportsMultiModal, SupportsPP): def _init_mlp1( self, config: PretrainedConfig, - quant_config: QuantizationConfig, + quant_config: QuantizationConfig | None, prefix: str = "", ) -> nn.Module: vit_hidden_size = config.vision_config.hidden_size diff --git a/vllm/model_executor/models/solar.py b/vllm/model_executor/models/solar.py index 430337c77d..cb175f7c7a 100644 --- a/vllm/model_executor/models/solar.py +++ b/vllm/model_executor/models/solar.py @@ -318,9 +318,11 @@ class SolarModel(nn.Module): for i in range(self.start_layer, self.end_layer): if i in self.config.bskcn_1: bskcn_h_1 = hidden_states.clone() + assert residual is not None bskcn_r_1 = residual.clone() if i in self.config.bskcn_2: bskcn_h_2 = hidden_states.clone() + assert residual is not None bskcn_r_2 = residual.clone() if i in self.config.bskcn_3: hidden_states = bskcn_h_1 * bskcn_tv + hidden_states * (1 - bskcn_tv) diff --git a/vllm/model_executor/models/step3_text.py b/vllm/model_executor/models/step3_text.py index 446e66ca8b..57c35ed8ce 100644 --- a/vllm/model_executor/models/step3_text.py +++ b/vllm/model_executor/models/step3_text.py @@ -10,7 +10,7 @@ import torch from torch import nn from vllm.compilation.decorators import support_torch_compile -from vllm.config import CacheConfig, ModelConfig, VllmConfig +from vllm.config import CacheConfig, VllmConfig from vllm.distributed import ( get_pp_group, get_tensor_model_parallel_world_size, @@ -55,7 +55,7 @@ logger = init_logger(__name__) class FusedMoEBlock(nn.Module): def __init__( self, - config: ModelConfig, + config: Step3TextConfig, quant_config: QuantizationConfig | None = None, prefix: str = "", ): @@ -384,7 +384,7 @@ class Step3TextModel(nn.Module): return hidden_states def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - qkv_params_mapping = [ + qkv_params_mapping: list[tuple[str, str, float, float]] = [ # (param_name, shard_name, relative_start_idx, relative_end_idx) ( ".qkv_proj", @@ -410,7 +410,7 @@ class Step3TextModel(nn.Module): / (self.config.share_q_dim + self.config.head_dim * 2), ), ] - stacked_params_mapping = [ + stacked_params_mapping: list[tuple[str, str, int]] = [ # (param_name, shard_name, shard_id) (".gate_up_proj", ".gate_proj", 0), (".gate_up_proj", ".up_proj", 1), @@ -421,7 +421,7 @@ class Step3TextModel(nn.Module): "base_layer." if any(".base_layer." in name for name in params_dict) else "" ) - expert_params_mapping = [ + expert_params_mapping: list[tuple[str, str, str]] = [ ( f".moe.experts.routed_experts.{base_layer}w13_weight", ".moe.gate_proj.weight", @@ -460,7 +460,7 @@ class Step3TextModel(nn.Module): break else: for mapping in expert_params_mapping: - param_name, weight_name, shard_id = mapping + param_name, weight_name, expert_shard_id = mapping if weight_name not in name: continue name = name.replace(weight_name, param_name) @@ -480,7 +480,7 @@ class Step3TextModel(nn.Module): param, loaded_weight_expert, name, - shard_id=shard_id, + shard_id=expert_shard_id, expert_id=expert_id, ) loaded_params.add(name) diff --git a/vllm/model_executor/models/step3_vl.py b/vllm/model_executor/models/step3_vl.py index 9f738b7c64..b060ed1d15 100644 --- a/vllm/model_executor/models/step3_vl.py +++ b/vllm/model_executor/models/step3_vl.py @@ -51,6 +51,7 @@ from .interfaces import ( SupportsEncoderCudaGraph, SupportsMultiModal, SupportsPP, + supports_pp, ) from .utils import ( AutoWeightsLoader, @@ -105,10 +106,10 @@ class Step3VLProcessingInfo(BaseProcessingInfo): return Step3VLImageProcessor(**kwargs) - def get_hf_processor(self) -> Step3VLProcessor: + def get_hf_processor(self, **kwargs: object) -> Step3VLProcessor: return Step3VLProcessor( tokenizer=self.get_tokenizer(), - image_processor=self.get_image_processor(), + image_processor=self.get_image_processor(**kwargs), ) def get_supported_mm_limits(self) -> Mapping[str, int | None]: @@ -169,8 +170,11 @@ class Step3VLMultiModalProcessor(BaseMultiModalProcessor[Step3VLProcessingInfo]) def get_replacement_step1o(item_idx: int): out_item = out_mm_kwargs["image"][item_idx] - num_patches = int(out_item["num_patches"].data) + num_patches_data = out_item["num_patches"].data patch_newline_mask = out_item["patch_newline_mask"].data + assert isinstance(num_patches_data, torch.Tensor) + assert isinstance(patch_newline_mask, torch.Tensor) + num_patches = int(num_patches_data.item()) image_repl_ids = hf_processor.get_image_repl_feature_ids( 1, num_patches, patch_newline_mask.tolist() ) @@ -514,7 +518,7 @@ class Step3VLForConditionalGeneration( def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: super().__init__() config = vllm_config.model_config.hf_config - multimodal_config = vllm_config.model_config.multimodal_config + multimodal_config = vllm_config.model_config.get_multimodal_config() self.config = config self.model_config = vllm_config.model_config @@ -560,14 +564,16 @@ class Step3VLForConditionalGeneration( ) with self._mark_language_model(vllm_config): - self.language_model = init_vllm_registered_model( + language_model = init_vllm_registered_model( vllm_config=vllm_config, hf_config=config.text_config, prefix=maybe_prefix(prefix, "language_model"), ) + self.language_model = language_model + assert supports_pp(language_model) self.make_empty_intermediate_tensors = ( - self.language_model.make_empty_intermediate_tensors + language_model.make_empty_intermediate_tensors ) @property @@ -626,6 +632,9 @@ class Step3VLForConditionalGeneration( return None if pixel_values is not None and patch_pixel_values is not None: + assert isinstance(pixel_values, torch.Tensor) + assert isinstance(patch_pixel_values, torch.Tensor) + assert isinstance(num_patches, torch.Tensor) return Step3VLImagePixelInputs( type="pixel_values", pixel_values=pixel_values.to(self.dtype), @@ -634,6 +643,7 @@ class Step3VLForConditionalGeneration( ) if image_embeds is not None: + assert isinstance(image_embeds, torch.Tensor) return Step3VLImageEmbeddingInputs( type="image_embeds", data=image_embeds.to(self.dtype), @@ -657,7 +667,7 @@ class Step3VLForConditionalGeneration( def _process_image_input( self, image_input: Step3VLImageInputs - ) -> tuple[torch.Tensor, ...]: + ) -> list[torch.Tensor]: if image_input["type"] == "image_embeds": image_features = image_input["data"] return [ @@ -686,6 +696,7 @@ class Step3VLForConditionalGeneration( for i, num_patch in enumerate(num_patches_list): cur_feature = [] if num_patch > 0: + assert patch_image_features is not None patch_slice = patch_image_features[ cur_patch_idx : cur_patch_idx + num_patch ] @@ -765,6 +776,7 @@ class Step3VLForConditionalGeneration( from vllm.v1.worker.encoder_cudagraph_defs import EncoderItemSpec num_patches = mm_kwargs.get("num_patches") + assert isinstance(num_patches, torch.Tensor) img_grid = ( self.config.vision_config.image_size // self.config.vision_config.patch_size diff --git a/vllm/model_executor/models/step3p5.py b/vllm/model_executor/models/step3p5.py index 5c6cee77d7..3a928aa59d 100644 --- a/vllm/model_executor/models/step3p5.py +++ b/vllm/model_executor/models/step3p5.py @@ -11,7 +11,7 @@ from torch import nn from torch.nn.parameter import Parameter from vllm.compilation.decorators import support_torch_compile -from vllm.config import CacheConfig, ModelConfig, VllmConfig +from vllm.config import CacheConfig, VllmConfig from vllm.distributed import ( get_dp_group, get_ep_group, @@ -47,6 +47,11 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader from vllm.sequence import IntermediateTensors from vllm.v1.attention.backend import AttentionType +if typing.TYPE_CHECKING: + from vllm.transformers_utils.configs.step3p5 import Step3p5Config +else: + from transformers import PretrainedConfig as Step3p5Config + from .interfaces import MixtureOfExperts, SupportsPP from .utils import ( AutoWeightsLoader, @@ -79,7 +84,7 @@ class FP32ReplicatedLinear(ReplicatedLinear): class Step3p5MLP(nn.Module): def __init__( self, - config: ModelConfig, + config: Step3p5Config, hidden_size: int, intermediate_size: int, hidden_act: str, @@ -147,9 +152,9 @@ class Step3p5Attention(nn.Module): # Step3p5 specific args sliding_window: int | None = None, use_head_wise_attn_gate: bool = False, - layer_types: list = None, - use_rope_layers: list = None, - yarn_only_types: list = None, + layer_types: list[str] | None = None, + use_rope_layers: list[bool] | None = None, + yarn_only_types: list[str] | None = None, swa_num_attention_heads: int | None = None, partial_rotary_factor: float = 1.0, ): @@ -162,8 +167,10 @@ class Step3p5Attention(nn.Module): enable_sliding_window = layer_types[self.layer_idx] == "sliding_attention" else: enable_sliding_window = self.layer_idx % 2 == 0 - if yarn_only_types and layer_types[self.layer_idx] not in yarn_only_types: - rope_scaling = None + if yarn_only_types: + assert layer_types is not None + if layer_types[self.layer_idx] not in yarn_only_types: + rope_scaling = None if sliding_window is not None and enable_sliding_window: sliding_window = sliding_window @@ -304,6 +311,7 @@ class FusedMoEBlock(nn.Module): self.ep_size = get_ep_group().device_group.size() config = vllm_config.model_config.hf_config + assert isinstance(config, Step3p5Config) quant_config = vllm_config.quant_config parallel_config = vllm_config.parallel_config @@ -349,6 +357,7 @@ class FusedMoEBlock(nn.Module): else None ) if swiglu_limit not in (None, 0): + assert swiglu_limit is not None swiglu_limit = float(swiglu_limit) assert swiglu_limit == 7.0, ( "Swiglu limit in fused moe block only support 7.0 now." @@ -408,6 +417,7 @@ class Step3p5DecoderLayer(nn.Module): ) -> None: super().__init__() config = vllm_config.model_config.hf_config + assert isinstance(config, Step3p5Config) self.hidden_size = config.hidden_size layer_idx = extract_layer_index(prefix) self.layer_idx = layer_idx @@ -419,20 +429,18 @@ class Step3p5DecoderLayer(nn.Module): num_attention_heads = None num_attention_groups = None head_dim = None + attention_other_setting = config.attention_other_setting + layer_types = config.layer_types if ( - getattr(config, "attention_other_setting", None) - and getattr(config, "layer_types", []) - and config.layer_types[layer_idx] - == config.attention_other_setting["attention_type"] + attention_other_setting + and layer_types + and layer_types[layer_idx] == attention_other_setting["attention_type"] ): - num_attention_heads = config.attention_other_setting[ - "num_attention_heads" - ] - num_attention_groups = config.attention_other_setting[ - "num_attention_groups" - ] - head_dim = config.attention_other_setting["head_dim"] + num_attention_heads = attention_other_setting["num_attention_heads"] + num_attention_groups = attention_other_setting["num_attention_groups"] + head_dim = attention_other_setting["head_dim"] partial_rotary_factors = getattr(config, "partial_rotary_factors", []) + assert config.max_position_embeddings is not None self.self_attn = Step3p5Attention( hidden_size=self.hidden_size, num_heads=num_attention_heads @@ -539,6 +547,7 @@ class Step3p5Model(nn.Module): self.vllm_config = vllm_config config = vllm_config.model_config.hf_config + assert isinstance(config, Step3p5Config) self.vocab_size = config.vocab_size self.config = config @@ -605,8 +614,8 @@ class Step3p5Model(nn.Module): def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: config = self.config assert config.num_attention_groups > 1, "Only support GQA" - qkv_params_mapping = [] - stacked_params_mapping = [ + qkv_params_mapping: list[tuple[str, str, float, float]] = [] + stacked_params_mapping: list[tuple[str, str, str | int]] = [ # (param_name, shard_name, shard_id) ("qkv_proj", "q_proj", "q"), ("qkv_proj", "k_proj", "k"), @@ -622,7 +631,7 @@ class Step3p5Model(nn.Module): ) # Old packed 3D format: .moe.gate_proj.weight [num_experts, out, in] - expert_params_mapping = [ + expert_params_mapping: list[tuple[str, str, str]] = [ ( f".moe.experts.routed_experts.{base_layer}w13_weight", ".moe.gate_proj.weight", @@ -758,16 +767,16 @@ class Step3p5Model(nn.Module): # Not an expert proj — use default loader # (e.g. share_expert weights if they matched) param = params_dict[local_name] - weight_loader = getattr( + fallback_loader = getattr( param, "weight_loader", default_weight_loader, ) - weight_loader(param, loaded_weight) + fallback_loader(param, loaded_weight) loaded_params.add(local_name) continue - for param_name, weight_name, shard_id in stacked_params_mapping: + for param_name, weight_name, stacked_shard_id in stacked_params_mapping: if weight_name not in local_name: continue if any( @@ -782,11 +791,11 @@ class Step3p5Model(nn.Module): continue param = params_dict[replaced_name] weight_loader = param.weight_loader - weight_loader(param, loaded_weight, shard_id) + weight_loader(param, loaded_weight, stacked_shard_id) loaded_params.add(replaced_name) break else: - for param_name, weight_name, shard_id in expert_params_mapping: + for param_name, weight_name, expert_shard_id in expert_params_mapping: if weight_name not in local_name: continue replaced_name = local_name.replace(weight_name, param_name) @@ -823,7 +832,7 @@ class Step3p5Model(nn.Module): param, loaded_weight_expert, replaced_name, - shard_id=shard_id, + shard_id=expert_shard_id, expert_id=expert_id, ) loaded_params.add(replaced_name) @@ -861,10 +870,10 @@ class Step3p5Model(nn.Module): if local_name not in params_dict: continue param = params_dict[local_name] - weight_loader = getattr( + fallback_loader = getattr( param, "weight_loader", default_weight_loader ) - weight_loader(param, loaded_weight) + fallback_loader(param, loaded_weight) loaded_params.add(local_name) return loaded_params @@ -918,6 +927,7 @@ class Step3p5ForCausalLM(nn.Module, SupportsPP, MixtureOfExperts): self.moe_layers.append(layer.moe.experts) assert len(self.moe_layers) > 0, "No MoE layers found in the model." + assert example_layer is not None self.num_moe_layers = len(self.moe_layers) self.num_expert_groups = 1 self.num_shared_experts = 0 diff --git a/vllm/model_executor/models/step3p5_mtp.py b/vllm/model_executor/models/step3p5_mtp.py index 87b7aab491..294b3559c0 100644 --- a/vllm/model_executor/models/step3p5_mtp.py +++ b/vllm/model_executor/models/step3p5_mtp.py @@ -305,7 +305,7 @@ class Step3p5MTP(nn.Module): ) return loaded_params - def _rewrite_spec_layer_name(self, spec_layer: int, name: str) -> str: + def _rewrite_spec_layer_name(self, spec_layer: int | None, name: str) -> str: """ Rewrite the weight name to match the format of the original model. Add .mtp_block for modules in transformer layer block for spec layer @@ -324,6 +324,7 @@ class Step3p5MTP(nn.Module): break if not spec_layer_weight: # treat rest weights as weights for transformer layer block + assert spec_layer is not None name = name.replace( f"model.layers.{spec_layer}.", f"model.layers.{spec_layer}.mtp_block." ) diff --git a/vllm/model_executor/models/step3p7.py b/vllm/model_executor/models/step3p7.py index 2a98ec53fc..0fa2a32306 100644 --- a/vllm/model_executor/models/step3p7.py +++ b/vllm/model_executor/models/step3p7.py @@ -9,6 +9,7 @@ from vllm.logger import init_logger from vllm.model_executor.layers.activation import get_act_fn from vllm.model_executor.layers.linear import ColumnParallelLinear +from .interfaces import supports_pp from .step3_vl import Step3VLForConditionalGeneration from .step_vl import PerceptionEncoder from .utils import WeightsMapper, init_vllm_registered_model, maybe_prefix @@ -41,7 +42,7 @@ class Step3p7ForConditionalGeneration(Step3VLForConditionalGeneration): super(Step3VLForConditionalGeneration, self).__init__() config = vllm_config.model_config.hf_config - multimodal_config = vllm_config.model_config.multimodal_config + multimodal_config = vllm_config.model_config.get_multimodal_config() quant_config = vllm_config.quant_config self.config = config @@ -66,14 +67,16 @@ class Step3p7ForConditionalGeneration(Step3VLForConditionalGeneration): ) with self._mark_language_model(vllm_config): - self.language_model = init_vllm_registered_model( + language_model = init_vllm_registered_model( vllm_config=vllm_config, hf_config=config.text_config, prefix=maybe_prefix(prefix, "language_model"), ) + self.language_model = language_model + assert supports_pp(language_model) self.make_empty_intermediate_tensors = ( - self.language_model.make_empty_intermediate_tensors + language_model.make_empty_intermediate_tensors ) def _get_vision_model_output( diff --git a/vllm/model_executor/models/step_vl.py b/vllm/model_executor/models/step_vl.py index 53dd3ebf48..247078e2ec 100644 --- a/vllm/model_executor/models/step_vl.py +++ b/vllm/model_executor/models/step_vl.py @@ -22,6 +22,7 @@ from vllm.model_executor.layers.linear import ( ) from vllm.model_executor.layers.quantization import QuantizationConfig +from .interfaces import supports_pp from .step3_vl import Step3VLForConditionalGeneration from .utils import WeightsMapper, init_vllm_registered_model, maybe_prefix from .vision import is_vit_use_data_parallel, run_dp_sharded_vision_model @@ -148,7 +149,7 @@ class PerceptionEncoderMLP(nn.Module): self, input_dim: int, hidden_dim: int, - act_layer: Callable[[], nn.Module], + act_layer: Callable[[torch.Tensor], torch.Tensor], quant_config: QuantizationConfig | None = None, prefix: str = "", ): @@ -257,7 +258,7 @@ class PerceptionEncoderVisionBlock(nn.Module): max_grid_height: int, max_grid_width: int, mlp_ratio: float = 4.0, - ls_init_value: float = None, + ls_init_value: float | None = None, act_layer: Callable = nn.GELU, norm_layer: Callable = nn.LayerNorm, use_cls_token: bool = False, @@ -310,7 +311,7 @@ class PerceptionEncoderVisionTransformer(nn.Module): max_grid_height: int, max_grid_width: int, mlp_ratio: float = 4.0, - ls_init_value: float = None, + ls_init_value: float | None = None, act_layer: Callable = nn.GELU, norm_layer: Callable = nn.LayerNorm, use_cls_token: bool = False, @@ -496,7 +497,7 @@ class StepVLForConditionalGeneration(Step3VLForConditionalGeneration): super(Step3VLForConditionalGeneration, self).__init__() config = vllm_config.model_config.hf_config - multimodal_config = vllm_config.model_config.multimodal_config + multimodal_config = vllm_config.model_config.get_multimodal_config() quant_config = vllm_config.quant_config self.config = config @@ -522,14 +523,16 @@ class StepVLForConditionalGeneration(Step3VLForConditionalGeneration): ) with self._mark_language_model(vllm_config): - self.language_model = init_vllm_registered_model( + language_model = init_vllm_registered_model( vllm_config=vllm_config, hf_config=config.text_config, prefix=maybe_prefix(prefix, "language_model"), ) + self.language_model = language_model + assert supports_pp(language_model) self.make_empty_intermediate_tensors = ( - self.language_model.make_empty_intermediate_tensors + language_model.make_empty_intermediate_tensors ) def _get_vision_model_output(