[Mypy] Fix typing for R/S models (#54157)
Signed-off-by: Taneem Ibrahim <taneem.ibrahim@gmail.com>
This commit is contained in:
@@ -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]",
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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),
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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."
|
||||
)
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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(
|
||||
|
||||
Reference in New Issue
Block a user