class Step5MTP(nn.Module, MixtureOfExperts, SupportsMultiModal):
supports_multimodal_embeddings = True
has_own_lm_head = True
packed_modules_mapping = STEP5_PACKED_MODULES_MAPPING
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
super().__init__()
# The target VL model maps quantization exclusions to language_model.*.
# Draft modules live at model.* and must not mutate the target config.
if vllm_config.quant_config is not None:
quant_config = deepcopy(vllm_config.quant_config)
quant_config.apply_vllm_mapper(
WeightsMapper(orig_to_new_prefix={"language_model.": ""})
)
vllm_config = replace(vllm_config, quant_config=quant_config)
self.config = _get_mtp_config(vllm_config)
self.vllm_config = vllm_config
self.model = Step5MultiTokenPredictor(
vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
)
# Expose the first draft head to the generic sharing policy. Each MTP
# layer retains its own checkpoint head for compute_logits.
self.lm_head = self.model.layers[
str(self.model.mtp_start_layer_idx)
].shared_head.head
self.num_moe_layers: int
self.num_expert_groups: int
self.num_shared_experts: int
self.num_logical_experts: int
self.num_physical_experts: int
self.num_local_physical_experts: int
self.num_routed_experts: int
self.num_redundant_experts: int
if not _is_step5_mtp_dense(self.model):
raise NotImplementedError("Step5 MTP supports dense draft layers only")
self.moe_layers = []
_set_step5_moe_protocol_metadata(self, None)
def embed_input_ids(
self,
input_ids: torch.Tensor,
multimodal_embeddings: MultiModalEmbeddings | None = None,
*,
is_multimodal: torch.Tensor | None = None,
) -> torch.Tensor:
inputs_embeds = self._embed_text_input_ids(
input_ids, self.model.embed_input_ids, is_multimodal=is_multimodal
)
if multimodal_embeddings is None or len(multimodal_embeddings) == 0:
return inputs_embeds
return _merge_multimodal_embeddings(
inputs_embeds=inputs_embeds,
multimodal_embeddings=multimodal_embeddings,
is_multimodal=_require_is_multimodal(is_multimodal),
)
def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
return self.model.get_input_embeddings(input_ids)
def forward(
self,
input_ids: torch.Tensor,
positions: torch.Tensor,
hidden_states: torch.Tensor,
intermediate_tensors: IntermediateTensors | None = None,
inputs_embeds: torch.Tensor | None = None,
spec_step_idx: int = 0,
) -> torch.Tensor:
return self.model(
input_ids, positions, hidden_states, inputs_embeds, spec_step_idx
)
def compute_logits(
self,
hidden_states: torch.Tensor,
spec_step_idx: int = 0,
) -> torch.Tensor | None:
return self.model.compute_logits(hidden_states, spec_step_idx)
def get_top_tokens(
self,
hidden_states: torch.Tensor,
spec_step_idx: int = 0,
) -> torch.Tensor:
return self.model.get_top_tokens(hidden_states, spec_step_idx)
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
from vllm.model_executor.model_loader.mtp_validation import (
is_mtp_completeness_check_enabled,
)
from vllm.model_executor.model_loader.weight_utils import (
default_weight_loader,
)
validate_completeness = is_mtp_completeness_check_enabled()
if validate_completeness:
_reset_fused_qkv_indexer_load_state(self)
params_dict = dict(self.named_parameters())
loaded_params: set[str] = set()
for name, loaded_weight in weights:
if "rotary_emb.inv_freq" in name:
continue
if name.startswith("layers."):
name = f"model.{name}"
spec_layer = get_spec_layer_idx_from_weight_name(self.config, name)
if "embed_tokens" not in name and spec_layer is None:
continue
name = self._rewrite_spec_layer_name(spec_layer, name)
for param_name, weight_name, shard_id in _STEP5_WEIGHT_SHARDS:
if weight_name not in name or "experts" in name or "moe" in name:
continue
replaced_name = name.replace(weight_name, param_name)
if replaced_name not in params_dict:
continue
name = replaced_name
param = params_dict[name]
weight_loader = param.weight_loader
weight_loader(param, loaded_weight, shard_id)
break
else:
# Skip loading extra bias for GPTQ models.
if (
name.endswith(".bias")
and name not in params_dict
or "tok_embeddings" in name
):
continue
if spec_layer is not None and ".transformer." in name:
name = name.replace(".transformer.", ".")
if "shared_head" in name:
name = name.replace("shared_head.output", "shared_head.head")
if "embed_tokens" in name:
name = "model.embed_tokens.weight"
if name not in params_dict:
logger.warning_once(
"Skipping unexpected MTP checkpoint weight: %s", name
)
continue
param = params_dict[name]
weight_loader = getattr(param, "weight_loader", default_weight_loader)
weight_loader(param, loaded_weight)
loaded_params.add(name)
if validate_completeness:
loaded_params.update(_validate_fused_qkv_indexer_weights(self))
missing_params = set(params_dict).difference(loaded_params)
if missing_params and validate_completeness:
param_name_example = min(missing_params)
raise RuntimeError(
"Some parameters like "
f"{param_name_example} are not in the checkpoint and will falsely "
"use random initialization"
)
return loaded_params
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
"""
spec_layer_weight_names = [
"embed_tokens",
"enorm",
"hnorm",
"eh_proj",
"shared_head",
]
spec_layer_weight = False
for weight_name in spec_layer_weight_names:
if weight_name in name:
spec_layer_weight = True
break
if not spec_layer_weight:
name = name.replace(
f"model.layers.{spec_layer}.", f"model.layers.{spec_layer}.mtp_block."
)
return name