Reduce AutoWeightsLoader kwargs (#53106)

Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
This commit is contained in:
Harry Mellor
2026-08-20 15:40:11 +00:00
committed by GitHub
parent 2dd17225c6
commit 1fe3a1571a
84 changed files with 577 additions and 568 deletions
@@ -104,7 +104,7 @@ def test_inkling_mapper_maps_modelopt_exclusions() -> None:
)
quant_config.apply_vllm_mapper(
_TmlForCausalLMBase.hf_to_vllm_mapper.get_unstacked_mapper()
_TmlForCausalLMBase.hf_to_vllm_mapper.get_rename_mapper()
)
assert quant_config.is_layer_excluded("model.layers.2.mlp.experts")
+3 -5
View File
@@ -55,11 +55,9 @@ def test_dspark_mla_checkpoint_weight_mapping(checkpoint_name, runtime_name, sha
def test_dspark_mla_shares_frozen_target_weights_and_skips_training_head():
assert not K3DSparkForCausalLM.has_own_embed_tokens
assert not K3DSparkForCausalLM.has_own_lm_head
assert set(K3DSparkForCausalLM.checkpoint_skip_substrs) == {
"confidence_head",
"embed_tokens",
"lm_head",
}
mapper = K3DSparkForCausalLM.hf_to_vllm_mapper
for name in ("confidence_head.weight", "embed_tokens.weight", "lm_head.weight"):
assert mapper._map_name(name) is None
@pytest.mark.cpu_test
+28 -76
View File
@@ -2,6 +2,7 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import pytest
import regex as re
import torch
from vllm.model_executor.layers.vocab_parallel_embedding import (
@@ -10,6 +11,7 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
)
from vllm.model_executor.models.utils import (
AutoWeightsLoader,
WeightsMapper,
_merge_multimodal_embeddings,
)
from vllm.platforms import current_platform
@@ -93,82 +95,6 @@ def test_module_with_child_containing_batchnorm_can_autoload():
assert new_mod.nested_mod.bn.num_batches_tracked.item() == 1
@pytest.mark.cpu_test
def test_module_skip_prefix():
"""Ensure the auto weight loader can skip prefix."""
mod = ModuleWithNestedBatchNorm()
# Run some data through the module with batchnorm
mod(torch.Tensor([[1, 2], [3, 4]]))
# Try to load the weights to a new instance
def weight_generator():
# weights needed to be filtered out
redundant_weights = {
"prefix.bn.weight": torch.Tensor([1, 2]),
"prefix.bn.bias": torch.Tensor([3, 4]),
}
yield from (mod.state_dict() | redundant_weights).items()
new_mod = ModuleWithNestedBatchNorm()
assert not torch.all(
new_mod.nested_mod.bn.running_mean == mod.nested_mod.bn.running_mean
)
assert not torch.all(
new_mod.nested_mod.bn.running_var == mod.nested_mod.bn.running_var
)
assert new_mod.nested_mod.bn.num_batches_tracked.item() == 0
loader = AutoWeightsLoader(new_mod, skip_prefixes=["prefix."])
loader.load_weights(weight_generator())
# Ensure the stats are updated
assert torch.all(
new_mod.nested_mod.bn.running_mean == mod.nested_mod.bn.running_mean
)
assert torch.all(new_mod.nested_mod.bn.running_var == mod.nested_mod.bn.running_var)
assert new_mod.nested_mod.bn.num_batches_tracked.item() == 1
@pytest.mark.cpu_test
def test_module_skip_substr():
"""Ensure the auto weight loader can skip prefix."""
mod = ModuleWithNestedBatchNorm()
# Run some data through the module with batchnorm
mod(torch.Tensor([[1, 2], [3, 4]]))
# Try to load the weights to a new instance
def weight_generator():
# weights needed to be filtered out
redundant_weights = {
"nested_mod.0.substr.weight": torch.Tensor([1, 2]),
"nested_mod.0.substr.bias": torch.Tensor([3, 4]),
"nested_mod.substr.weight": torch.Tensor([1, 2]),
"nested_mod.substr.bias": torch.Tensor([3, 4]),
}
yield from (mod.state_dict() | redundant_weights).items()
new_mod = ModuleWithNestedBatchNorm()
assert not torch.all(
new_mod.nested_mod.bn.running_mean == mod.nested_mod.bn.running_mean
)
assert not torch.all(
new_mod.nested_mod.bn.running_var == mod.nested_mod.bn.running_var
)
assert new_mod.nested_mod.bn.num_batches_tracked.item() == 0
loader = AutoWeightsLoader(new_mod, skip_substrs=["substr."])
loader.load_weights(weight_generator())
# Ensure the stats are updated
assert torch.all(
new_mod.nested_mod.bn.running_mean == mod.nested_mod.bn.running_mean
)
assert torch.all(new_mod.nested_mod.bn.running_var == mod.nested_mod.bn.running_var)
assert new_mod.nested_mod.bn.num_batches_tracked.item() == 1
VOCAB_SIZE = 16
HIDDEN_SIZE = 2
@@ -265,3 +191,29 @@ def test_merge_multimodal_embeddings_no_sync():
_merge_multimodal_embeddings(
inputs_embeds, multimodal_embeddings, is_multimodal
)
@pytest.mark.cpu_test
def test_get_rename_mapper_keeps_only_renames():
"""`None` means "do not load", which is meaningless to the consumers of
this mapper (LoRA name parsing, quantization config layer lists), and
applying it would silently shrink their lists."""
mapper = WeightsMapper(
orig_to_new_regex={re.compile(r"^drop_regex\."): None},
orig_to_new_substr={"drop_substr": None, "keep_substr": "kept"},
orig_to_new_stacked={".q_proj": (".qkv_proj", "q")},
orig_to_new_prefix={"drop_prefix.": None, "keep_prefix.": "kept."},
orig_to_new_suffix={".drop_suffix": None},
)
renames = mapper.get_rename_mapper()
assert renames.orig_to_new_regex == {}
assert renames.orig_to_new_substr == {"keep_substr": "kept"}
assert renames.orig_to_new_stacked == {}
assert renames.orig_to_new_prefix == {"keep_prefix.": "kept."}
assert renames.orig_to_new_suffix == {}
# Names the full mapper drops now survive unchanged.
for name in ("drop_regex.w", "drop_substr.w", "drop_prefix.w", "w.drop_suffix"):
assert mapper._map_name(name) is None
assert renames._map_name(name) == name
+2 -2
View File
@@ -214,9 +214,9 @@ def test_modelopt_mixed_precision_composes_gemma4_mappers():
)
config.apply_vllm_mapper(
Gemma4ForConditionalGeneration.hf_to_vllm_mapper.get_unstacked_mapper()
Gemma4ForConditionalGeneration.hf_to_vllm_mapper.get_rename_mapper()
)
config.apply_vllm_mapper(Gemma4ForCausalLM.hf_to_vllm_mapper.get_unstacked_mapper())
config.apply_vllm_mapper(Gemma4ForCausalLM.hf_to_vllm_mapper.get_rename_mapper())
expected_prefix = "language_model.model.layers.0.moe.experts"
assert set(config.quantized_layers) == {
+4 -5
View File
@@ -128,14 +128,13 @@ class WorkerLoRAManager:
# loading weights, throwing an exception if validation fails.
peft_helper.validate_legal(self.lora_config)
# For some models like Qwen2VL, we need to use hf_to_vllm_mapper
# to ensure correct loading of lora weights. Drop the QKV/MLP fusion
# substr maps so constituent names (e.g. `q_proj`) survive for the
# LoRA manager to pack, while keeping genuine renames/prefixes.
# For some models like Qwen2VL, we need to use hf_to_vllm_mapper to ensure
# correct loading of lora weights. We only need to know about renames for
# this, so we use get_rename_mapper() to ignore stacking and deletions.
model = self._adapter_manager.model
hf_to_vllm_mapper = getattr(model, "hf_to_vllm_mapper", None)
if hf_to_vllm_mapper is not None:
hf_to_vllm_mapper = hf_to_vllm_mapper.get_unstacked_mapper()
hf_to_vllm_mapper = hf_to_vllm_mapper.get_rename_mapper()
# Get model-defined prefixes to skip during LoRA loading.
lora_skip_prefixes = getattr(model, "lora_skip_prefixes", None)
+1 -1
View File
@@ -281,6 +281,6 @@ def configure_quant_config(
# pass mappings by reference to quant_config
if hf_to_vllm_mapper is not None:
quant_config.apply_vllm_mapper(hf_to_vllm_mapper.get_unstacked_mapper())
quant_config.apply_vllm_mapper(hf_to_vllm_mapper.get_rename_mapper())
if packed_mapping is not None:
quant_config.packed_modules_mapping = packed_mapping
+8 -5
View File
@@ -236,7 +236,8 @@ def _create_pooling_model_cls(orig_cls: type[_T]) -> type[_T]:
def default_load_weights(weights):
loader = AutoWeightsLoader(self)
return loader.load_weights(weights)
mapper = getattr(self, "hf_to_vllm_mapper", None)
return loader.load_weights(weights, mapper=mapper)
load_weights = getattr(super(), "load_weights", default_load_weights)
return load_weights(mapped_weights)
@@ -584,10 +585,11 @@ def load_weights_using_from_2_way_softmax(
)
loaded_weights.add(score_weight_name)
lm_head_name = "lm_head.weight"
lm_head_name: str | None = "lm_head.weight"
if hf_to_vllm_mapper := getattr(model, "hf_to_vllm_mapper", None):
lm_head_name = hf_to_vllm_mapper._map_name(lm_head_name)
loaded_weights.discard(lm_head_name)
if lm_head_name is not None:
loaded_weights.discard(lm_head_name)
return loaded_weights
@@ -649,10 +651,11 @@ def load_weights_no_post_processing(model, weights: Iterable[tuple[str, torch.Te
)
loaded_weights.add(score_weight_name)
lm_head_name = "lm_head.weight"
lm_head_name: str | None = "lm_head.weight"
if hf_to_vllm_mapper := getattr(model, "hf_to_vllm_mapper", None):
lm_head_name = hf_to_vllm_mapper._map_name(lm_head_name)
loaded_weights.discard(lm_head_name)
if lm_head_name is not None:
loaded_weights.discard(lm_head_name)
return loaded_weights
+6 -9
View File
@@ -219,6 +219,11 @@ class AIMv2Model(torch.nn.Module):
require_post_norm=require_post_norm,
prefix=f"{prefix}.trunk",
)
# post_trunk_norm is optional (absent for clip-skip backbones).
if self.trunk.post_trunk_norm is None:
self.hf_to_vllm_mapper = self.hf_to_vllm_mapper | WeightsMapper(
orig_to_new_prefix={"trunk.post_trunk_norm.": None}
)
def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
x = self.preprocessor(pixel_values)
@@ -227,13 +232,5 @@ class AIMv2Model(torch.nn.Module):
return x
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
loader = AutoWeightsLoader(
self,
# post_trunk_norm is optional (absent for clip-skip backbones).
skip_prefixes=(
["trunk.post_trunk_norm."]
if self.trunk.post_trunk_norm is None
else None
),
)
loader = AutoWeightsLoader(self)
return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
+5 -2
View File
@@ -285,7 +285,10 @@ class ArceeForCausalLM(
".q_proj": (".qkv_proj", "q"),
".k_proj": (".qkv_proj", "k"),
".v_proj": (".qkv_proj", "v"),
}
},
orig_to_new_substr={
"gate_proj": None,
},
)
# Map fused module names to their submodule components
# (for quantization and LoRA)
@@ -356,7 +359,7 @@ class ArceeForCausalLM(
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
"""Load weights into the model (delegates to inner model and handles
tied embeddings)."""
loader = AutoWeightsLoader(self, skip_substrs=["gate_proj"])
loader = AutoWeightsLoader(self)
# AutoWeightLoader handles weight name remapping, including fusing
# separate q_proj, k_proj, v_proj into qkv_proj
return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
+7 -15
View File
@@ -13,18 +13,13 @@ from vllm.config import VllmConfig
from vllm.config.multimodal import BaseDummyOptions, ImageDummyOptions
from vllm.inputs import MultiModalDataDict
from vllm.model_executor.layers.activation import get_act_fn
from vllm.model_executor.layers.fused_moe import (
FusedMoEFactory,
)
from vllm.model_executor.layers.fused_moe import FusedMoEFactory
from vllm.model_executor.layers.linear import ColumnParallelLinear, RowParallelLinear
from vllm.model_executor.layers.logits_processor import LogitsProcessor
from vllm.model_executor.layers.quantization import QuantizationConfig
from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
from vllm.multimodal import MULTIMODAL_REGISTRY
from vllm.multimodal.inputs import (
MultiModalFieldConfig,
MultiModalKwargsItems,
)
from vllm.multimodal.inputs import MultiModalFieldConfig, MultiModalKwargsItems
from vllm.multimodal.parse import MultiModalDataItems
from vllm.multimodal.processing import (
BaseDummyInputsBuilder,
@@ -42,11 +37,7 @@ from .idefics2_vision_model import (
)
from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsQuant
from .llama import LlamaDecoderLayer, LlamaMLP, LlamaModel
from .utils import (
AutoWeightsLoader,
WeightsMapper,
maybe_prefix,
)
from .utils import AutoWeightsLoader, WeightsMapper, maybe_prefix
class AriaImagePixelInputs(TensorSchema):
@@ -88,16 +79,17 @@ class AriaVisionTransformer(Idefics3VisionTransformer, SupportsQuant):
self.post_layernorm = nn.Identity()
hf_to_vllm_mapper = WeightsMapper(
# NOTE: post_layernorm is not used in Aria.
orig_to_new_substr={"post_layernorm": None},
orig_to_new_stacked={
".q_proj": (".qkv_proj", "q"),
".k_proj": (".qkv_proj", "k"),
".v_proj": (".qkv_proj", "v"),
}
},
)
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
# NOTE: post_layernorm is not used in Aria.
loader = AutoWeightsLoader(self, skip_substrs=["post_layernorm"])
loader = AutoWeightsLoader(self)
return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
+3 -9
View File
@@ -338,14 +338,9 @@ class BagelForConditionalGeneration(
The image generation part is not supported in vLLM.
"""
# Weight mapping from HF to vLLM
# pos_embed is handled by the PositionEmbedding module
hf_to_vllm_mapper = WeightsMapper(
orig_to_new_prefix={
"language_model.": "language_model.",
"vit_model.": "vit_model.",
"connector.": "connector.",
"vit_pos_embed.": "vit_pos_embed.",
}
orig_to_new_prefix={"vit_pos_embed.pos_embed": None}
)
@classmethod
@@ -580,6 +575,5 @@ class BagelForConditionalGeneration(
filtered_weights.append((name, tensor))
# Skip vit_pos_embed.pos_embed as it's handled by PositionEmbedding module
loader = AutoWeightsLoader(self, skip_prefixes=["vit_pos_embed.pos_embed"])
loader = AutoWeightsLoader(self)
return loader.load_weights(filtered_weights, mapper=self.hf_to_vllm_mapper)
+10 -3
View File
@@ -2,6 +2,7 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from collections.abc import Iterable, Set
from dataclasses import replace
import torch
from torch import nn
@@ -382,6 +383,7 @@ class BertModel(nn.Module, SupportsQuant):
".self.key": (".self.qkv_proj", "k"),
".self.value": (".self.qkv_proj", "v"),
},
orig_to_new_prefix={"pooler.": None},
)
def __init__(
@@ -416,13 +418,16 @@ class BertModel(nn.Module, SupportsQuant):
return self.encoder(hidden_states)
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
loader = AutoWeightsLoader(self, skip_prefixes=["pooler."])
loader = AutoWeightsLoader(self)
return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
class BertPoolingModel(BertModel):
is_pooling_model = True
# Unlike `BertModel`, this model has a pooler to load weights into.
hf_to_vllm_mapper = replace(BertModel.hf_to_vllm_mapper, orig_to_new_prefix={})
def __init__(
self,
*,
@@ -488,11 +493,13 @@ class BertEmbeddingModel(nn.Module, SupportsQuant):
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
weights_list = list(weights)
orig_to_new_prefix: dict[str, str | None] = {"lm_head.": None}
has_model_prefix = any(name.startswith("model.") for name, _ in weights_list)
if not has_model_prefix:
mapper = WeightsMapper(orig_to_new_prefix={"": "model."})
orig_to_new_prefix[""] = "model."
mapper = WeightsMapper(orig_to_new_prefix=orig_to_new_prefix)
loader = AutoWeightsLoader(self, skip_prefixes=["lm_head."])
loader = AutoWeightsLoader(self)
return loader.load_weights(weights_list, mapper=mapper)
def _build_model(self, vllm_config: VllmConfig, prefix: str = "") -> BertModel:
+4 -4
View File
@@ -313,6 +313,9 @@ class BlipVisionModel(nn.Module, SupportsQuant):
)
else:
self.post_layernorm = None
self.hf_to_vllm_mapper = self.hf_to_vllm_mapper | WeightsMapper(
orig_to_new_prefix={"post_layernorm.": None}
)
def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
hidden_states = self.embeddings(pixel_values)
@@ -324,10 +327,7 @@ class BlipVisionModel(nn.Module, SupportsQuant):
return self.post_layernorm(hidden_states)
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
skip_prefixes: list[str] = []
if self.post_layernorm is None:
skip_prefixes.append("post_layernorm.")
loader = AutoWeightsLoader(self, skip_prefixes=skip_prefixes)
loader = AutoWeightsLoader(self)
# omit layers when num_hidden_layers_override is set
def _filter(ws):
+7 -6
View File
@@ -669,6 +669,9 @@ class CLIPVisionTransformer(nn.Module):
self.post_layernorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)
else:
self.post_layernorm = None
self.hf_to_vllm_mapper = self.hf_to_vllm_mapper | WeightsMapper(
orig_to_new_prefix={"post_layernorm.": None}
)
@property
def dtype(self):
@@ -707,10 +710,7 @@ class CLIPVisionTransformer(nn.Module):
return encoder_outputs
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
skip_prefixes: list[str] = []
if self.post_layernorm is None:
skip_prefixes.append("post_layernorm.")
loader = AutoWeightsLoader(self, skip_prefixes=skip_prefixes)
loader = AutoWeightsLoader(self)
# Drop layers beyond num_hidden_layers_override.
def _filter(ws):
@@ -775,6 +775,8 @@ class CLIPVisionModel(nn.Module):
class CLIPEmbeddingModel(nn.Module, SupportsMultiModal, SupportsQuant):
is_pooling_model = True
hf_to_vllm_mapper = WeightsMapper(orig_to_new_substr={".position_ids": None})
packed_modules_mapping = {"qkv_proj": ["q_proj", "k_proj", "v_proj"]}
@classmethod
@@ -981,8 +983,7 @@ class CLIPEmbeddingModel(nn.Module, SupportsMultiModal, SupportsQuant):
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
loader = AutoWeightsLoader(
self,
skip_substrs=[".position_ids"],
ignore_unexpected_prefixes=["logit_scale."],
)
return loader.load_weights(weights)
return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
+3 -2
View File
@@ -474,7 +474,8 @@ class Cohere2MoeForCausalLM(nn.Module, SupportsPP, SupportsQuant):
".mlp.up_proj": (".mlp.gate_up_proj", 1),
".shared_experts.gate_proj": (".shared_experts.gate_up_proj", 0),
".shared_experts.up_proj": (".shared_experts.gate_up_proj", 1),
}
},
orig_to_new_prefix={"lm_head.": None},
)
packed_modules_mapping = {
"qkv_proj": [
@@ -529,5 +530,5 @@ class Cohere2MoeForCausalLM(nn.Module, SupportsPP, SupportsQuant):
return self.logits_processor(self.model.embed_tokens, hidden_states)
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
loader = AutoWeightsLoader(self, skip_prefixes=["lm_head."])
loader = AutoWeightsLoader(self)
return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
+10 -9
View File
@@ -2014,7 +2014,15 @@ class CohereAsrForConditionalGeneration(
}
hf_to_vllm_mapper = WeightsMapper(
orig_to_new_substr={".fc1.": ".mlp.fc1.", ".fc2.": ".mlp.fc2."}
orig_to_new_substr={
".fc1.": ".mlp.fc1.",
".fc2.": ".mlp.fc2.",
"model.conv.batch_norm.num_batches_tracked": None,
},
orig_to_new_prefix={
"model.preprocessor.featurizer.fb": None,
"model.preprocessor.featurizer.window": None,
},
)
supports_transcription_only = True
@@ -2273,14 +2281,7 @@ class CohereAsrForConditionalGeneration(
return name, loaded_weight
loader = AutoWeightsLoader(
self,
skip_prefixes=[
"model.preprocessor.featurizer.fb",
"model.preprocessor.featurizer.window",
],
skip_substrs=["model.conv.batch_norm.num_batches_tracked"],
)
loader = AutoWeightsLoader(self)
return loader.load_weights(
map(transform, weights), mapper=self.hf_to_vllm_mapper
+6 -8
View File
@@ -22,6 +22,7 @@ from vllm.model_executor.models.commandr import (
from .utils import (
AutoWeightsLoader,
WeightsMapper,
get_draft_quant_config,
maybe_prefix,
process_eagle_weight,
@@ -145,6 +146,10 @@ class EagleCohereForCausalLM(CohereForCausalLM):
# use tied embeddings so these weights are absent from the draft file.
self.has_own_embed_tokens = False
self.has_own_lm_head = False
if self.config.tie_word_embeddings:
self.hf_to_vllm_mapper = self.hf_to_vllm_mapper | WeightsMapper(
orig_to_new_prefix={"model.embed_tokens.": None}
)
target_layer_num = vllm_config.model_config.get_num_layers(
vllm_config.parallel_config
)
@@ -181,14 +186,7 @@ class EagleCohereForCausalLM(CohereForCausalLM):
process_eagle_weight(self, name)
return name, weight
loader = AutoWeightsLoader(
self,
skip_prefixes=(
["lm_head.", "model.embed_tokens."]
if self.config.tie_word_embeddings
else None
),
)
loader = AutoWeightsLoader(self)
loaded_weight_names = loader.load_weights(
map(_track_and_forward, weights), mapper=self.hf_to_vllm_mapper
+5 -4
View File
@@ -358,11 +358,12 @@ class ColBERTJinaRobertaModel(ColBERTMixin, nn.Module):
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
other_weights, colbert_loaded = self._load_colbert_weights(weights)
mapper = WeightsMapper(orig_to_new_prefix={"roberta.": "model."})
# Skip HF pooler weights (model.pooler.*) as they not used in ColBERT
loader = AutoWeightsLoader(self, skip_prefixes=["model.pooler."])
# HF pooler weights are not used in ColBERT
mapper = WeightsMapper(
orig_to_new_prefix={"roberta.": "model.", "model.pooler.": None}
)
loader = AutoWeightsLoader(self)
loaded = loader.load_weights(other_weights, mapper=mapper)
return loaded | colbert_loaded
+2 -4
View File
@@ -140,6 +140,7 @@ class ColQwen3_5Model(
hf_to_vllm_mapper = WeightsMapper(
orig_to_new_prefix={
"language_model.": "language_model.model.",
"mtp.": None,
}
)
@@ -238,10 +239,7 @@ class ColQwen3_5Model(
else:
model_weights.append((name, weight))
loader = AutoWeightsLoader(
self,
skip_prefixes=["mtp."],
)
loader = AutoWeightsLoader(self)
loaded = loader.load_weights(model_weights, mapper=self.hf_to_vllm_mapper)
for name, weight in proj_weights:
+2 -3
View File
@@ -415,6 +415,7 @@ class CohereForCausalLM(nn.Module, SupportsLoRA, SupportsPP, SupportsQuant):
# (e.g. "*.weight_quantizer._double_scale"); drop them before loading.
# See #41925.
orig_to_new_substr={"_quantizer.": None},
orig_to_new_prefix={"lm_head": None},
)
packed_modules_mapping = {
"qkv_proj": ["q_proj", "k_proj", "v_proj"],
@@ -472,7 +473,5 @@ class CohereForCausalLM(nn.Module, SupportsLoRA, SupportsPP, SupportsQuant):
return logits
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
loader = AutoWeightsLoader(
self, skip_prefixes=["lm_head", "rotary_emb.inv_freq"]
)
loader = AutoWeightsLoader(self)
return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
+1 -4
View File
@@ -253,8 +253,5 @@ class EagleDeepseekV3ForCausalLM(DeepseekV3ForCausalLM):
process_eagle_weight(self, name)
return name, loaded_weight
loader = AutoWeightsLoader(
self,
skip_prefixes=None,
)
loader = AutoWeightsLoader(self)
loader.load_weights(map(transform, weights))
@@ -384,18 +384,17 @@ class Eagle3DeepseekV2ForCausalLM(LocalArgmaxMixin, DeepseekV2ForCausalLM):
model_weights[name] = loaded_weight
process_eagle_weight(self, name)
skip_substrs = []
orig_to_new_substr: dict[str, str | None] = {}
if not includes_draft_id_mapping:
skip_substrs.append("draft_id_to_target_id")
orig_to_new_substr["draft_id_to_target_id"] = None
if not includes_embed_tokens:
skip_substrs.append("embed_tokens")
orig_to_new_substr["embed_tokens"] = None
loader = AutoWeightsLoader(
self,
skip_prefixes=None,
skip_substrs=skip_substrs,
loader = AutoWeightsLoader(self)
loader.load_weights(
model_weights.items(),
mapper=WeightsMapper(orig_to_new_substr=orig_to_new_substr),
)
loader.load_weights(model_weights.items())
# Aliases for compatibility
+3 -6
View File
@@ -403,7 +403,8 @@ class Ernie4_5_MoeModel(nn.Module):
".mlp.up_proj": (".mlp.gate_up_proj", 1),
".shared_experts.gate_proj": (".shared_experts.gate_up_proj", 0),
".shared_experts.up_proj": (".shared_experts.gate_up_proj", 1),
}
},
orig_to_new_substr={"mtp": None},
)
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
@@ -497,11 +498,7 @@ class Ernie4_5_MoeModel(nn.Module):
yield name, loaded_weight
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
loader = AutoWeightsLoader(
self,
skip_substrs=["mtp"],
ignore_unexpected_suffixes=[".bias", "_bias"],
)
loader = AutoWeightsLoader(self)
return loader.load_weights(
self._preprocess(weights), mapper=self.hf_to_vllm_mapper
)
+11 -5
View File
@@ -52,7 +52,12 @@ from vllm.multimodal import MULTIMODAL_REGISTRY
from .qwen2_vl import Qwen2VLDummyInputsBuilder as Exaone4_5_DummyInputsBuilder
from .qwen2_vl import Qwen2VLMultiModalProcessor as Exaone4_5_MultiModalProcessor
from .utils import AutoWeightsLoader, init_vllm_registered_model, maybe_prefix
from .utils import (
AutoWeightsLoader,
WeightsMapper,
init_vllm_registered_model,
maybe_prefix,
)
logger = init_logger(__name__)
@@ -317,6 +322,10 @@ class Exaone4_5_ProcessingInfo(Qwen2VLProcessingInfo):
dummy_inputs=Exaone4_5_DummyInputsBuilder,
)
class Exaone4_5_ForConditionalGeneration(Qwen2_5_VLForConditionalGeneration):
hf_to_vllm_mapper = Qwen2_5_VLForConditionalGeneration.hf_to_vllm_mapper | (
WeightsMapper(orig_to_new_prefix={"mtp.": None})
)
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
nn.Module.__init__(self)
@@ -353,10 +362,7 @@ class Exaone4_5_ForConditionalGeneration(Qwen2_5_VLForConditionalGeneration):
)
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
loader = AutoWeightsLoader(
self,
skip_prefixes=(["mtp."]),
)
loader = AutoWeightsLoader(self)
return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
@classmethod
+2 -2
View File
@@ -471,7 +471,8 @@ class ExaoneMoeForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
".mlp.up_proj": (".mlp.gate_up_proj", 1),
".shared_experts.gate_proj": (".shared_experts.gate_up_proj", 0),
".shared_experts.up_proj": (".shared_experts.gate_up_proj", 1),
}
},
orig_to_new_prefix={"mtp.": None},
)
packed_modules_mapping = {
"qkv_proj": [
@@ -560,7 +561,6 @@ class ExaoneMoeForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
loader = AutoWeightsLoader(
self,
skip_prefixes=["mtp."],
# Skip loading extra parameters for GPTQ/modelopt models.
ignore_unexpected_suffixes=[
".bias",
+5 -4
View File
@@ -332,7 +332,10 @@ class FireRedASR2ForConditionalGeneration(
"net.0": "pre_layer_norm",
"net.1": "linear_expand",
"net.4": "linear_project",
}
},
orig_to_new_prefix={
"model.encoder.audio_encoder.positional_encoding.pe": None,
},
)
supports_transcription_only = True
@@ -477,8 +480,6 @@ class FireRedASR2ForConditionalGeneration(
return logits
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
loader = AutoWeightsLoader(
self, skip_prefixes=["model.encoder.audio_encoder.positional_encoding.pe"]
)
loader = AutoWeightsLoader(self)
return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
+10 -12
View File
@@ -589,7 +589,15 @@ class FireRedLIDForConditionalGeneration(
"net.0": "pre_layer_norm",
"net.1": "linear_expand",
"net.4": "linear_project",
}
},
orig_to_new_prefix={
# Position encoding buffers are rebuilt at init, and the output
# projection is tied to the embedding.
"model.encoder.positional_encoding.pe": None,
"model.decoder.positional_encoding.pe": None,
"model.decoder.tgt_word_prj.weight": None,
"proj_out.": None,
},
)
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
@@ -778,15 +786,5 @@ class FireRedLIDForConditionalGeneration(
return text.strip()
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
loader = AutoWeightsLoader(
self,
skip_prefixes=[
# Position encoding buffers are rebuilt at init
"model.encoder.positional_encoding.pe",
"model.decoder.positional_encoding.pe",
# Tied output projection (shared with embedding)
"model.decoder.tgt_word_prj.weight",
"proj_out.",
],
)
loader = AutoWeightsLoader(self)
return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
+4 -2
View File
@@ -760,6 +760,8 @@ class FunAudioChatMultiModalProcessor(
dummy_inputs=FunAudioChatDummyInputsBuilder,
)
class FunAudioChatForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP):
hf_to_vllm_mapper = WeightsMapper(orig_to_new_prefix={"audio_invert_tower.": None})
@classmethod
def get_placeholder_str(cls, modality: str, i: int) -> str | None:
if modality.startswith("audio"):
@@ -969,5 +971,5 @@ class FunAudioChatForConditionalGeneration(nn.Module, SupportsMultiModal, Suppor
return self.language_model.compute_logits(hidden_states)
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
loader = AutoWeightsLoader(self, skip_prefixes=["audio_invert_tower."])
return loader.load_weights(weights)
loader = AutoWeightsLoader(self)
return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
+10 -7
View File
@@ -1056,6 +1056,14 @@ class Gemma3nTextModel(nn.Module, SupportsQuant):
class Gemma3nForCausalLM(nn.Module):
hf_to_vllm_mapper = WeightsMapper(
orig_to_new_substr={
"embed_audio.": None,
"embed_vision.": None,
"audio_tower.": None,
"vision_tower.": None,
}
)
packed_modules_mapping = {
"qkv_proj": [
"q_proj",
@@ -1112,10 +1120,5 @@ class Gemma3nForCausalLM(nn.Module):
return logits
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
loader = AutoWeightsLoader(
self,
skip_substrs=(
["embed_audio.", "embed_vision.", "audio_tower.", "vision_tower."]
),
)
return loader.load_weights(weights)
loader = AutoWeightsLoader(self)
return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
+13 -10
View File
@@ -1693,13 +1693,16 @@ class Gemma4ForCausalLM(
yield name, weight
# Skip multimodal weights — handled by the multimodal wrapper.
skip = [
"audio_tower.",
"vision_tower.",
"embed_audio.",
"embed_vision.",
]
loader = AutoWeightsLoader(self, skip_substrs=skip)
return loader.load_weights(_weight_iterator())
# Drop multimodal weights, which are handled by the multimodal wrapper.
# `_weight_iterator` already applies this model's renames by hand, so
# `hf_to_vllm_mapper` is deliberately not passed here.
mapper = WeightsMapper(
orig_to_new_substr={
"audio_tower.": None,
"vision_tower.": None,
"embed_audio.": None,
"embed_vision.": None,
}
)
loader = AutoWeightsLoader(self)
return loader.load_weights(_weight_iterator(), mapper=mapper)
+13 -14
View File
@@ -45,11 +45,7 @@ from vllm.v1.attention.backend import AttentionType
from .interfaces import SupportsLoRA, SupportsPP
from .llama import LlamaMLP as Glm4MLP
from .llama import LlamaModel
from .utils import (
AutoWeightsLoader,
PPMissingLayer,
maybe_prefix,
)
from .utils import AutoWeightsLoader, PPMissingLayer, WeightsMapper, maybe_prefix
class Glm4Attention(nn.Module):
@@ -269,6 +265,16 @@ class Glm4ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
self.model.make_empty_intermediate_tensors
)
# Drop the speculative (MTP) layers, which are loaded by the draft
# model instead. They are appended after the main decoder layers.
num_nextn_layers = getattr(config, "num_nextn_predict_layers", 0)
self.hf_to_vllm_mapper = WeightsMapper(
orig_to_new_prefix={
f"model.layers.{config.num_hidden_layers + i}.": None
for i in range(num_nextn_layers)
}
)
def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
return self.model.embed_input_ids(input_ids)
@@ -292,12 +298,5 @@ class Glm4ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
return logits
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
# Skip the speculative (MTP) layers, which are loaded by the
# draft model instead.
num_nextn_layers = getattr(self.config, "num_nextn_predict_layers", 0)
skip_prefixes = [
f"model.layers.{self.config.num_hidden_layers + i}."
for i in range(num_nextn_layers)
]
loader = AutoWeightsLoader(self, skip_prefixes=skip_prefixes)
return loader.load_weights(weights)
loader = AutoWeightsLoader(self)
return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
+6 -3
View File
@@ -916,6 +916,10 @@ class GlmAsrForConditionalGeneration(
):
supported_languages = ISO639_1_SUPPORTED_LANGS
hf_to_vllm_mapper = WeightsMapper(
orig_to_new_prefix={"audio_tower.embed_positions": None}
)
packed_modules_mapping = {
"qkv_proj": ["q_proj", "k_proj", "v_proj"],
"gate_up_proj": ["gate_proj", "up_proj"],
@@ -1083,9 +1087,8 @@ class GlmAsrForConditionalGeneration(
return self.language_model.compute_logits(hidden_states)
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
skip_prefixes = ["audio_tower.embed_positions"]
loader = AutoWeightsLoader(self, skip_prefixes=skip_prefixes)
return loader.load_weights(weights)
loader = AutoWeightsLoader(self)
return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
@classmethod
def _get_audio_token(cls, model_config: ModelConfig) -> str:
+9 -4
View File
@@ -52,6 +52,7 @@ from vllm.sequence import IntermediateTensors
from .interfaces import SupportsCrossEncoding, SupportsPP
from .utils import (
AutoWeightsLoader,
WeightsMapper,
make_empty_intermediate_tensors_factory,
make_layers,
maybe_prefix,
@@ -182,6 +183,11 @@ class GPT2Block(nn.Module):
@support_torch_compile
class GPT2Model(nn.Module):
# Drop attention mask buffers; NOTE: "c_attn.bias" must not be dropped.
hf_to_vllm_mapper = WeightsMapper(
orig_to_new_substr={".attn.bias": None, ".attn.masked_bias": None}
)
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
super().__init__()
@@ -253,11 +259,10 @@ class GPT2Model(nn.Module):
yield name, loaded_weight
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
# Skip attention mask buffers; NOTE: "c_attn.bias" must not be skipped.
loader = AutoWeightsLoader(
self, skip_substrs=[".attn.bias", ".attn.masked_bias"]
loader = AutoWeightsLoader(self)
return loader.load_weights(
self._transpose_conv1d(weights), mapper=self.hf_to_vllm_mapper
)
return loader.load_weights(self._transpose_conv1d(weights))
class GPT2LMHeadModel(nn.Module, SupportsPP):
+6 -2
View File
@@ -243,7 +243,11 @@ class GPTJForCausalLM(nn.Module, SupportsPP):
".q_proj": (".qkv_proj", "q"),
".k_proj": (".qkv_proj", "k"),
".v_proj": (".qkv_proj", "v"),
}
},
orig_to_new_substr={
"attn.bias": None,
"attn.masked_bias": None,
},
)
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
@@ -291,5 +295,5 @@ class GPTJForCausalLM(nn.Module, SupportsPP):
return logits
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
loader = AutoWeightsLoader(self, skip_substrs=["attn.bias", "attn.masked_bias"])
loader = AutoWeightsLoader(self)
return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
+8 -3
View File
@@ -48,6 +48,7 @@ from vllm.sequence import IntermediateTensors
from .interfaces import SupportsPP
from .utils import (
AutoWeightsLoader,
WeightsMapper,
make_empty_intermediate_tensors_factory,
make_layers,
maybe_prefix,
@@ -196,6 +197,10 @@ class GPTNeoXLayer(nn.Module):
@support_torch_compile
class GPTNeoXModel(nn.Module):
hf_to_vllm_mapper = WeightsMapper(
orig_to_new_substr={"attention.bias": None, "attention.masked_bias": None}
)
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
super().__init__()
@@ -262,10 +267,10 @@ class GPTNeoXModel(nn.Module):
yield name, loaded_weight
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
loader = AutoWeightsLoader(
self, skip_substrs=["attention.bias", "attention.masked_bias"]
loader = AutoWeightsLoader(self)
return loader.load_weights(
self._repack_qkv(weights), mapper=self.hf_to_vllm_mapper
)
return loader.load_weights(self._repack_qkv(weights))
class GPTNeoXForCausalLM(nn.Module, SupportsPP):
@@ -19,7 +19,6 @@
"""PyTorch Idefics2 model."""
from collections.abc import Iterable
from typing import ClassVar
import torch
from torch import nn
@@ -360,7 +359,7 @@ class Idefics2Encoder(nn.Module):
class Idefics2VisionTransformer(nn.Module):
hf_to_vllm_mapper: ClassVar[WeightsMapper] = WeightsMapper(
hf_to_vllm_mapper = WeightsMapper(
orig_to_new_stacked={
".q_proj": (".qkv_proj", "q"),
".k_proj": (".qkv_proj", "k"),
@@ -408,6 +407,13 @@ class Idefics2VisionTransformer(nn.Module):
if require_post_norm
else nn.Identity()
)
# head is a pooling header absent from this model.
orig_to_new_prefix: dict[str, str | None] = {"head.": None}
if not require_post_norm:
orig_to_new_prefix["post_layernorm."] = None
self.hf_to_vllm_mapper = self.hf_to_vllm_mapper | WeightsMapper(
orig_to_new_prefix=orig_to_new_prefix
)
def get_input_embeddings(self):
return self.embeddings
@@ -469,11 +475,7 @@ class Idefics2VisionTransformer(nn.Module):
return last_hidden_state
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
# head is a pooling header absent from this model.
skip_prefixes = ["head."]
if not self.require_post_norm:
skip_prefixes.append("post_layernorm.")
loader = AutoWeightsLoader(self, skip_prefixes=skip_prefixes)
loader = AutoWeightsLoader(self)
layer_count = len(self.encoder.layers)
+2 -3
View File
@@ -1150,7 +1150,7 @@ class SupportsLateInteraction(Protocol):
class SupportsQuant:
"""The interface required for all models that support quantization."""
hf_to_vllm_mapper: ClassVar["WeightsMapper | None"] = None
hf_to_vllm_mapper: "WeightsMapper | None" = None
packed_modules_mapping: ClassVar[dict[str, list[str]]]
quant_config: QuantizationConfig | None = None
@@ -1184,8 +1184,7 @@ class SupportsQuant:
if self.quant_config is None:
return
if (hf_to_vllm_mapper := self.hf_to_vllm_mapper) is not None:
unstacked_mapper = hf_to_vllm_mapper.get_unstacked_mapper()
self.quant_config.apply_vllm_mapper(unstacked_mapper)
self.quant_config.apply_vllm_mapper(hf_to_vllm_mapper.get_rename_mapper())
if packed_modules_mapping := getattr(self, "packed_modules_mapping", None):
self.quant_config.packed_modules_mapping.update(packed_modules_mapping)
+9 -8
View File
@@ -616,17 +616,18 @@ class InternS1ProForConditionalGeneration(
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
"""load weights"""
skip_prefixes = ["model.time_series."]
orig_to_new_prefix: dict[str, str | None] = {
"model.visual.": "visual.",
"lm_head.": "language_model.lm_head.",
"model.language_model.": "language_model.model.",
"model.time_series.": None,
}
if self.visual is None:
skip_prefixes.append("visual.")
orig_to_new_prefix["visual."] = None
# FIXME(Isotr0py): See if we can avoid tighing FoPE to PP layers
weights_mapper = WeightsMapper(
orig_to_new_prefix={
"model.visual.": "visual.",
"lm_head.": "language_model.lm_head.",
"model.language_model.": "language_model.model.",
},
orig_to_new_prefix=orig_to_new_prefix,
orig_to_new_suffix=self.get_frope_params_map(),
)
loader = AutoWeightsLoader(self, skip_prefixes=skip_prefixes)
loader = AutoWeightsLoader(self)
return loader.load_weights(weights, mapper=weights_mapper)
@@ -13,7 +13,7 @@ from .qwen3_vl import (
Qwen3VLMultiModalProcessor,
Qwen3VLProcessingInfo,
)
from .utils import AutoWeightsLoader
from .utils import AutoWeightsLoader, WeightsMapper
class InternS2PreviewProcessingInfo(Qwen3VLProcessingInfo):
@@ -30,9 +30,13 @@ class InternS2PreviewProcessingInfo(Qwen3VLProcessingInfo):
dummy_inputs=Qwen3VLDummyInputsBuilder,
)
class InternS2PreviewForConditionalGeneration(Qwen3_5MoeForConditionalGeneration):
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
loader = AutoWeightsLoader(
self,
skip_prefixes=["mtp.", "model.time_series.", "time_series."],
# `mtp.` is already dropped by `Qwen3_5ForConditionalGeneration`.
hf_to_vllm_mapper = Qwen3_5MoeForConditionalGeneration.hf_to_vllm_mapper | (
WeightsMapper(
orig_to_new_prefix={"model.time_series.": None, "time_series.": None}
)
)
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
loader = AutoWeightsLoader(self)
return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
+27 -18
View File
@@ -59,7 +59,12 @@ from .interfaces import (
SupportsMultiModal,
SupportsPP,
)
from .utils import AutoWeightsLoader, init_vllm_registered_model, maybe_prefix
from .utils import (
AutoWeightsLoader,
WeightsMapper,
init_vllm_registered_model,
maybe_prefix,
)
class InternVLImagePixelInputs(TensorSchema):
@@ -554,6 +559,25 @@ class InternVLChatModel(
):
supports_encoder_tp_data = True
hf_to_vllm_mapper = WeightsMapper(
orig_to_new_prefix=dict.fromkeys(
[
"action_embed",
"temporal_embed",
"track_embed",
"track_embed_decoder",
"box_token",
"cg_criterion",
"cg_model",
"loc_encoder",
"loc_decoder",
"sam",
"temporal_token",
"track_token",
]
)
)
@classmethod
def get_placeholder_str(cls, modality: str, i: int) -> str | None:
if modality.startswith("image"):
@@ -865,23 +889,8 @@ class InternVLChatModel(
return self.language_model.compute_logits(hidden_states)
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
# unused modules appear in OpenGVLab/InternVideo2_5_Chat_8B
skip_prefixes = [
"action_embed",
"temporal_embed",
"track_embed",
"track_embed_decoder",
"box_token",
"cg_criterion",
"cg_model",
"loc_encoder",
"loc_decoder",
"sam",
"temporal_token",
"track_token",
]
loader = AutoWeightsLoader(self, skip_prefixes=skip_prefixes)
return loader.load_weights(weights)
loader = AutoWeightsLoader(self)
return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
def get_mm_mapping(self) -> MultiModelKeys:
"""
+4 -2
View File
@@ -43,6 +43,8 @@ logger = logging.getLogger(__name__)
class JinaForRanking(nn.Module, SupportsLateInteraction):
is_pooling_model = True
hf_to_vllm_mapper = WeightsMapper(orig_to_new_prefix={"lm_head.": None})
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
super().__init__()
config = vllm_config.model_config.hf_config
@@ -87,8 +89,8 @@ class JinaForRanking(nn.Module, SupportsLateInteraction):
return hidden_states
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
loader = AutoWeightsLoader(self, skip_prefixes=(["lm_head."]))
return loader.load_weights(weights)
loader = AutoWeightsLoader(self)
return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
class JinaForRankingPool(StepPool):
+3 -2
View File
@@ -721,7 +721,8 @@ class KeyeSiglipVisionModel(nn.Module):
".q_proj": (".qkv_proj", "q"),
".k_proj": (".qkv_proj", "k"),
".v_proj": (".qkv_proj", "v"),
}
},
orig_to_new_prefix={"vision_model.head.": None},
)
def __init__(
@@ -782,7 +783,7 @@ class KeyeSiglipVisionModel(nn.Module):
)
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
loader = AutoWeightsLoader(self, skip_prefixes=["vision_model.head."])
loader = AutoWeightsLoader(self)
return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
+8 -15
View File
@@ -380,6 +380,12 @@ class KimiAudioForConditionalGeneration(
"model.embed_tokens.": "language_model.model.embed_tokens.",
"model.norm.": "language_model.model.norm.",
"lm_head.": "language_model.lm_head.",
# MIMO/TTS weights and any `model.` residue no rule above
# claimed: this model only does ASR (speech-to-text).
"model.": None,
"mimo_layers.": None,
"mimo_output.": None,
"mimo_norm.": None,
},
orig_to_new_substr={
".fc1.": ".mlp.fc1.",
@@ -573,21 +579,8 @@ class KimiAudioForConditionalGeneration(
return self.language_model.compute_logits(hidden_states)
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
"""Load weights, skipping MIMO layers (TTS-only) for ASR."""
# Filter out MIMO/TTS weights since we only do ASR (speech-to-text)
skipped_patterns = [
# Audio tower
"model.",
# MIMO/TTS
"mimo_layers.",
"mimo_output.",
"mimo_norm.",
]
# Load main model weights (LLM + projector) with mapper
loader = AutoWeightsLoader(self, skip_prefixes=skipped_patterns)
loaded = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
return loaded
loader = AutoWeightsLoader(self)
return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
@classmethod
def get_speech_to_text_config(
+5 -4
View File
@@ -484,6 +484,10 @@ class Siglip2Model(torch.nn.Module):
require_post_norm=require_post_norm,
prefix=maybe_prefix(prefix, "vision_model"),
)
if self.vision_model.post_layernorm is None:
self.hf_to_vllm_mapper = self.hf_to_vllm_mapper | WeightsMapper(
orig_to_new_prefix={"vision_model.post_layernorm.": None}
)
def forward(
self,
@@ -510,10 +514,7 @@ class Siglip2Model(torch.nn.Module):
)
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
skip_prefixes = []
if self.vision_model.post_layernorm is None:
skip_prefixes.append("vision_model.post_layernorm.")
loader = AutoWeightsLoader(self, skip_prefixes=skip_prefixes)
loader = AutoWeightsLoader(self)
# Drop layers omitted by num_hidden_layers_override.
layer_count = len(self.vision_model.encoder.layers)
+1 -2
View File
@@ -26,7 +26,6 @@
from collections.abc import Iterable
from itertools import islice
from typing import ClassVar
import torch
from torch import nn
@@ -343,7 +342,7 @@ class LlamaDecoderLayer(nn.Module):
},
)
class LlamaModel(nn.Module, EagleModelMixin):
hf_to_vllm_mapper: ClassVar[WeightsMapper] = WeightsMapper(
hf_to_vllm_mapper = WeightsMapper(
orig_to_new_stacked={
# weight_name: (param_name, shard_id)
".q_proj": (".qkv_proj", "q"),
+1 -5
View File
@@ -206,9 +206,5 @@ class EagleLlama4ForCausalLM(Llama4ForCausalLM, SupportsMultiModalEmbeddings):
process_eagle_weight(self, name)
return name, weight
loader = AutoWeightsLoader(
self,
# lm_head is tied with target model (Llama4ForCausalLM)
skip_prefixes=([]),
)
loader = AutoWeightsLoader(self)
loader.load_weights(map(transform, weights))
+1 -4
View File
@@ -175,8 +175,5 @@ class EagleLlamaForCausalLM(LlamaForCausalLM):
process_eagle_weight(self, name)
return name, loaded_weight
loader = AutoWeightsLoader(
self,
skip_prefixes=None,
)
loader = AutoWeightsLoader(self)
loader.load_weights(map(transform, weights))
+9 -10
View File
@@ -415,18 +415,17 @@ class Eagle3LlamaForCausalLM(LlamaForCausalLM):
"Please provide mask_hidden in the weights."
)
skip_substrs = ["mask_hidden"]
orig_to_new_substr = {"mask_hidden": None}
if not includes_draft_id_mapping:
skip_substrs.append("draft_id_to_target_id")
orig_to_new_substr["draft_id_to_target_id"] = None
if not includes_embed_tokens:
skip_substrs.append("embed_tokens")
orig_to_new_substr["embed_tokens"] = None
if not self.model.use_aux_hidden_state:
skip_substrs.append("fc.")
orig_to_new_substr["fc."] = None
if not self.model.norm_before_fc:
skip_substrs.append("input_norm.")
loader = AutoWeightsLoader(
self,
skip_prefixes=None,
skip_substrs=skip_substrs,
orig_to_new_substr["input_norm."] = None
loader = AutoWeightsLoader(self)
loader.load_weights(
model_weights.items(),
mapper=WeightsMapper(orig_to_new_substr=orig_to_new_substr),
)
loader.load_weights(model_weights.items())
@@ -31,7 +31,7 @@ from vllm.v1.worker.gpu.states import RequestState
from .interfaces import SupportsLoRA, SupportsPP
from .longcat_flash import FlashConfig, FlashModel
from .utils import AutoWeightsLoader, PPMissingLayer, maybe_prefix
from .utils import AutoWeightsLoader, PPMissingLayer, WeightsMapper, maybe_prefix
def uses_ngram_embedding(config: FlashConfig) -> bool:
@@ -204,6 +204,9 @@ class FlashNgramModel(FlashModel):
class LongcatFlashNgramForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
"""LongCat-Flash-Lite for causal LM (MRV2-only, n-gram embedding)."""
# MTP weights are not part of this model.
hf_to_vllm_mapper = WeightsMapper(orig_to_new_prefix={"model.mtp.": None})
packed_modules_mapping = {
"qkv_proj": ["q_proj", "k_proj", "v_proj"],
"gate_up_proj": ["gate_proj", "up_proj"],
@@ -266,10 +269,9 @@ class LongcatFlashNgramForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
# AutoWeightsLoader routes ``model.*`` to FlashNgramModel.load_weights
# (which handles the ngram split) and ``lm_head.*`` to the head. MTP
# weights are not part of this model.
loader = AutoWeightsLoader(self, skip_prefixes=["model.mtp."])
return loader.load_weights(weights)
# (which handles the ngram split) and ``lm_head.*`` to the head.
loader = AutoWeightsLoader(self)
return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
class LongcatNgramModelState(DefaultModelState):
+6 -4
View File
@@ -41,7 +41,7 @@ from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
from vllm.model_executor.models.qwen2 import Qwen2ForCausalLM, Qwen2Model
from vllm.sequence import IntermediateTensors
from .utils import AutoWeightsLoader, PPMissingLayer, maybe_prefix
from .utils import AutoWeightsLoader, PPMissingLayer, WeightsMapper, maybe_prefix
logger = init_logger(__name__)
@@ -87,6 +87,9 @@ class MiMoModel(Qwen2Model):
class MiMoForCausalLM(Qwen2ForCausalLM, nn.Module):
# MTP layers are loaded by the draft model, not the main model.
hf_to_vllm_mapper = WeightsMapper(orig_to_new_prefix={"model.mtp_layers.": None})
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
nn.Module.__init__(self)
config = vllm_config.model_config.hf_config
@@ -119,9 +122,8 @@ class MiMoForCausalLM(Qwen2ForCausalLM, nn.Module):
)
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
# MTP layers are loaded by the draft model, not the main model.
loader = AutoWeightsLoader(self, skip_prefixes=["model.mtp_layers."])
return loader.load_weights(weights)
loader = AutoWeightsLoader(self)
return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
def compute_logits(
self,
+2 -1
View File
@@ -1189,6 +1189,7 @@ class MiMoV2OmniForCausalLM(nn.Module, SupportsMultiModal, SupportsPP, SupportsQ
# mapping for original checkpoint
"lm_head.": "language_model.lm_head.",
"model.": "language_model.model.",
"audio_tokenizer.": None,
}
)
@@ -1490,6 +1491,6 @@ class MiMoV2OmniForCausalLM(nn.Module, SupportsMultiModal, SupportsPP, SupportsQ
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
audio_loaded: set[str] = set()
loader = AutoWeightsLoader(self, skip_prefixes=["audio_tokenizer."])
loader = AutoWeightsLoader(self)
auto_loaded = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
return audio_loaded | auto_loaded
+11 -3
View File
@@ -71,7 +71,12 @@ from .minicpmv import (
MiniCPMVProcessingInfo,
_minicpmv_field_config,
)
from .utils import AutoWeightsLoader, cast_overflow_tensors, maybe_prefix
from .utils import (
AutoWeightsLoader,
WeightsMapper,
cast_overflow_tensors,
maybe_prefix,
)
CPU_DEVICE = torch.device("cpu")
@@ -672,6 +677,9 @@ class MiniCPMWhisperEncoder(WhisperEncoder):
class MiniCPMOBaseModel:
"""Base mixin class for MiniCPM-O models with audio support."""
# Unlike the vision-only MiniCPM-V models, audio weights are loaded here.
hf_to_vllm_mapper = WeightsMapper(orig_to_new_prefix={"tts": None})
packed_modules_mapping = {
"qkv_proj": [
"q_proj",
@@ -718,8 +726,8 @@ class MiniCPMOBaseModel:
return model
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
loader = AutoWeightsLoader(self, skip_prefixes=["tts"])
loaded = loader.load_weights(weights)
loader = AutoWeightsLoader(self)
loaded = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
self._ensure_resampler_device()
return loaded
+13 -7
View File
@@ -109,7 +109,7 @@ from .interfaces import (
SupportsMultiModal,
SupportsPP,
)
from .utils import AutoWeightsLoader, flatten_bn, maybe_prefix
from .utils import AutoWeightsLoader, WeightsMapper, flatten_bn, maybe_prefix
# For profile run
_MAX_FRAMES_PER_VIDEO = 16
@@ -1567,6 +1567,10 @@ class MiniCPMV2_5(MiniCPMVBaseModel, SupportsLoRA):
class MiniCPMV2_6(MiniCPMVBaseModel, SupportsLoRA):
hf_to_vllm_mapper = WeightsMapper(
# The vision-only models have no audio tower or TTS head to load weights into.
orig_to_new_prefix={"apm.": None, "audio": None, "tts": None}
)
packed_modules_mapping = {
"qkv_proj": [
"q_proj",
@@ -1660,13 +1664,14 @@ class MiniCPMV2_6(MiniCPMVBaseModel, SupportsLoRA):
return self.resampler(vision_embedding, tgt_sizes)
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
loader = AutoWeightsLoader(self, skip_prefixes=["apm.", "audio", "tts"])
loaded = loader.load_weights(weights)
loader = AutoWeightsLoader(self)
loaded = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
self._ensure_resampler_device()
return loaded
class MiniCPMV4_0(MiniCPMVBaseModel, SupportsLoRA):
hf_to_vllm_mapper = MiniCPMV2_6.hf_to_vllm_mapper
packed_modules_mapping = {
"qkv_proj": [
"q_proj",
@@ -1758,13 +1763,14 @@ class MiniCPMV4_0(MiniCPMVBaseModel, SupportsLoRA):
return self.resampler(vision_embedding, tgt_sizes)
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
loader = AutoWeightsLoader(self, skip_prefixes=["apm.", "audio", "tts"])
loaded = loader.load_weights(weights)
loader = AutoWeightsLoader(self)
loaded = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
self._ensure_resampler_device()
return loaded
class MiniCPMV4_5(MiniCPMVBaseModel, SupportsLoRA):
hf_to_vllm_mapper = MiniCPMV2_6.hf_to_vllm_mapper
packed_modules_mapping = {
"qkv_proj": [
"q_proj",
@@ -1861,8 +1867,8 @@ class MiniCPMV4_5(MiniCPMVBaseModel, SupportsLoRA):
return self.resampler(vision_embedding, tgt_sizes, all_temporal_ids)
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
loader = AutoWeightsLoader(self, skip_prefixes=["apm.", "audio", "tts"])
loaded = loader.load_weights(weights)
loader = AutoWeightsLoader(self)
loaded = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
self._ensure_resampler_device()
return loaded
+2 -1
View File
@@ -958,6 +958,7 @@ class MiniCPMV4_6ForConditionalGeneration(
"model.merger.": "merger.",
"model.language_model.": "language_model.model.",
"lm_head.": "language_model.lm_head.",
"mtp.": None,
}
)
@@ -1281,7 +1282,7 @@ class MiniCPMV4_6ForConditionalGeneration(
self,
weights: Iterable[tuple[str, torch.Tensor]],
) -> set[str]:
loader = AutoWeightsLoader(self, skip_prefixes=["mtp."])
loader = AutoWeightsLoader(self)
return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
def get_mm_mapping(self) -> MultiModelKeys:
+9 -11
View File
@@ -371,6 +371,14 @@ class MiniMaxM2Model(nn.Module, EagleModelMixin):
["hidden_states", "residual"], config.hidden_size
)
# Drop spec-decode (MTP) layers; they are appended after the main
# decoder layers and have no destination in the main model.
if num_mtp := getattr(config, "num_mtp_modules", 0):
base = config.num_hidden_layers
self.hf_to_vllm_mapper = self.hf_to_vllm_mapper | WeightsMapper(
orig_to_new_prefix={f"layers.{base + i}.": None for i in range(num_mtp)}
)
def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
return self.embed_tokens(input_ids)
@@ -413,17 +421,7 @@ class MiniMaxM2Model(nn.Module, EagleModelMixin):
return hidden_states
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
# Skip spec-decode (MTP) layers; they are appended after the main
# decoder layers and have no destination in the main model.
skip_prefixes = None
num_mtp = getattr(self.config, "num_mtp_modules", 0)
if num_mtp:
base = self.config.num_hidden_layers
skip_prefixes = [f"layers.{base + i}." for i in range(num_mtp)]
loader = AutoWeightsLoader(
self,
skip_prefixes=skip_prefixes,
)
loader = AutoWeightsLoader(self)
return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
+4 -2
View File
@@ -434,6 +434,8 @@ class ModernBertPredictionHead(nn.Module):
class ModernBertForTokenClassification(nn.Module):
is_pooling_model = True
hf_to_vllm_mapper = WeightsMapper(orig_to_new_prefix={"drop": None})
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
super().__init__()
config = vllm_config.model_config.hf_config
@@ -456,8 +458,8 @@ class ModernBertForTokenClassification(nn.Module):
return self.model.embed_input_ids(input_ids)
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
loader = AutoWeightsLoader(self, skip_prefixes=["drop"])
loaded_params = loader.load_weights(weights)
loader = AutoWeightsLoader(self)
loaded_params = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
return loaded_params
def forward(
+2 -4
View File
@@ -1450,6 +1450,7 @@ class MossAudioModel(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA):
"language_model.embed_tokens.": "language_model.model.embed_tokens.",
"language_model.layers.": "language_model.model.layers.",
"language_model.norm.": "language_model.model.norm.",
"audio_encoder.embed_positions": None,
},
orig_to_new_stacked={
".gate_proj": (".gate_up_proj", 0),
@@ -1865,8 +1866,5 @@ class MossAudioModel(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA):
return self.language_model.compute_logits(hidden_states)
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
loader = AutoWeightsLoader(
self,
skip_prefixes=["audio_encoder.embed_positions"],
)
loader = AutoWeightsLoader(self)
return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
+2 -2
View File
@@ -708,7 +708,7 @@ class NemotronHForCausalLM(
is_non_gated_moe: bool = True
hf_to_vllm_mapper = WeightsMapper(
orig_to_new_prefix={"backbone": "model"},
orig_to_new_prefix={"backbone": "model", "mtp": None},
orig_to_new_substr={"A_log": "A", "embeddings": "embed_tokens"},
orig_to_new_stacked={
".q_proj": (".qkv_proj", "q"),
@@ -887,5 +887,5 @@ class NemotronHForCausalLM(
return logits
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
loader = AutoWeightsLoader(self, skip_prefixes=["mtp"])
loader = AutoWeightsLoader(self)
return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
+8 -5
View File
@@ -89,6 +89,12 @@ class NemotronVLProcessingInfo(BaseInternVLProcessingInfo):
dummy_inputs=BaseInternVLDummyInputsBuilder[NemotronVLProcessingInfo],
)
class LlamaNemotronVLChatModel(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA):
# Ignore registered buffers, see
# https://huggingface.co/nvidia/C-RADIOv2-H/blob/main/input_conditioner.py#L28
hf_to_vllm_mapper = WeightsMapper(
orig_to_new_substr={"norm_mean": None, "norm_std": None}
)
@classmethod
def get_placeholder_str(cls, modality: str, i: int) -> str | None:
if modality.startswith("image"):
@@ -377,11 +383,8 @@ class LlamaNemotronVLChatModel(nn.Module, SupportsMultiModal, SupportsPP, Suppor
return self.language_model.compute_logits(hidden_states)
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
## Ignore registered_buffers
## see https://huggingface.co/nvidia/C-RADIOv2-H/blob/main/input_conditioner.py#L28 # noqa: E501
skip_substrs = ["norm_mean", "norm_std"]
loader = AutoWeightsLoader(self, skip_substrs=skip_substrs)
return loader.load_weights(weights)
loader = AutoWeightsLoader(self)
return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
def get_mm_mapping(self) -> MultiModelKeys:
"""
+11 -13
View File
@@ -887,7 +887,16 @@ class SiglipVisionModel(nn.Module):
".q_proj": (".qkv_proj", "q"),
".k_proj": (".qkv_proj", "k"),
".v_proj": (".qkv_proj", "v"),
}
},
# The SigLIP attention pooling head and packing pos embedding are
# present in the checkpoint but absent from this vision tower.
orig_to_new_substr={
"head.attention": None,
"head.layernorm": None,
"head.mlp": None,
"head.probe": None,
"packing_position_embedding": None,
},
)
def __init__(
@@ -934,18 +943,7 @@ class SiglipVisionModel(nn.Module):
)
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
# Skip the SigLIP attention pooling head and packing pos embedding
# present in the checkpoint but absent from this vision tower.
loader = AutoWeightsLoader(
self,
skip_substrs=[
"head.attention",
"head.layernorm",
"head.mlp",
"head.probe",
"packing_position_embedding",
],
)
loader = AutoWeightsLoader(self)
return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
+2 -1
View File
@@ -1016,6 +1016,7 @@ class Phi4MMForCausalLM(nn.Module, SupportsLoRA, SupportsMultiModal):
hf_to_vllm_mapper = WeightsMapper(
orig_to_new_substr={
"base_layer.": "",
"lora": None,
},
orig_to_new_prefix={
"model.embed_tokens_extend.audio_embed.audio_projection.vision.": "embed_tokens_extend.audio_projection_for_vision.", # noqa: E501
@@ -1258,7 +1259,7 @@ class Phi4MMForCausalLM(nn.Module, SupportsLoRA, SupportsMultiModal):
return logits
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> None:
loader = AutoWeightsLoader(self, skip_substrs=["lora"])
loader = AutoWeightsLoader(self)
return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
def get_mm_mapping(self) -> MultiModelKeys:
@@ -1112,6 +1112,8 @@ class Qwen2_5OmniThinkerForConditionalGeneration(
"thinker.lm_head.": "language_model.lm_head.",
"thinker.model.": "language_model.model.",
"thinker.": "",
"talker.": None,
"token2wav.": None,
}
)
packed_modules_mapping = {
@@ -1591,7 +1593,7 @@ class Qwen2_5OmniThinkerForConditionalGeneration(
return self.language_model.compute_logits(hidden_states)
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
loader = AutoWeightsLoader(self, skip_prefixes=["talker.", "token2wav."])
loader = AutoWeightsLoader(self)
return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
def get_mm_mapping(self) -> MultiModelKeys:
+8 -9
View File
@@ -319,7 +319,7 @@ class Qwen3_5ForCausalLMBase(
# `model.language_model.` prefix inherited from the VL training stack.
# Strip it so both prefixed and clean checkpoints load correctly.
hf_to_vllm_mapper = WeightsMapper(
orig_to_new_prefix={"model.language_model.": "model."},
orig_to_new_prefix={"model.language_model.": "model.", "mtp.": None},
)
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
@@ -430,10 +430,7 @@ class Qwen3_5ForCausalLMBase(
return self.logits_processor(self.lm_head, hidden_states)
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
loader = AutoWeightsLoader(
self,
skip_prefixes=["mtp."],
)
loader = AutoWeightsLoader(self)
return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
def get_mrope_input_positions(
@@ -470,6 +467,11 @@ class Qwen3_5MoeForCausalLM(Qwen3_5ForCausalLMBase, QwenNextMixtureOfExperts):
class Qwen3_5ForConditionalGeneration(Qwen3VLForConditionalGeneration, IsHybrid):
supports_multimodal_pruning = True
hf_to_vllm_mapper = (
Qwen3VLForConditionalGeneration.hf_to_vllm_mapper
| WeightsMapper(orig_to_new_prefix={"mtp.": None})
)
packed_modules_mapping = Qwen3VLForConditionalGeneration.packed_modules_mapping | {
"in_proj_qkvz": ["in_proj_qkv", "in_proj_z"],
"in_proj_ba": ["in_proj_b", "in_proj_a"],
@@ -590,10 +592,7 @@ class Qwen3_5ForConditionalGeneration(Qwen3VLForConditionalGeneration, IsHybrid)
return hidden_states
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
loader = AutoWeightsLoader(
self,
skip_prefixes=["mtp."],
)
loader = AutoWeightsLoader(self)
return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
@classmethod
+3 -4
View File
@@ -372,6 +372,8 @@ class Qwen3ASRForConditionalGeneration(
"model.language_model.": "language_model.model.",
"model.multi_modal_projector.linear_1.": "audio_tower.proj1.",
"model.multi_modal_projector.linear_2.": "audio_tower.proj2.",
"talker.": None,
"code2wav.": None,
}
)
@@ -545,10 +547,7 @@ class Qwen3ASRForConditionalGeneration(
return self.language_model.compute_logits(hidden_states)
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
loader = AutoWeightsLoader(
self,
skip_prefixes=["talker.", "code2wav."],
)
loader = AutoWeightsLoader(self)
loaded_weights = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
return loaded_weights
@@ -59,6 +59,8 @@ class Qwen3ASRForcedAlignerForTokenClassification(
"thinker.lm_head.": "classifier.",
"thinker.model.": "language_model.model.",
"thinker.": "",
"talker.": None,
"code2wav.": None,
}
)
@@ -113,8 +115,5 @@ class Qwen3ASRForcedAlignerForTokenClassification(
return self.classifier(hidden_states)
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
loader = AutoWeightsLoader(
self,
skip_prefixes=["talker.", "code2wav."],
)
loader = AutoWeightsLoader(self)
return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
+8 -11
View File
@@ -835,21 +835,18 @@ class DFlashQwen3ForCausalLM(Qwen3ForCausalLM):
model_weights["model.mask_embedding"] = mask_embedding
self.model.has_separate_mask_embedding = True
skip_substrs = []
orig_to_new_substr = {}
if not includes_draft_id_mapping:
skip_substrs.append("draft_id_to_target_id")
orig_to_new_substr["draft_id_to_target_id"] = None
if not includes_embed_tokens:
skip_substrs.append("embed_tokens")
orig_to_new_substr["embed_tokens"] = None
if not self.model.use_aux_hidden_state:
skip_substrs.append("fc.")
orig_to_new_substr["fc."] = None
if not self.model.has_separate_mask_embedding:
skip_substrs.append("mask_embedding")
loader = AutoWeightsLoader(
self,
skip_prefixes=None,
skip_substrs=skip_substrs,
)
loader.load_weights(model_weights.items())
orig_to_new_substr["mask_embedding"] = None
mapper = WeightsMapper(orig_to_new_substr=orig_to_new_substr)
loader = AutoWeightsLoader(self)
loader.load_weights(model_weights.items(), mapper=mapper)
self.model._build_fused_kv_buffers()
def _read_mask_embedding(self) -> torch.Tensor | None:
+14 -8
View File
@@ -30,7 +30,12 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
)
from .qwen3_dflash import DFlashQwen3ForCausalLM, DFlashQwen3Model
from .utils import AutoWeightsLoader, maybe_prefix, process_eagle_weight
from .utils import (
AutoWeightsLoader,
WeightsMapper,
maybe_prefix,
process_eagle_weight,
)
logger = init_logger(__name__)
@@ -275,16 +280,17 @@ class Qwen3DSparkForCausalLM(DFlashQwen3ForCausalLM):
# mask_embedding is an unused placeholder param; DSpark masks via the vocab row.
# embed_tokens / lm_head are optional; when omitted they are shared from
# the target by load_dspark_model, so skip the unloaded params here.
skip_substrs = ["mask_embedding"]
orig_to_new_substr = {"mask_embedding": None}
if not includes_embed_tokens:
skip_substrs.append("embed_tokens")
orig_to_new_substr["embed_tokens"] = None
if not includes_lm_head:
skip_substrs.append("lm_head")
orig_to_new_substr["lm_head"] = None
if not includes_draft_id_mapping:
skip_substrs.append("draft_id_to_target_id")
orig_to_new_substr["draft_id_to_target_id"] = None
if self.model.confidence_head is None or not includes_confidence_head:
self.model.confidence_head = None
skip_substrs.append("confidence_head")
loader = AutoWeightsLoader(self, skip_substrs=skip_substrs)
loader.load_weights(model_weights.items())
orig_to_new_substr["confidence_head"] = None
mapper = WeightsMapper(orig_to_new_substr=orig_to_new_substr)
loader = AutoWeightsLoader(self)
loader.load_weights(model_weights.items(), mapper=mapper)
self.model._build_fused_kv_buffers()
+8 -11
View File
@@ -425,18 +425,15 @@ class Eagle3Qwen3ForCausalLM(Qwen3ForCausalLM):
"Please provide mask_hidden in the weights."
)
skip_substrs = ["mask_hidden"]
orig_to_new_substr = {"mask_hidden": None}
if not includes_draft_id_mapping:
skip_substrs.append("draft_id_to_target_id")
orig_to_new_substr["draft_id_to_target_id"] = None
if not includes_embed_tokens:
skip_substrs.append("embed_tokens")
orig_to_new_substr["embed_tokens"] = None
if not self.model.use_aux_hidden_state:
skip_substrs.append("fc.")
orig_to_new_substr["fc."] = None
if not self.model.norm_before_fc:
skip_substrs.append("input_norm.")
loader = AutoWeightsLoader(
self,
skip_prefixes=None,
skip_substrs=skip_substrs,
)
loader.load_weights(model_weights.items())
orig_to_new_substr["input_norm."] = None
mapper = WeightsMapper(orig_to_new_substr=orig_to_new_substr)
loader = AutoWeightsLoader(self)
loader.load_weights(model_weights.items(), mapper=mapper)
+5 -2
View File
@@ -727,6 +727,9 @@ class Qwen3NextForCausalLM(
IsHybrid,
SupportsEagle3,
):
# MTP weights are loaded by the draft model, not this one.
hf_to_vllm_mapper = WeightsMapper(orig_to_new_prefix={"mtp.": None})
packed_modules_mapping = {
"qkv_proj": [
"q_proj",
@@ -833,5 +836,5 @@ class Qwen3NextForCausalLM(
return self.logits_processor(self.lm_head, hidden_states)
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
loader = AutoWeightsLoader(self, skip_prefixes=["mtp."])
return loader.load_weights(weights)
loader = AutoWeightsLoader(self)
return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
@@ -1587,6 +1587,8 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
"thinker.lm_head.": "language_model.lm_head.",
"thinker.model.": "language_model.model.",
"thinker.": "",
"talker.": None,
"code2wav.": None,
}
)
@@ -1943,10 +1945,7 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
return self.language_model.compute_logits(hidden_states)
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
loader = AutoWeightsLoader(
self,
skip_prefixes=["talker.", "code2wav."],
)
loader = AutoWeightsLoader(self)
loaded_weights = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
return loaded_weights
+5 -3
View File
@@ -147,19 +147,21 @@ class RobertaEmbeddingModel(BertEmbeddingModel):
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
weights_list = list(weights)
orig_to_new_prefix: dict[str, str | None] = {"lm_head.": None}
has_roberta_prefix = any(
name.startswith("roberta.") for name, _ in weights_list
)
if has_roberta_prefix:
# For models with the `roberta.` prefix e.g.
# `FacebookAI/roberta-base`
mapper = WeightsMapper(orig_to_new_prefix={"roberta.": "model."})
orig_to_new_prefix["roberta."] = "model."
else:
# For models without the `roberta.` prefix e.g.
# `sentence-transformers/stsb-roberta-base-v2`
mapper = WeightsMapper(orig_to_new_prefix={"": "model."})
orig_to_new_prefix[""] = "model."
mapper = WeightsMapper(orig_to_new_prefix=orig_to_new_prefix)
loader = AutoWeightsLoader(self, skip_prefixes=["lm_head."])
loader = AutoWeightsLoader(self)
return loader.load_weights(weights_list, mapper=mapper)
+14 -8
View File
@@ -738,6 +738,16 @@ class SiglipVisionTransformer(nn.Module):
)
self.last_hs_proc = partial(self.maybe_layer_norm_and_apply_head)
drops = dict[str, None]()
if self.post_layernorm is None:
drops["post_layernorm."] = None
if self.head is None:
drops["head."] = None
if drops:
self.hf_to_vllm_mapper = self.hf_to_vllm_mapper | WeightsMapper(
orig_to_new_prefix=drops
)
@property
def dtype(self):
return next(self.parameters()).dtype
@@ -797,12 +807,7 @@ class SiglipVisionTransformer(nn.Module):
return encoder_outputs
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
skip_prefixes = []
if self.post_layernorm is None:
skip_prefixes.append("post_layernorm.")
if self.head is None:
skip_prefixes.append("head.")
loader = AutoWeightsLoader(self, skip_prefixes=skip_prefixes)
loader = AutoWeightsLoader(self)
layer_count = len(self.encoder.layers)
@@ -912,6 +917,8 @@ class SiglipTextEmbeddings(nn.Module):
class SiglipEmbeddingModel(nn.Module, SupportsMultiModal, SupportsQuant):
is_pooling_model = True
hf_to_vllm_mapper = WeightsMapper(orig_to_new_substr={".position_ids": None})
packed_modules_mapping = {"qkv_proj": ["q_proj", "k_proj", "v_proj"]}
@classmethod
@@ -1154,8 +1161,7 @@ class SiglipEmbeddingModel(nn.Module, SupportsMultiModal, SupportsQuant):
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
loader = AutoWeightsLoader(
self,
skip_substrs=[".position_ids"],
ignore_unexpected_prefixes=["logit_scale.", "logit_bias."],
)
return loader.load_weights(weights)
return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
+27 -17
View File
@@ -35,7 +35,12 @@ from .internvl import (
BaseInternVLMultiModalProcessor,
BaseInternVLProcessingInfo,
)
from .utils import AutoWeightsLoader, init_vllm_registered_model, maybe_prefix
from .utils import (
AutoWeightsLoader,
WeightsMapper,
init_vllm_registered_model,
maybe_prefix,
)
class SkyworkR1VImagePixelInputs(TensorSchema):
@@ -120,6 +125,25 @@ class SkyworkR1VProcessingInfo(BaseInternVLProcessingInfo):
dummy_inputs=BaseInternVLDummyInputsBuilder,
)
class SkyworkR1VChatModel(nn.Module, SupportsMultiModal, SupportsPP):
hf_to_vllm_mapper = WeightsMapper(
orig_to_new_prefix=dict.fromkeys(
[
"action_embed",
"temporal_embed",
"track_embed",
"track_embed_decoder",
"box_token",
"cg_criterion",
"cg_model",
"loc_encoder",
"loc_decoder",
"sam",
"temporal_token",
"track_token",
]
)
)
@classmethod
def get_placeholder_str(cls, modality: str, i: int) -> str | None:
if modality.startswith("image"):
@@ -381,19 +405,5 @@ class SkyworkR1VChatModel(nn.Module, SupportsMultiModal, SupportsPP):
return self.language_model.compute_logits(hidden_states)
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
skip_prefixes = [
"action_embed",
"temporal_embed",
"track_embed",
"track_embed_decoder",
"box_token",
"cg_criterion",
"cg_model",
"loc_encoder",
"loc_decoder",
"sam",
"temporal_token",
"track_token",
]
loader = AutoWeightsLoader(self, skip_prefixes=skip_prefixes)
return loader.load_weights(weights)
loader = AutoWeightsLoader(self)
return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
@@ -122,10 +122,6 @@ class Base(
self.tp_group = get_tp_group()
# Attrs for weight loading (see self.load_weights)
self.skip_prefixes: list[str] = []
"""Skip loading weights whose qualname starts with these prefixes."""
self.skip_substrs: list[str] = []
"""Skip loading weights whose qualname contains these substrings."""
self.ignore_unexpected_prefixes: list[str] = []
"""Ignore unexpected weights whose qualname starts with these prefixes."""
self.ignore_unexpected_suffixes: list[str] = []
@@ -731,8 +727,6 @@ class Base(
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
loader = AutoWeightsLoader(
self,
skip_prefixes=self.skip_prefixes,
skip_substrs=self.skip_substrs,
ignore_unexpected_prefixes=self.ignore_unexpected_prefixes,
ignore_unexpected_suffixes=self.ignore_unexpected_suffixes,
)
@@ -30,26 +30,24 @@ class LegacyMixin:
def __init__(self, *, vllm_config: "VllmConfig", prefix: str = ""):
super().__init__(vllm_config=vllm_config, prefix=prefix)
# Skip unsupported/unwanted output embeddings layers
self.skip_prefixes.extend(
[
"model.lm_head.",
"model.predictions.",
"model.qa_outputs.",
"model.embeddings_project.",
"model.discriminator_predictions.",
]
# Drop unsupported/unwanted output embeddings layers.
self.hf_to_vllm_mapper.orig_to_new_prefix.update(
{
"model.lm_head.": None,
"model.predictions.": None,
"model.qa_outputs.": None,
"model.embeddings_project.": None,
"model.discriminator_predictions.": None,
}
)
# Some encoder models have the position_ids buffer in the checkpoint.
# vLLM will always pass position_ids as an argument, so we skip loading
# the buffer if it exists
self.skip_substrs.append("position_ids")
# Some encoder models have the position_ids buffer in the checkpoint. vLLM will
# always pass position_ids as an argument, so we drop the buffer if it exists.
self.hf_to_vllm_mapper.orig_to_new_substr["position_ids"] = None
# Some encoder models have the bias of the final classifier layer
# in the checkpoint. vLLM does not use this bias, so we skip loading
# it if it exists
self.skip_substrs.append("score.bias")
# Some encoder models have the bias of the final classifier layer in the
# checkpoint. vLLM does not use this bias, so we drop it if it exists.
self.hf_to_vllm_mapper.orig_to_new_substr["score.bias"] = None
# roberta-like models an extra padding in positions.
# FIXME(Isotr0py): This is quite hacky for roberta edge case,
+34 -36
View File
@@ -160,14 +160,26 @@ class WeightsMapper:
if (out_name := self._map_name(name)) is not None
}
def get_unstacked_mapper(self) -> "WeightsMapper":
"""Mapper variant that drops stacked maps, keeping all genuine renames/prefixes.
def get_rename_mapper(self) -> "WeightsMapper":
"""Mapper variant keeping only the renames.
Consumers that reference the checkpoint's *unstacked* module names (LoRA name
parsing and the quantization config's layer lists) need the constituent names
(e.g. `q_proj`) to survive rather than being rewritten to the stacked vLLM name
(`qkv_proj`)."""
return replace(self, orig_to_new_stacked={})
This is what consumers that *name* modules rather than load them need:
LoRA name parsing and the quantization config's layer lists.
Stacked maps are dropped so that constituent names (e.g. `q_proj`) survive
rather than being rewritten to the stacked vLLM name (`qkv_proj`). Mappings to
`None` are dropped because "do not load this weight" is meaningless to such a
consumer, and applying it would silently shrink a quantization config's ignore
list or make LoRA name parsing fail."""
remove_none = lambda d: {k: v for k, v in d.items() if v is not None}
return replace(
self,
orig_to_new_regex=remove_none(self.orig_to_new_regex),
orig_to_new_substr=remove_none(self.orig_to_new_substr),
orig_to_new_stacked={},
orig_to_new_prefix=remove_none(self.orig_to_new_prefix),
orig_to_new_suffix=remove_none(self.orig_to_new_suffix),
)
def _get_tied_embedding_params(module: nn.Module) -> dict[str, str]:
@@ -201,32 +213,28 @@ class AutoWeightsLoader:
"""
# Models trained using early version ColossalAI or quantized by
# GPTQModel may include these tensors in checkpoint. Skip them.
ROTARY_EMBEDS_UNUSED_WEIGHTS = [
"rotary_pos_emb.inv_freq",
"rotary_emb.inv_freq",
"rotary_emb.cos_cached",
"rotary_emb.sin_cached",
]
# GPTQModel may include these tensors in checkpoint. Drop them.
REMOVE_UNUSED_ROTARY_EMBEDS_MAPPER = WeightsMapper(
orig_to_new_substr={
"rotary_pos_emb.inv_freq": None,
"rotary_emb.inv_freq": None,
"rotary_emb.cos_cached": None,
"rotary_emb.sin_cached": None,
}
)
def __init__(
self,
module: nn.Module,
*,
skip_prefixes: list[str] | None = None,
skip_substrs: list[str] | None = None,
ignore_unexpected_prefixes: list[str] | None = None,
ignore_unexpected_suffixes: list[str] | None = None,
) -> None:
super().__init__()
self.module = module
self.skip_prefixes = skip_prefixes or []
self.skip_substrs = skip_substrs or []
self.ignore_unexpected_prefixes = ignore_unexpected_prefixes or []
self.ignore_unexpected_suffixes = ignore_unexpected_suffixes or []
# update default skip_substrs
self.skip_substrs += self.ROTARY_EMBEDS_UNUSED_WEIGHTS
# Weight tying makes two qualnames point at the same `nn.Parameter`
# (e.g. `lm_head.weight` and `model.embed_tokens.weight`). Loading both
@@ -265,11 +273,7 @@ class AutoWeightsLoader:
return ".".join((prefix, rest))
def _can_skip(self, qualname: str) -> bool:
return (
qualname in self.aliased_params
or any(qualname.startswith(p) for p in self.skip_prefixes)
or any(substr in qualname for substr in self.skip_substrs)
)
return qualname in self.aliased_params
def _can_ignore_unexpected(self, qualname: str) -> bool:
iup = (qualname.startswith(p) for p in self.ignore_unexpected_prefixes)
@@ -375,11 +379,6 @@ class AutoWeightsLoader:
prefix = self._get_qualname(base_prefix, child_prefix)
if child_prefix in child_modules:
if self._can_skip(prefix + "."):
logger.debug("Skipping module %s", prefix)
continue
yield from self._load_module(
prefix, child_modules[child_prefix], child_weights
)
@@ -393,9 +392,7 @@ class AutoWeightsLoader:
prefix, child_params[child_prefix], child_weights
)
else:
can_skip_module = self._can_skip(prefix + ".")
can_skip_param = self._can_skip(prefix)
if can_skip_module or can_skip_param:
if self._can_skip(prefix):
logger.debug("Skipping missing %s", prefix)
continue
@@ -429,18 +426,19 @@ class AutoWeightsLoader:
# Ignore unexpected biases (typically from GPTQ models)
self.ignore_unexpected_suffixes.append(".bias")
mapper = mapper or WeightsMapper()
# Many models store quant_config in the base model instead of the causal model.
# We look at the causal model's direct children for this reason.
modules = (self.module, *self.module.children())
iterator = (m.quant_config for m in modules if hasattr(m, "quant_config"))
if quant_config := next(iterator, None):
# Get mappings and ignore prefixes for KV cache quantization scales
mapper = mapper or WeightsMapper()
mapper |= quant_config.get_cache_scale_mapper()
ignore_unexpected_suffixes = quant_config._ignore_unexpected_suffixes
self.ignore_unexpected_suffixes.extend(ignore_unexpected_suffixes)
if mapper is not None:
weights = mapper.apply(weights)
mapper |= self.REMOVE_UNUSED_ROTARY_EMBEDS_MAPPER
weights = mapper.apply(weights)
weights = self._filter_skipped(weights)
autoloaded_weights = set(self._load_module("", self.module, weights))
+2 -1
View File
@@ -838,6 +838,7 @@ class WhisperForConditionalGeneration(
".encoder_attn.k_proj": (".encoder_attn.kv_proj", 0),
".encoder_attn.v_proj": (".encoder_attn.kv_proj", 1),
},
orig_to_new_prefix={"proj_out.": None},
)
# Whisper only supports audio-conditioned generation.
@@ -1048,7 +1049,7 @@ class WhisperForConditionalGeneration(
return logits
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
loader = AutoWeightsLoader(self, skip_prefixes=["proj_out."])
loader = AutoWeightsLoader(self)
# add fake zeros bias for k_proj to state_dict
weights = _create_fake_bias_for_k_proj(weights, ".k_proj.weight")
+4 -3
View File
@@ -947,11 +947,12 @@ def _make_deepseek_v4_weights_mapper(
# When shared experts are fused into the routed MXFP4 grouped GEMM, the
# shared_experts tensors are redirected to routed expert slots ; leave
# their names untouched here.
substr_map = (
orig_to_new_substr: dict[str, str | None] = (
{}
if fuse_shared_experts
else {".shared_experts.w2": ".shared_experts.down_proj"}
)
orig_to_new_substr["mtp."] = None
return WeightsMapper(
orig_to_new_prefix={
"layers.": "model.layers.",
@@ -967,7 +968,7 @@ def _make_deepseek_v4_weights_mapper(
".ffn.gate.bias": ".ffn.gate.e_score_correction_bias",
".input_scale": ".input_scale_2",
},
orig_to_new_substr=substr_map,
orig_to_new_substr=orig_to_new_substr,
)
@@ -1038,7 +1039,7 @@ class DeepseekV4ForCausalLM(nn.Module, SupportsPP, SupportsEagle3):
return getattr(self.model, "_mtp_hidden_buffer", None)
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
loader = AutoWeightsLoader(self, skip_substrs=["mtp."])
loader = AutoWeightsLoader(self)
return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
def process_weights_after_loading(self) -> None:
+2 -1
View File
@@ -1392,6 +1392,7 @@ def _make_deepseek_v4_weights_mapper(expert_dtype: str) -> WeightsMapper:
},
orig_to_new_substr={
".shared_experts.w2": ".shared_experts.down_proj",
"mtp.": None,
},
)
@@ -1517,7 +1518,7 @@ class DeepseekV4ForCausalLM(
return getattr(self.model, "_mtp_hidden_buffer", None)
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
loader = AutoWeightsLoader(self, skip_substrs=["mtp."])
loader = AutoWeightsLoader(self)
loaded_params = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
self.process_weights_after_loading()
return loaded_params
+2 -1
View File
@@ -1299,6 +1299,7 @@ def _make_deepseek_v4_weights_mapper(expert_dtype: str) -> WeightsMapper:
},
orig_to_new_substr={
".shared_experts.w2": ".shared_experts.down_proj",
"mtp.": None,
},
)
@@ -1364,7 +1365,7 @@ class DeepseekV4ForCausalLM(nn.Module, SupportsPP, SupportsEagle3):
return getattr(self.model, "_mtp_hidden_buffer", None)
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
loader = AutoWeightsLoader(self, skip_substrs=["mtp."])
loader = AutoWeightsLoader(self)
loaded_params = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
self.model.finalize_mega_moe_weights()
return loaded_params
+11 -6
View File
@@ -141,6 +141,16 @@ class Dots3NoteForCausalLM(nn.Module, SupportsMultiModal, SupportsPP):
self.language_model.make_empty_intermediate_tensors
)
orig_to_new_prefix = dict[str, None]()
if self.visual is None:
orig_to_new_prefix["visual."] = None
if self.audio_tower is None:
orig_to_new_prefix["audio_tower."] = None
if orig_to_new_prefix:
self.hf_to_vllm_mapper = self.hf_to_vllm_mapper | WeightsMapper(
orig_to_new_prefix=orig_to_new_prefix
)
def _process_image_input(
self,
pixel_values: torch.Tensor,
@@ -283,12 +293,7 @@ class Dots3NoteForCausalLM(nn.Module, SupportsMultiModal, SupportsPP):
return self.language_model.compute_logits(hidden_states)
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
skip_prefixes = []
if self.visual is None:
skip_prefixes.append("visual.")
if self.audio_tower is None:
skip_prefixes.append("audio_tower.")
return AutoWeightsLoader(self, skip_prefixes=skip_prefixes).load_weights(
return AutoWeightsLoader(self).load_weights(
weights,
mapper=self.hf_to_vllm_mapper,
)
+3 -2
View File
@@ -656,8 +656,9 @@ def _load_inkling_weights(
yield name, weight
loader = AutoWeightsLoader(module, skip_prefixes=["model.mtp."])
loaded |= loader.load_weights(_iter_loadable_weights())
mapper = WeightsMapper(orig_to_new_prefix={"model.mtp.": None})
loader = AutoWeightsLoader(module)
loaded |= loader.load_weights(_iter_loadable_weights(), mapper=mapper)
# Post-load MoE fixups (default input scales, zeroed EP-padding experts).
for moe_name, moe in moe_modules.items():
+3 -2
View File
@@ -693,8 +693,9 @@ def _load_inkling_weights(
# The release checkpoint also carries auxiliary prediction-head weights;
# they are not part of the causal LM served by this implementation.
loader = AutoWeightsLoader(module, skip_prefixes=["model.mtp."])
loaded |= loader.load_weights(_iter_loadable_weights())
mapper = WeightsMapper(orig_to_new_prefix={"model.mtp.": None})
loader = AutoWeightsLoader(module)
loaded |= loader.load_weights(_iter_loadable_weights(), mapper=mapper)
# Post-load MoE fixups (default input scales, zeroed EP-padding experts).
for moe_name, moe in moe_modules.items():
+8 -8
View File
@@ -399,9 +399,14 @@ class K3DSparkForCausalLM(nn.Module):
has_own_embed_tokens = False
has_own_lm_head = False
draft_id_to_target_id = None
checkpoint_skip_substrs = ("confidence_head", "embed_tokens", "lm_head")
hf_to_vllm_mapper = WeightsMapper(
# confidence_head is training-only. The frozen target embedding and LM
# head are shared after this draft-specific checkpoint is loaded.
orig_to_new_substr={
"confidence_head": None,
"embed_tokens": None,
"lm_head": None,
},
orig_to_new_prefix={"": "model."},
orig_to_new_stacked={
".gate_proj": (".gate_up_proj", 0),
@@ -477,12 +482,7 @@ class K3DSparkForCausalLM(nn.Module):
return self.model.markov_head.bias(markov_embed, self.logits_processor)
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
# confidence_head is training-only. The frozen target embedding and LM
# head are shared after this draft-specific checkpoint is loaded.
loader = AutoWeightsLoader(
self,
skip_substrs=list(self.checkpoint_skip_substrs),
)
loader = AutoWeightsLoader(self)
# read: 1. all weights. 2. context kv weights
weights = _duplicate_context_kv_weights(weights, len(self.model.layers))
loaded_weights = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)