[Mypy] Fix typing for R/S models (#54157)

Signed-off-by: Taneem Ibrahim <taneem.ibrahim@gmail.com>
This commit is contained in:
Taneem Ibrahim
2026-09-10 09:43:08 -04:00
committed by GitHub
parent 73fb19151f
commit f9083cb83a
15 changed files with 146 additions and 85 deletions
-2
View File
@@ -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]",
+26 -8
View File
@@ -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
+5 -2
View File
@@ -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)
+6 -3
View File
@@ -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)
+4 -6
View File
@@ -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
+6 -4
View File
@@ -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,
@@ -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),
+14 -7
View File
@@ -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
+2
View File
@@ -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)
+7 -7
View File
@@ -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)
+19 -7
View File
@@ -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
+39 -29
View File
@@ -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
+2 -1
View File
@@ -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."
)
+6 -3
View File
@@ -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(
+9 -6
View File
@@ -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(