5051698e41 · tduncan/Nvidia-Optimization-Stack-VLLM

tests/model_executor/test_weight_utils.py

+121

@@ -160,5 +160,126 @@ class TestMaybeRemapKvScaleName:

160 assert result is None

161

162

163 +class TestKvCacheScaleMapper:

164 + """The `WeightsMapper` returned by `get_cache_scale_mapper` replaces the

165 + per-model `maybe_remap_kv_scale_name` calls. It must remap the same set of

166 + checkpoint formats (the non-`params_dict`-dependent ones) and be idempotent

167 + so it composes safely with a model's own qkv/gate_up `hf_to_vllm_mapper`."""

168 +

169 + def _mapper(self):

170 + # `get_cache_scale_mapper` does not use `self`; call it on the base

171 + # class to get the default (non-config-specific) mapper.

172 + from vllm.model_executor.layers.quantization.base_config import (

173 + QuantizationConfig,

174 + )

175 +

176 + return QuantizationConfig.get_cache_scale_mapper()

177 +

178 + def _map(self, name: str) -> str | None:

179 + return self._mapper()._map_name(name)

180 +

181 + @pytest.mark.parametrize(

182 + "name,expected",

183 + [

184 + # Qwen3-MoE / llm-compressor fused qkv_proj

185 + (

186 + "model.layers.0.self_attn.qkv_proj.k_scale",

187 + "model.layers.0.self_attn.attn.k_scale",

188 + ),

189 + (

190 + "model.layers.0.self_attn.qkv_proj.v_scale",

191 + "model.layers.0.self_attn.attn.v_scale",

192 + ),

193 + # ModelOpt / NVFP4 k_proj/v_proj

194 + (

195 + "model.layers.0.self_attn.k_proj.k_scale",

196 + "model.layers.0.self_attn.attn.k_scale",

197 + ),

198 + (

199 + "model.layers.0.self_attn.v_proj.v_scale",

200 + "model.layers.0.self_attn.attn.v_scale",

201 + ),

202 + # deprecated fused kv_scale and bare scales

203 + (

204 + "model.layers.0.self_attn.kv_scale",

205 + "model.layers.0.self_attn.attn.k_scale",

206 + ),

207 + (

208 + "model.layers.0.self_attn.k_scale",

209 + "model.layers.0.self_attn.attn.k_scale",

210 + ),

211 + # NemotronH mixer

212 + (

213 + "model.layers.0.mixer.k_proj.k_scale",

214 + "model.layers.0.mixer.attn.k_scale",

215 + ),

216 + # already in vLLM form -> unchanged (idempotent)

217 + (

218 + "model.layers.0.self_attn.attn.k_scale",

219 + "model.layers.0.self_attn.attn.k_scale",

220 + ),

221 + # non-kv scales must not be touched

222 + (

223 + "model.layers.0.self_attn.k_proj.weight_scale",

224 + "model.layers.0.self_attn.k_proj.weight_scale",

225 + ),

226 + (

227 + "model.layers.0.self_attn.k_proj.input_scale",

228 + "model.layers.0.self_attn.k_proj.input_scale",

229 + ),

230 + # regular weights untouched

231 + (

232 + "model.layers.0.self_attn.q_proj.weight",

233 + "model.layers.0.self_attn.q_proj.weight",

234 + ),

235 + ],

236 + )

237 + def test_remap(self, name, expected):

238 + assert self._map(name) == expected

239 +

240 + @pytest.mark.parametrize(

241 + "name",

242 + [

243 + "model.layers.0.self_attn.k_scale",

244 + "model.layers.0.self_attn.k_proj.k_scale",

245 + "model.layers.0.self_attn.qkv_proj.v_scale",

246 + "model.layers.0.mixer.k_proj.k_scale",

247 + ],

248 + )

249 + def test_idempotent(self, name):

250 + once = self._map(name)

251 + assert once is not None

252 + assert self._map(once) == once

253 +

254 + def test_composes_with_qkv_mapper(self):

255 + """Applied together with a model's qkv/gate_up mapper, the regex scale

256 + rules run before the substr rename, so scales are normalized to `.attn.`

257 + and regular projections are still fused correctly."""

258 + from vllm.model_executor.models.utils import WeightsMapper

259 +

260 + model_mapper = WeightsMapper(

261 + orig_to_new_substr={

262 + ".q_proj": ".qkv_proj.q",

263 + ".k_proj": ".qkv_proj.k",

264 + ".v_proj": ".qkv_proj.v",

265 + }

266 + )

267 + # AutoWeightsLoader does `mapper |= cache_scale_mapper`

268 + combined = model_mapper | self._mapper()

269 +

270 + assert (

271 + combined._map_name("model.layers.0.self_attn.q_proj.weight")

272 + == "model.layers.0.self_attn.qkv_proj.q.weight"

273 + )

274 + assert (

275 + combined._map_name("model.layers.0.self_attn.k_proj.k_scale")

276 + == "model.layers.0.self_attn.attn.k_scale"

277 + )

278 + assert (

279 + combined._map_name("model.layers.0.self_attn.k_scale")

280 + == "model.layers.0.self_attn.attn.k_scale"

281 + )

282 +

283 +

284 if __name__ == "__main__":

285 test_download_weights_from_hf()

vllm/lora/worker_manager.py

+5 -1

@@ -128,9 +128,13 @@ class WorkerLoRAManager:

128 peft_helper.validate_legal(self.lora_config)

129

130 # For some models like Qwen2VL, we need to use hf_to_vllm_mapper

131 - # to ensure correct loading of lora weights.

131 + # to ensure correct loading of lora weights. Drop the QKV/MLP fusion

132 + # substr maps so constituent names (e.g. `q_proj`) survive for the

133 + # LoRA manager to pack, while keeping genuine renames/prefixes.

134 model = self._adapter_manager.model

135 hf_to_vllm_mapper = getattr(model, "hf_to_vllm_mapper", None)

136 + if hf_to_vllm_mapper is not None:

137 + hf_to_vllm_mapper = hf_to_vllm_mapper.get_unstacked_mapper()

138

139 # Get model-defined prefixes to skip during LoRA loading.

140 lora_skip_prefixes = getattr(model, "lora_skip_prefixes", None)

vllm/model_executor/layers/linear.py

+85 -37

@@ -3,9 +3,12 @@

3

4 import itertools

5 from abc import abstractmethod

6 +from collections.abc import Iterable

7 +from typing import Any

8

9 import torch

10 from torch.nn.parameter import Parameter

11 +from typing_extensions import TypeIs

12

13 import vllm.envs as envs

14 from vllm.distributed import (

@@ -632,31 +635,31 @@ class MergedColumnParallelLinear(ColumnParallelLinear):

635 disable_tp=disable_tp,

636 )

637

635 - def validate_shard_id(self, loaded_shard_id: int | tuple[int, ...] | None):

636 - if loaded_shard_id is None:

637 - return

638 - if isinstance(loaded_shard_id, tuple):

639 - for idx in loaded_shard_id:

638 + def validate_shard_id(self, shard_id: Any) -> TypeIs[int | tuple[int, ...] | None]:

639 + if isinstance(shard_id, int):

640 + if shard_id < 0 or shard_id >= len(self.output_sizes):

641 + raise ValueError(

642 + f"Shard id should be between 0 and {len(self.output_sizes) - 1}. "

643 + f"Got shard id {shard_id}."

644 + )

645 + return True

646 + if shard_id is None:

647 + return True

648 + if isinstance(shard_id, tuple):

649 + for idx in shard_id:

650 if not (0 <= idx < len(self.output_sizes)):

651 raise ValueError(

652 f"Shard id index {idx} should be between 0 and "

643 - f"{len(self.output_sizes) - 1}. Got shard id {loaded_shard_id}."

653 + f"{len(self.output_sizes) - 1}. Got shard id {shard_id}."

654 )

645 - if len(loaded_shard_id) > 1 and any(

646 - b - a != 1 for a, b in zip(loaded_shard_id[:-1], loaded_shard_id[1:])

655 + if len(shard_id) > 1 and any(

656 + b - a != 1 for a, b in zip(shard_id[:-1], shard_id[1:])

657 ):

658 raise ValueError(

659 "Shard id with multiple indices should be consecutive. "

650 - f"Got shard id {loaded_shard_id}."

660 + f"Got shard id {shard_id}."

661 )

652 - return

653 - elif isinstance(loaded_shard_id, int):

654 - if loaded_shard_id < 0 or loaded_shard_id >= len(self.output_sizes):

655 - raise ValueError(

656 - f"Shard id should be between 0 and {len(self.output_sizes) - 1}. "

657 - f"Got shard id {loaded_shard_id}."

658 - )

659 - return

662 + return True

663 raise ValueError("This line should not be reached")

664

665 def weight_loader(

@@ -910,6 +913,31 @@ class MergedColumnParallelLinear(ColumnParallelLinear):

913 tp_rank=self.tp_rank,

914 )

915

916 + def load_weights(

917 + self, weights: Iterable[tuple[str, torch.Tensor]]

918 + ) -> Iterable[str]:

919 + for name, loaded_weight in weights:

920 + shard_id = getattr(loaded_weight, "shard_id", None)

921 + self.validate_shard_id(shard_id)

922 + # Load into self if name is not an attr of self or its submodules

923 + param: Parameter

924 + if "." in name:

925 + submodule, _, attr = name.rpartition(".")

926 + param = getattr(self.get_submodule(submodule), attr, self)

927 + else:

928 + param = getattr(self, name, self)

929 + if param is None and name == "bias":

930 + continue

931 + param.weight_loader(param, loaded_weight, shard_id)

932 + logger.debug(

933 + "Loaded shard %s with shape %s into %s.%s",

934 + shard_id,

935 + loaded_weight.shape,

936 + self.prefix,

937 + name,

938 + )

939 + yield name

940 +

941

942 class QKVParallelLinear(ColumnParallelLinear):

943 """Linear layers for the attention's QKV transformation.

@@ -996,17 +1024,13 @@ class QKVParallelLinear(ColumnParallelLinear):

1024 disable_tp=disable_tp,

1025 )

1026

999 - def validate_shard_id(self, loaded_shard_id: str | None):

1000 - if loaded_shard_id is None:

1001 - return

1002 - if isinstance(loaded_shard_id, str):

1003 - if loaded_shard_id not in ["q", "k", "v"]:

1004 - raise ValueError(

1005 - "Shard id for QKVParallelLinear should be 'q', 'k', or 'v', "

1006 - f"got shard id {loaded_shard_id}."

1007 - )

1008 - return

1009 - raise ValueError("This line should not be reached")

1027 + def validate_shard_id(self, shard_id: Any) -> TypeIs[str | None]:

1028 + if shard_id in {"q", "k", "v"} or shard_id is None:

1029 + return True

1030 + raise ValueError(

1031 + "Shard id for QKVParallelLinear should be 'q', 'k', or 'v', "

1032 + f"got shard id {shard_id}."

1033 + )

1034

1035 def _get_shard_offset_mapping(self, loaded_shard_id: str):

1036 shard_offset_mapping = {

@@ -1302,6 +1326,31 @@ class QKVParallelLinear(ColumnParallelLinear):

1326 assert param_data.shape == loaded_weight.shape

1327 param_data.copy_(loaded_weight)

1328

1329 + def load_weights(

1330 + self, weights: Iterable[tuple[str, torch.Tensor]]

1331 + ) -> Iterable[str]:

1332 + for name, loaded_weight in weights:

1333 + shard_id = getattr(loaded_weight, "shard_id", None)

1334 + self.validate_shard_id(shard_id)

1335 + # Load into self if name is not an attr of self or its submodules

1336 + param: Parameter

1337 + if "." in name:

1338 + submodule, _, attr = name.rpartition(".")

1339 + param = getattr(self.get_submodule(submodule), attr, self)

1340 + else:

1341 + param = getattr(self, name, self)

1342 + if param is None and name == "bias":

1343 + continue

1344 + param.weight_loader(param, loaded_weight, shard_id)

1345 + logger.debug(

1346 + "Loaded shard %s with shape %s into %s.%s",

1347 + shard_id,

1348 + loaded_weight.shape,

1349 + self.prefix,

1350 + name,

1351 + )

1352 + yield name

1353 +

1354

1355 class MinimaxM3QKVParallelLinearWithIndexer(QKVParallelLinear):

1356 """QKV projection fused with a lightning-indexer's index_q/index_k.

@@ -1387,15 +1436,14 @@ class MinimaxM3QKVParallelLinearWithIndexer(QKVParallelLinear):

1436 prefix=prefix,

1437 )

1438

1390 - def validate_shard_id(self, loaded_shard_id: str | None) -> None:

1391 - if loaded_shard_id is None:

1392 - return

1393 - if loaded_shard_id not in ("q", "k", "v", "index_q", "index_k"):

1394 - raise ValueError(

1395 - "Shard id for MinimaxM3QKVParallelLinearWithIndexer must be one of "

1396 - "'q', 'k', 'v', 'index_q', 'index_k'; got "

1397 - f"{loaded_shard_id}."

1398 - )

1439 + def validate_shard_id(self, shard_id: Any) -> TypeIs[str | None]:

1440 + if shard_id in {"q", "k", "v", "index_q", "index_k"} or shard_id is None:

1441 + return True

1442 + raise ValueError(

1443 + "Shard id for MinimaxM3QKVParallelLinearWithIndexer must be one of "

1444 + "'q', 'k', 'v', 'index_q', 'index_k'; got "

1445 + f"{shard_id}."

1446 + )

1447

1448 def _get_shard_offset_mapping(self, loaded_shard_id: str) -> int | None:

1449 h = self.head_size

vllm/model_executor/layers/quantization/base_config.py

+46 -5

@@ -5,6 +5,7 @@ import inspect

5 from abc import ABC, abstractmethod

6 from typing import TYPE_CHECKING, Any

7

8 +import regex as re

9 import torch

10 from torch import nn

11 from transformers import PretrainedConfig

@@ -19,10 +20,12 @@ else:

20 class QuantizeMethodBase(ABC):

21 """Base class for different quantized methods."""

22

22 - # Whether this method creates weights on meta device for online quantization.

23 - # When True, weights are created on meta device and quantized layer-wise

24 - # in process_weights_after_loading, reducing peak memory during loading.

23 uses_meta_device: bool = False

24 + """

25 + Whether this method creates weights on meta device for online quantization.

26 + When True, weights are created on meta device and quantized layer-wise

27 + in process_weights_after_loading, reducing peak memory during loading.

28 + """

29

30 @abstractmethod

31 def create_weights(

@@ -84,6 +87,18 @@ def method_has_implemented_embedding(method_class: type[QuantizeMethodBase]) ->

87 class QuantizationConfig(ABC):

88 """Base class for quantization configs."""

89

90 + _ignore_unexpected_suffixes = (

91 + ".q_scale",

92 + ".k_scale",

93 + ".v_scale",

94 + ".q_zero_point",

95 + ".k_zero_point",

96 + ".v_zero_point",

97 + )

98 + """Suffixes of quantization parameters that may be present in the checkpoint but

99 + not in the model, and should be ignored if unexpected during loading. These are used

100 + after remapping, so should be in vLLM format (e.g. .q_scale, not .q.scale)."""

101 +

102 def __init__(self):

103 super().__init__()

104 # mapping is updated by models as they initialize

@@ -176,14 +191,40 @@ class QuantizationConfig(ABC):

191 """

192 raise NotImplementedError

193

179 - def get_cache_scale_mapper(self) -> "WeightsMapper | None":

194 + @staticmethod

195 + def get_cache_scale_mapper() -> "WeightsMapper":

196 """Mapping from checkpoint KV-cache scale names to vLLM scale names.

197

198 Returning a mapper here causes `AutoWeightsLoader` to apply it to the

199 weight stream automatically; individual model `load_weights` methods

200 do not need to know about KV-cache scales.

201 """

186 - return None

202 + from vllm.model_executor.models.utils import WeightsMapper

203 +

204 + orig_to_new_regex = {

205 + # Deprecated fused kv_scale -> attn.k_scale

206 + re.compile(r"\.kv_scale$"): r".attn.k_scale",

207 + # ModelOpt: .self_attn.{k,v}_proj.{k,v}_scale -> .self_attn.attn.*

208 + re.compile(r"\.self_attn\.[kv]_proj\.([kv])_scale$"): (

209 + r".self_attn.attn.\1_scale"

210 + ),

211 + # Fused QKV / qkqkv proj: .self_attn.qk(qk)v_proj.{k,v}_scale -> attn

212 + re.compile(r"\.self_attn\.qk(?:qk)?v_proj\.([kv])_scale$"): (

213 + r".self_attn.attn.\1_scale"

214 + ),

215 + # NemotronH: .mixer.{k,v}_proj.{k,v}_scale -> .mixer.attn.*

216 + re.compile(r"\.mixer\.[kv]_proj\.([kv])_scale$"): r".mixer.attn.\1_scale",

217 + # HYV3: .self_attn.q.scale -> .self_attn.attn.q_scale

218 + re.compile(r"\.self_attn\.q\.scale$"): r".self_attn.attn.q_scale",

219 + # HYV3: .self_attn.{k,v}_cache.scale -> .self_attn.attn.{k,v}_scale

220 + re.compile(r"\.self_attn\.([kv])_cache\.scale$"): (

221 + r".self_attn.attn.\1_scale"

222 + ),

223 + # Default: .{q,k,v}_scale -> .attn.{q,k,v}_scale (unless already .attn)

224 + re.compile(r"(?<!\.attn)\.([qkv])_scale$"): r".attn.\1_scale",

225 + re.compile(r"(?<!\.attn)\.([qkv])_zero_point$"): r".attn.\1_zero_point",

226 + }

227 + return WeightsMapper(orig_to_new_regex=orig_to_new_regex)

228

229 def apply_vllm_mapper( # noqa: B027

230 self, hf_to_vllm_mapper: "WeightsMapper"

vllm/model_executor/layers/quantization/fp8.py

+10 -9

@@ -219,18 +219,19 @@ class Fp8Config(QuantizationConfig):

219 return Fp8KVCacheMethod(self)

220 return None

221

222 - def get_cache_scale_mapper(self) -> "WeightsMapper":

222 + @staticmethod

223 + def get_cache_scale_mapper() -> "WeightsMapper":

224 """Map compressed-tensors KV-cache scale names to vLLM names."""

225 from vllm.model_executor.models.utils import WeightsMapper

226

226 - return WeightsMapper(

227 - orig_to_new_suffix={

228 - ".k_proj.output_scale": ".attn.k_scale",

229 - ".v_proj.output_scale": ".attn.v_scale",

230 - ".q_proj.output_scale": ".attn.q_scale",

231 - ".self_attn.prob_output_scale": ".self_attn.attn.prob_scale",

232 - }

233 - )

227 + orig_to_new_suffix = {

228 + ".k_proj.output_scale": ".attn.k_scale",

229 + ".v_proj.output_scale": ".attn.v_scale",

230 + ".q_proj.output_scale": ".attn.q_scale",

231 + ".self_attn.prob_output_scale": ".self_attn.attn.prob_scale",

232 + }

233 + cache_scale_mapper = WeightsMapper(orig_to_new_suffix=orig_to_new_suffix)

234 + return cache_scale_mapper | QuantizationConfig.get_cache_scale_mapper()

235

236

237 class CopyNumelCounter(TorchDispatchMode):

vllm/model_executor/layers/quantization/quark/quark.py

+10 -9

@@ -679,16 +679,17 @@ class QuarkConfig(QuantizationConfig):

679

680 return scheme

681

682 - def get_cache_scale_mapper(self) -> "WeightsMapper":

682 + @staticmethod

683 + def get_cache_scale_mapper() -> "WeightsMapper":

684 """Map Quark KV-cache scale names to vLLM names."""

684 - return WeightsMapper(

685 - orig_to_new_suffix={

686 - ".k_proj.output_scale": ".attn.k_scale",

687 - ".v_proj.output_scale": ".attn.v_scale",

688 - ".q_proj.output_scale": ".attn.q_scale",

689 - ".self_attn.prob_output_scale": ".self_attn.attn.prob_scale",

690 - }

691 - )

685 + orig_to_new_suffix = {

686 + ".k_proj.output_scale": ".attn.k_scale",

687 + ".v_proj.output_scale": ".attn.v_scale",

688 + ".q_proj.output_scale": ".attn.q_scale",

689 + ".self_attn.prob_output_scale": ".self_attn.attn.prob_scale",

690 + }

691 + cache_scale_mapper = WeightsMapper(orig_to_new_suffix=orig_to_new_suffix)

692 + return cache_scale_mapper | QuantizationConfig.get_cache_scale_mapper()

693

694

695 class QuarkLinearMethod(LinearMethodBase):

vllm/model_executor/model_loader/bitsandbytes_loader.py

+2 -1

@@ -576,7 +576,8 @@ class BitsAndBytesModelLoader(BaseModelLoader):

576 # For some models like Molmo, we need to use hf_to_vllm_mapper

577 # to ensure correct loading of weights.

578 if hf_to_vllm_mapper := getattr(model, "hf_to_vllm_mapper", None):

579 - self.weight_mapper = lambda name: hf_to_vllm_mapper._map_name(name)

579 + unstacked_mapper = hf_to_vllm_mapper.get_unstacked_mapper()

580 + self.weight_mapper = lambda name, m=unstacked_mapper: m._map_name(name)

581

582 self._get_bnb_target_modules(model)

583 self._classify_module_sharding(model)

vllm/model_executor/model_loader/reload/layerwise.py

+10 -1

@@ -131,8 +131,11 @@ def initialize_online_processing(layer: torch.nn.Module):

131 # Track loading progress to determine when to process/copy

132 info.load_numel = 0

133 info.load_numel_total = get_layer_size(layer)

134 + _wrap_parameters_weight_loader(layer)

135

135 - # Wrap each parameter's weight loader

136 +

137 +def _wrap_parameters_weight_loader(layer: torch.nn.Module) -> None:

138 + """Wrap each parameter's weight loader."""

139 # Note that nested wrapping will occur for shared tensors

140 for name, tensor in get_layer_tensors(layer).items():

141 if name in SKIP_TENSORS:

@@ -168,6 +171,12 @@ def make_online_process_loader(layer: torch.nn.Module, param_name: str) -> Calla

171 logger.debug("%s: Excessive loading", layer.__class__.__name__)

172 return

173

174 + # Re-run on each load: layers may register parameters later (e.g., `bias`).

175 + # Wrap late parameters and refresh `load_numel_total` so processing waits

176 + # until all parameters are loaded.

177 + info.load_numel_total = get_layer_size(layer)

178 + _wrap_parameters_weight_loader(layer)

179 +

180 # Bind and normalize arguments

181 bound_args = loader_signature.bind(*args, **kwargs)

182 bound_args.apply_defaults()

vllm/model_executor/model_loader/utils.py

+1 -1

@@ -290,6 +290,6 @@ def configure_quant_config(

290

291 # pass mappings by reference to quant_config

292 if hf_to_vllm_mapper is not None:

293 - quant_config.apply_vllm_mapper(hf_to_vllm_mapper)

293 + quant_config.apply_vllm_mapper(hf_to_vllm_mapper.get_unstacked_mapper())

294 if packed_mapping is not None:

295 quant_config.packed_modules_mapping = packed_mapping

vllm/model_executor/models/arcee.py

+10 -67

@@ -26,10 +26,6 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (

26 ParallelLMHead,

27 VocabParallelEmbedding,

28 )

29 -from vllm.model_executor.model_loader.weight_utils import (

30 - default_weight_loader,

31 - maybe_remap_kv_scale_name,

32 -)

29 from vllm.sequence import IntermediateTensors

30

31 from .interfaces import (

@@ -42,7 +38,7 @@ from .interfaces import (

38 from .utils import (

39 AutoWeightsLoader,

40 PPMissingLayer,

45 - is_pp_missing_parameter,

41 + WeightsMapper,

42 make_empty_intermediate_tensors_factory,

43 make_layers,

44 maybe_prefix,

@@ -276,67 +272,6 @@ class ArceeModel(nn.Module, EagleModelMixin):

272 return hidden_states, aux_hidden_states

273 return hidden_states

274

279 - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

280 - """Load weights, mapping q/k/v projections to fused qkv_proj."""

281 - stacked_params_mapping = [

282 - (".qkv_proj", ".q_proj", "q"),

283 - (".qkv_proj", ".k_proj", "k"),

284 - (".qkv_proj", ".v_proj", "v"),

285 - ]

286 -

287 - params_dict = dict(self.named_parameters())

288 - loaded_params: set[str] = set()

289 -

290 - for name, loaded_weight in weights:

291 - if "rotary_emb.inv_freq" in name:

292 - continue

293 - if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name:

294 - continue

295 -

296 - if "scale" in name or "zero_point" in name:

297 - remapped_name = maybe_remap_kv_scale_name(name, params_dict)

298 - if remapped_name is None:

299 - continue

300 - name = remapped_name

301 -

302 - mapped = False

303 - for param_name, weight_name, shard_id in stacked_params_mapping:

304 - if weight_name not in name:

305 - continue

306 -

307 - name = name.replace(weight_name, param_name)

308 -

309 - if name.endswith(".bias") and name not in params_dict:

310 - mapped = True

311 - break

312 -

313 - if is_pp_missing_parameter(name, self):

314 - mapped = True

315 - break

316 -

317 - param = params_dict[name]

318 - weight_loader = param.weight_loader # type: ignore[attr-defined]

319 - weight_loader(param, loaded_weight, shard_id)

320 - loaded_params.add(name)

321 - mapped = True

322 - break

323 -

324 - if mapped:

325 - continue

326 -

327 - if name.endswith(".bias") and name not in params_dict:

328 - continue

329 -

330 - if is_pp_missing_parameter(name, self):

331 - continue

332 -

333 - param = params_dict[name]

334 - weight_loader = getattr(param, "weight_loader", default_weight_loader)

335 - weight_loader(param, loaded_weight)

336 - loaded_params.add(name)

337 -

338 - return loaded_params

339 -

275

276 class ArceeForCausalLM(

277 nn.Module, SupportsLoRA, SupportsPP, SupportsEagle, SupportsEagle3

@@ -344,6 +279,14 @@ class ArceeForCausalLM(

279 """Arcee Model for causal language modeling, integrated with vLLM

280 runtime."""

281

282 + hf_to_vllm_mapper = WeightsMapper(

283 + orig_to_new_stacked={

284 + # weight_name: (param_name, shard_id)

285 + ".q_proj": (".qkv_proj", "q"),

286 + ".k_proj": (".qkv_proj", "k"),

287 + ".v_proj": (".qkv_proj", "v"),

288 + }

289 + )

290 # Map fused module names to their submodule components

291 # (for quantization and LoRA)

292 packed_modules_mapping = {

@@ -420,4 +363,4 @@ class ArceeForCausalLM(

363 )

364 # AutoWeightLoader handles weight name remapping, including fusing

365 # separate q_proj, k_proj, v_proj into qkv_proj

423 - return loader.load_weights(weights)

366 + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

vllm/model_executor/models/chatglm.py

+6 -47

@@ -30,7 +30,6 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (

30 ParallelLMHead,

31 VocabParallelEmbedding,

32 )

33 -from vllm.model_executor.model_loader.weight_utils import default_weight_loader

33 from vllm.sequence import IntermediateTensors

34 from vllm.transformers_utils.configs.chatglm import ChatGLMConfig

35

@@ -38,7 +37,6 @@ from .interfaces import SupportsLoRA, SupportsPP, SupportsQuant

37 from .utils import (

38 AutoWeightsLoader,

39 WeightsMapper,

41 - is_pp_missing_parameter,

40 make_empty_intermediate_tensors_factory,

41 make_layers,

42 maybe_prefix,

@@ -316,12 +314,9 @@ class GLMTransformer(nn.Module):

314

315 @support_torch_compile

316 class ChatGLMModel(nn.Module, SupportsQuant):

319 - packed_modules_mapping = {

320 - "linear_proj.merged_proj": [

321 - "linear_proj.gate_proj",

322 - "linear_proj.dense_h_to_4h",

323 - ]

324 - }

317 + hf_to_vllm_mapper = WeightsMapper(

318 + orig_to_new_substr={".word_embeddings": ""},

319 + )

320

321 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

322 super().__init__()

@@ -386,47 +381,11 @@ class ChatGLMModel(nn.Module, SupportsQuant):

381 return hidden_states

382

383 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

389 - stacked_params_mapping = [

390 - # (param_name, shard_name, shard_id)

391 - ("linear_proj.merged_proj", "linear_proj.gate_proj", 0),

392 - ("linear_proj.merged_proj", "linear_proj.dense_h_to_4h", 1),

393 - ]

394 - params_dict = dict(self.named_parameters())

395 - loaded_params: set[str] = set()

396 -

397 - for name, loaded_weight in weights:

398 - for param_name, weight_name, shard_id in stacked_params_mapping:

399 - if weight_name not in name:

400 - continue

401 - name = name.replace(weight_name, param_name)

402 - # Skip loading extra bias for GPTQ models.

403 - if name.endswith(".bias") and name not in params_dict:

404 - continue

405 - if is_pp_missing_parameter(name, self):

406 - continue

407 - param = params_dict[name]

408 - weight_loader = param.weight_loader

409 - weight_loader(param, loaded_weight, shard_id)

410 - break

411 - else:

412 - if "rotary_pos_emb.inv_freq" in name:

413 - continue

414 - if name.endswith(".bias") and name not in params_dict:

415 - continue

416 - if is_pp_missing_parameter(name, self):

417 - continue

418 - param = params_dict[name]

419 - weight_loader = getattr(param, "weight_loader", default_weight_loader)

420 - weight_loader(param, loaded_weight)

421 - loaded_params.add(name)

422 - return loaded_params

384 + loader = AutoWeightsLoader(self)

385 + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

386

387

388 class ChatGLMBaseModel(nn.Module):

426 - hf_to_vllm_mapper = WeightsMapper(

427 - orig_to_new_substr={".word_embeddings": ""},

428 - )

429 -

389 def __init__(

390 self,

391 *,

@@ -467,7 +426,7 @@ class ChatGLMBaseModel(nn.Module):

426

427 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):

428 loader = AutoWeightsLoader(self)

470 - return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

429 + return loader.load_weights(weights)

430

431

432 class ChatGLMForCausalLM(ChatGLMBaseModel, SupportsLoRA, SupportsPP, SupportsQuant):

vllm/model_executor/models/cohere_eagle.py

+3 -38

@@ -14,7 +14,6 @@ from vllm.model_executor.layers.linear import ReplicatedLinear

14 from vllm.model_executor.layers.logits_processor import LogitsProcessor

15 from vllm.model_executor.layers.quantization.base_config import QuantizationConfig

16 from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding

17 -from vllm.model_executor.model_loader.weight_utils import default_weight_loader

17 from vllm.model_executor.models.commandr import (

18 CohereDecoderLayer,

19 CohereForCausalLM,

@@ -134,42 +133,6 @@ class CohereEagleModel(nn.Module):

133 hidden_states, _ = self.norm(hidden_states, residual)

134 return hidden_states, hidden_states

135

137 - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

138 - stacked_params_mapping = [

139 - # (param_name, shard_name, shard_id)

140 - (".qkv_proj", ".q_proj", "q"),

141 - (".qkv_proj", ".k_proj", "k"),

142 - (".qkv_proj", ".v_proj", "v"),

143 - (".gate_up_proj", ".gate_proj", 0),

144 - (".gate_up_proj", ".up_proj", 1),

145 - ]

146 - params_dict = dict(self.named_parameters())

147 - loaded_params: set[str] = set()

148 -

149 - for name, loaded_weight in weights:

150 - if "rotary_emb.inv_freq" in name:

151 - continue

152 -

153 - for param_name, weight_name, shard_id in stacked_params_mapping:

154 - if weight_name not in name:

155 - continue

156 - name = name.replace(weight_name, param_name)

157 - if name.endswith(".bias") and name not in params_dict:

158 - continue

159 - param = params_dict[name]

160 - weight_loader = param.weight_loader

161 - weight_loader(param, loaded_weight, shard_id)

162 - break

163 - else:

164 - if name.endswith(".bias") and name not in params_dict:

165 - continue

166 -

167 - param = params_dict[name]

168 - weight_loader = getattr(param, "weight_loader", default_weight_loader)

169 - weight_loader(param, loaded_weight)

170 - loaded_params.add(name)

171 - return loaded_params

172 -

136

137 class EagleCohereForCausalLM(CohereForCausalLM):

138 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

@@ -225,7 +188,9 @@ class EagleCohereForCausalLM(CohereForCausalLM):

188 ),

189 )

190

228 - loaded_weight_names = loader.load_weights(map(_track_and_forward, weights))

191 + loaded_weight_names = loader.load_weights(

192 + map(_track_and_forward, weights), mapper=self.hf_to_vllm_mapper

193 + )

194

195 # Embed tokens are tied with the target model and therefore not

196 # present in the EAGLE checkpoint; mark them as loaded explicitly to

vllm/model_executor/models/commandr.py

+12 -54

@@ -45,8 +45,6 @@ from vllm.model_executor.layers.quantization import QuantizationConfig

45 from vllm.model_executor.layers.rotary_embedding import get_rope

46 from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding

47 from vllm.model_executor.model_loader.weight_utils import (

48 - default_weight_loader,

49 - maybe_remap_kv_scale_name,

48 row_parallel_weight_loader,

49 )

50 from vllm.model_executor.utils import set_weight_attrs

@@ -58,7 +56,6 @@ from .utils import (

56 AutoWeightsLoader,

57 WeightsMapper,

58 extract_layer_index,

61 - is_pp_missing_parameter,

59 make_empty_intermediate_tensors_factory,

60 make_layers,

61 maybe_prefix,

@@ -341,60 +338,21 @@ class CohereModel(nn.Module):

338 hidden_states, _ = self.norm(hidden_states, residual)

339 return hidden_states

340

344 - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

345 - stacked_params_mapping = [

346 - # (param_name, shard_name, shard_id)

347 - ("qkv_proj", "q_proj", "q"),

348 - ("qkv_proj", "k_proj", "k"),

349 - ("qkv_proj", "v_proj", "v"),

350 - ("gate_up_proj", "gate_proj", 0),

351 - ("gate_up_proj", "up_proj", 1),

352 - ]

353 - params_dict = dict(self.named_parameters())

354 - loaded_params: set[str] = set()

355 - for name, loaded_weight in weights:

356 - for param_name, shard_name, shard_id in stacked_params_mapping:

357 - if shard_name not in name:

358 - continue

359 - name = name.replace(shard_name, param_name)

360 - # Skip loading extra bias for GPTQ models.

361 - if name.endswith(".bias") and name not in params_dict:

362 - continue

363 - if is_pp_missing_parameter(name, self):

364 - continue

365 - param = params_dict[name]

366 - weight_loader = param.weight_loader

367 - weight_loader(param, loaded_weight, shard_id)

368 - break

369 - else:

370 - # Skip loading extra bias for GPTQ models.

371 - if name.endswith(".bias") and name not in params_dict:

372 - continue

373 - # Remapping the name of FP8 kv-scale.

374 - name = maybe_remap_kv_scale_name(name, params_dict)

375 - if name is None:

376 - continue

377 -

378 - if is_pp_missing_parameter(name, self):

379 - continue

380 - param = params_dict[name]

381 - weight_loader = getattr(param, "weight_loader", default_weight_loader)

382 - weight_loader(param, loaded_weight)

383 - loaded_params.add(name)

384 - return loaded_params

385 -

341

342 class CohereForCausalLM(nn.Module, SupportsLoRA, SupportsPP, SupportsQuant):

343 + hf_to_vllm_mapper = WeightsMapper(

344 + orig_to_new_stacked={

345 + # weight_name: (param_name, shard_id)

346 + ".q_proj": (".qkv_proj", "q"),

347 + ".k_proj": (".qkv_proj", "k"),

348 + ".v_proj": (".qkv_proj", "v"),

349 + ".gate_proj": (".gate_up_proj", 0),

350 + ".up_proj": (".gate_up_proj", 1),

351 + }

352 + )

353 packed_modules_mapping = {

389 - "qkv_proj": [

390 - "q_proj",

391 - "k_proj",

392 - "v_proj",

393 - ],

394 - "gate_up_proj": [

395 - "gate_proj",

396 - "up_proj",

397 - ],

354 + "qkv_proj": ["q_proj", "k_proj", "v_proj"],

355 + "gate_up_proj": ["gate_proj", "up_proj"],

356 }

357 # LoRA specific attributes

358 embedding_modules = {"embed_tokens": "input_embeddings"}

vllm/model_executor/models/exaone.py

+14 -67

@@ -50,17 +50,13 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (

50 ParallelLMHead,

51 VocabParallelEmbedding,

52 )

53 -from vllm.model_executor.model_loader.weight_utils import (

54 - default_weight_loader,

55 - maybe_remap_kv_scale_name,

56 -)

53 from vllm.sequence import IntermediateTensors

54

55 from .interfaces import SupportsLoRA, SupportsPP

56 from .utils import (

57 AutoWeightsLoader,

58 PPMissingLayer,

63 - is_pp_missing_parameter,

59 + WeightsMapper,

60 make_empty_intermediate_tensors_factory,

61 make_layers,

62 maybe_prefix,

@@ -370,70 +366,21 @@ class ExaoneModel(nn.Module):

366 hidden_states, _ = self.ln_f(hidden_states, residual)

367 return hidden_states

368

373 - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

374 - stacked_params_mapping = [

375 - # (param_name, shard_name, shard_id)

376 - (".qkv_proj", ".q_proj", "q"),

377 - (".qkv_proj", ".k_proj", "k"),

378 - (".qkv_proj", ".v_proj", "v"),

379 - (".gate_up_proj", ".c_fc_0", 0),

380 - (".gate_up_proj", ".c_fc_1", 1),

381 - ]

382 - params_dict = dict(self.named_parameters())

383 - loaded_params: set[str] = set()

384 - for name, loaded_weight in weights:

385 - if "rotary_emb.inv_freq" in name:

386 - continue

387 - if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name:

388 - # Models trained using ColossalAI may include these tensors in

389 - # the checkpoint. Skip them.

390 - continue

391 - for param_name, weight_name, shard_id in stacked_params_mapping:

392 - if weight_name not in name:

393 - continue

394 - name = name.replace(weight_name, param_name)

395 - # Skip loading extra bias for GPTQ models.

396 - if name.endswith(".bias") and name not in params_dict:

397 - continue

398 -

399 - if is_pp_missing_parameter(name, self):

400 - continue

401 -

402 - param = params_dict[name]

403 - weight_loader = param.weight_loader

404 - weight_loader(param, loaded_weight, shard_id)

405 -

406 - break

407 - else:

408 - # Skip loading extra bias for GPTQ models.

409 - if name.endswith(".bias") and name not in params_dict:

410 - continue

411 - # Remapping the name of FP8 kv-scale.

412 - name = maybe_remap_kv_scale_name(name, params_dict)

413 - if name is None:

414 - continue

415 -

416 - if is_pp_missing_parameter(name, self):

417 - continue

418 -

419 - param = params_dict[name]

420 - weight_loader = getattr(param, "weight_loader", default_weight_loader)

421 - weight_loader(param, loaded_weight)

422 - loaded_params.add(name)

423 - return loaded_params

424 -

369

370 class ExaoneForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

371 + hf_to_vllm_mapper = WeightsMapper(

372 + orig_to_new_stacked={

373 + # weight_name: (param_name, shard_id)

374 + ".q_proj": (".qkv_proj", "q"),

375 + ".k_proj": (".qkv_proj", "k"),

376 + ".v_proj": (".qkv_proj", "v"),

377 + ".c_fc_0": (".gate_up_proj", 0),

378 + ".c_fc_1": (".gate_up_proj", 1),

379 + }

380 + )

381 packed_modules_mapping = {

428 - "qkv_proj": [

429 - "q_proj",

430 - "k_proj",

431 - "v_proj",

432 - ],

433 - "gate_up_proj": [

434 - "c_fc_0",

435 - "c_fc_1",

436 - ],

382 + "qkv_proj": ["q_proj", "k_proj", "v_proj"],

383 + "gate_up_proj": ["c_fc_0", "c_fc_1"],

384 }

385

386 # LoRA specific attributes

@@ -506,4 +453,4 @@ class ExaoneForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

453 # processed with quantization, LoRA, fine-tuning, etc.

454 skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),

455 )

509 - return loader.load_weights(weights)

456 + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

vllm/model_executor/models/exaone4.py

+14 -67

@@ -46,10 +46,6 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (

46 ParallelLMHead,

47 VocabParallelEmbedding,

48 )

49 -from vllm.model_executor.model_loader.weight_utils import (

50 - default_weight_loader,

51 - maybe_remap_kv_scale_name,

52 -)

49 from vllm.sequence import IntermediateTensors

50 from vllm.transformers_utils.config import set_default_rope_theta

51

@@ -57,8 +53,8 @@ from .interfaces import SupportsLoRA, SupportsPP

53 from .utils import (

54 AutoWeightsLoader,

55 PPMissingLayer,

56 + WeightsMapper,

57 extract_layer_index,

61 - is_pp_missing_parameter,

58 make_empty_intermediate_tensors_factory,

59 make_layers,

60 maybe_prefix,

@@ -368,70 +364,21 @@ class Exaone4Model(nn.Module):

364 hidden_states = self.norm(hidden_states)

365 return hidden_states

366

371 - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

372 - stacked_params_mapping = [

373 - # (param_name, shard_name, shard_id)

374 - (".qkv_proj", ".q_proj", "q"),

375 - (".qkv_proj", ".k_proj", "k"),

376 - (".qkv_proj", ".v_proj", "v"),

377 - (".gate_up_proj", ".gate_proj", 0),

378 - (".gate_up_proj", ".up_proj", 1),

379 - ]

380 - params_dict = dict(self.named_parameters())

381 - loaded_params: set[str] = set()

382 - for name, loaded_weight in weights:

383 - if "rotary_emb.inv_freq" in name:

384 - continue

385 - if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name:

386 - # Models trained using ColossalAI may include these tensors in

387 - # the checkpoint. Skip them.

388 - continue

389 - for param_name, weight_name, shard_id in stacked_params_mapping:

390 - if weight_name not in name:

391 - continue

392 - name = name.replace(weight_name, param_name)

393 - # Skip loading extra bias for GPTQ models.

394 - if name.endswith(".bias") and name not in params_dict:

395 - continue

396 -

397 - if is_pp_missing_parameter(name, self):

398 - continue

399 -

400 - param = params_dict[name]

401 - weight_loader = param.weight_loader

402 - weight_loader(param, loaded_weight, shard_id)

403 -

404 - break

405 - else:

406 - # Skip loading extra bias for GPTQ models.

407 - if name.endswith(".bias") and name not in params_dict:

408 - continue

409 - # Remapping the name of FP8 kv-scale.

410 - name = maybe_remap_kv_scale_name(name, params_dict)

411 - if name is None:

412 - continue

413 -

414 - if is_pp_missing_parameter(name, self):

415 - continue

416 -

417 - param = params_dict[name]

418 - weight_loader = getattr(param, "weight_loader", default_weight_loader)

419 - weight_loader(param, loaded_weight)

420 - loaded_params.add(name)

421 - return loaded_params

422 -

367

368 class Exaone4ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

369 + hf_to_vllm_mapper = WeightsMapper(

370 + orig_to_new_stacked={

371 + # weight_name: (param_name, shard_id)

372 + ".q_proj": (".qkv_proj", "q"),

373 + ".k_proj": (".qkv_proj", "k"),

374 + ".v_proj": (".qkv_proj", "v"),

375 + ".gate_proj": (".gate_up_proj", 0),

376 + ".up_proj": (".gate_up_proj", 1),

377 + }

378 + )

379 packed_modules_mapping = {

426 - "qkv_proj": [

427 - "q_proj",

428 - "k_proj",

429 - "v_proj",

430 - ],

431 - "gate_up_proj": [

432 - "gate_proj",

433 - "up_proj",

434 - ],

380 + "qkv_proj": ["q_proj", "k_proj", "v_proj"],

381 + "gate_up_proj": ["gate_proj", "up_proj"],

382 }

383

384 # LoRA specific attributes

@@ -503,4 +450,4 @@ class Exaone4ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

450 # processed with quantization, LoRA, fine-tuning, etc.

451 skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),

452 )

506 - return loader.load_weights(weights)

453 + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

vllm/model_executor/models/fairseq2_llama.py

+2 -4

@@ -79,10 +79,8 @@ class Fairseq2LlamaForCausalLM(LlamaForCausalLM):

79 skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),

80 )

81 return loader.load_weights(

82 - (

83 - self.reshape_fairseq2_weights(name, loaded_weight, params)

84 - for name, loaded_weight in weights

85 - )

82 + self.reshape_fairseq2_weights(name, loaded_weight, params)

83 + for name, loaded_weight in weights

84 )

85

86 def flag_sharded_weights(self, params: dict[str, Parameter]):

vllm/model_executor/models/gemma.py

+16 -52

@@ -42,13 +42,12 @@ from vllm.model_executor.layers.logits_processor import LogitsProcessor

42 from vllm.model_executor.layers.quantization import QuantizationConfig

43 from vllm.model_executor.layers.rotary_embedding import get_rope

44 from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding

45 -from vllm.model_executor.model_loader.weight_utils import default_weight_loader

45 from vllm.sequence import IntermediateTensors

46

48 -from .interfaces import SupportsLoRA, SupportsPP

47 +from .interfaces import SupportsLoRA, SupportsPP, SupportsQuant

48 from .utils import (

49 AutoWeightsLoader,

51 - is_pp_missing_parameter,

50 + WeightsMapper,

51 make_empty_intermediate_tensors_factory,

52 make_layers,

53 maybe_prefix,

@@ -324,56 +323,21 @@ class GemmaModel(nn.Module):

323 hidden_states, _ = self.norm(hidden_states, residual)

324 return hidden_states

325

327 - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

328 - stacked_params_mapping = [

329 - # (param_name, shard_name, shard_id)

330 - ("qkv_proj", "q_proj", "q"),

331 - ("qkv_proj", "k_proj", "k"),

332 - ("qkv_proj", "v_proj", "v"),

333 - ("gate_up_proj", "gate_proj", 0),

334 - ("gate_up_proj", "up_proj", 1),

335 - ]

336 - params_dict = dict(self.named_parameters())

337 - loaded_params: set[str] = set()

338 - for name, loaded_weight in weights:

339 - for param_name, shard_name, shard_id in stacked_params_mapping:

340 - if shard_name not in name:

341 - continue

342 - name = name.replace(shard_name, param_name)

343 - # Skip loading extra bias for GPTQ models.

344 - if name.endswith(".bias") and name not in params_dict:

345 - continue

346 - if is_pp_missing_parameter(name, self):

347 - continue

348 - param = params_dict[name]

349 - weight_loader = param.weight_loader

350 - weight_loader(param, loaded_weight, shard_id)

351 - break

352 - else:

353 - # Skip loading extra bias for GPTQ models.

354 - if name.endswith(".bias") and name not in params_dict:

355 - continue

356 - if is_pp_missing_parameter(name, self):

357 - continue

358 - param = params_dict[name]

359 - weight_loader = getattr(param, "weight_loader", default_weight_loader)

360 - weight_loader(param, loaded_weight)

361 - loaded_params.add(name)

362 -

363 - return loaded_params

364 -

326

366 -class GemmaForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

327 +class GemmaForCausalLM(nn.Module, SupportsLoRA, SupportsPP, SupportsQuant):

328 + hf_to_vllm_mapper = WeightsMapper(

329 + orig_to_new_stacked={

330 + # weight_name: (param_name, shard_id)

331 + ".q_proj": (".qkv_proj", "q"),

332 + ".k_proj": (".qkv_proj", "k"),

333 + ".v_proj": (".qkv_proj", "v"),

334 + ".gate_proj": (".gate_up_proj", 0),

335 + ".up_proj": (".gate_up_proj", 1),

336 + }

337 + )

338 packed_modules_mapping = {

368 - "qkv_proj": [

369 - "q_proj",

370 - "k_proj",

371 - "v_proj",

372 - ],

373 - "gate_up_proj": [

374 - "gate_proj",

375 - "up_proj",

376 - ],

339 + "qkv_proj": ["q_proj", "k_proj", "v_proj"],

340 + "gate_up_proj": ["gate_proj", "up_proj"],

341 }

342

343 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

@@ -421,4 +385,4 @@ class GemmaForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

385 self,

386 skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),

387 )

424 - return loader.load_weights(weights)

388 + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

vllm/model_executor/models/gemma2.py

+14 -57

@@ -39,17 +39,13 @@ from vllm.model_executor.layers.logits_processor import LogitsProcessor

39 from vllm.model_executor.layers.quantization import QuantizationConfig

40 from vllm.model_executor.layers.rotary_embedding import get_rope

41 from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding

42 -from vllm.model_executor.model_loader.weight_utils import (

43 - default_weight_loader,

44 - maybe_remap_kv_scale_name,

45 -)

42 from vllm.sequence import IntermediateTensors

43

44 from .interfaces import SupportsLoRA, SupportsPP

45 from .utils import (

46 AutoWeightsLoader,

47 + WeightsMapper,

48 extract_layer_index,

52 - is_pp_missing_parameter,

49 make_empty_intermediate_tensors_factory,

50 make_layers,

51 maybe_prefix,

@@ -316,60 +312,21 @@ class Gemma2Model(nn.Module):

312 hidden_states, _ = self.norm(hidden_states, residual)

313 return hidden_states

314

319 - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

320 - stacked_params_mapping = [

321 - # (param_name, shard_name, shard_id)

322 - ("qkv_proj", "q_proj", "q"),

323 - ("qkv_proj", "k_proj", "k"),

324 - ("qkv_proj", "v_proj", "v"),

325 - ("gate_up_proj", "gate_proj", 0),

326 - ("gate_up_proj", "up_proj", 1),

327 - ]

328 - params_dict = dict(self.named_parameters())

329 - loaded_params: set[str] = set()

330 - for name, loaded_weight in weights:

331 - for param_name, shard_name, shard_id in stacked_params_mapping:

332 - if shard_name not in name:

333 - continue

334 - name = name.replace(shard_name, param_name)

335 - # Skip loading extra bias for GPTQ models.

336 - if name.endswith(".bias") and name not in params_dict:

337 - continue

338 - if is_pp_missing_parameter(name, self):

339 - continue

340 - param = params_dict[name]

341 - weight_loader = param.weight_loader

342 - weight_loader(param, loaded_weight, shard_id)

343 - break

344 - else:

345 - # Skip loading extra bias for GPTQ models.

346 - if name.endswith(".bias") and name not in params_dict:

347 - continue

348 - # Remapping the name of FP8 kv-scale.

349 - name = maybe_remap_kv_scale_name(name, params_dict)

350 - if name is None:

351 - continue

352 - if is_pp_missing_parameter(name, self):

353 - continue

354 - param = params_dict[name]

355 - weight_loader = getattr(param, "weight_loader", default_weight_loader)

356 - weight_loader(param, loaded_weight)

357 - loaded_params.add(name)

358 -

359 - return loaded_params

360 -

315

316 class Gemma2ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

317 + hf_to_vllm_mapper = WeightsMapper(

318 + orig_to_new_stacked={

319 + # weight_name: (param_name, shard_id)

320 + ".q_proj": (".qkv_proj", "q"),

321 + ".k_proj": (".qkv_proj", "k"),

322 + ".v_proj": (".qkv_proj", "v"),

323 + ".gate_proj": (".gate_up_proj", 0),

324 + ".up_proj": (".gate_up_proj", 1),

325 + }

326 + )

327 packed_modules_mapping = {

364 - "qkv_proj": [

365 - "q_proj",

366 - "k_proj",

367 - "v_proj",

368 - ],

369 - "gate_up_proj": [

370 - "gate_proj",

371 - "up_proj",

372 - ],

328 + "qkv_proj": ["q_proj", "k_proj", "v_proj"],

329 + "gate_up_proj": ["gate_proj", "up_proj"],

330 }

331

332 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

@@ -418,4 +375,4 @@ class Gemma2ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

375 self,

376 skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),

377 )

421 - return loader.load_weights(weights)

378 + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

vllm/model_executor/models/gemma3.py

+12 -63

@@ -44,18 +44,14 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (

44 ParallelLMHead,

45 VocabParallelEmbedding,

46 )

47 -from vllm.model_executor.model_loader.weight_utils import (

48 - default_weight_loader,

49 - maybe_remap_kv_scale_name,

50 -)

47 from vllm.sequence import IntermediateTensors

48 from vllm.v1.attention.backend import AttentionType

49

50 from .interfaces import SupportsLoRA, SupportsPP

51 from .utils import (

52 AutoWeightsLoader,

53 + WeightsMapper,

54 extract_layer_index,

58 - is_pp_missing_parameter,

55 make_empty_intermediate_tensors_factory,

56 make_layers,

57 maybe_prefix,

@@ -365,65 +361,18 @@ class Gemma3Model(nn.Module):

361 hidden_states, _ = self.norm(hidden_states, residual)

362 return hidden_states

363

368 - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

369 - stacked_params_mapping = [

370 - # (param_name, shard_name, shard_id)

371 - ("qkv_proj", "q_proj", "q"),

372 - ("qkv_proj", "k_proj", "k"),

373 - ("qkv_proj", "v_proj", "v"),

374 - ("gate_up_proj", "gate_proj", 0),

375 - ("gate_up_proj", "up_proj", 1),

376 - ]

377 - params_dict = dict(self.named_parameters())

378 - loaded_params: set[str] = set()

379 - for name, loaded_weight in weights:

380 - # Check if this is a scale parameter that needs remapping first

381 - if name.endswith((".k_scale", ".v_scale", ".q_scale", ".prob_scale")):

382 - # Try to remap the scale name first

383 - remapped_name = maybe_remap_kv_scale_name(name, params_dict)

384 - if remapped_name is not None and remapped_name in params_dict:

385 - # Successfully remapped, use the remapped name

386 - param = params_dict[remapped_name]

387 - weight_loader = getattr(

388 - param, "weight_loader", default_weight_loader

389 - )

390 - weight_loader(param, loaded_weight)

391 - loaded_params.add(remapped_name)

392 - continue

393 - # If remapping failed, continue with normal processing

394 -

395 - for param_name, shard_name, shard_id in stacked_params_mapping:

396 - if shard_name not in name:

397 - continue

398 - name = name.replace(shard_name, param_name)

399 - # Skip loading extra bias for GPTQ models.

400 - if name.endswith(".bias") and name not in params_dict:

401 - continue

402 - if is_pp_missing_parameter(name, self):

403 - continue

404 - param = params_dict[name]

405 - weight_loader = param.weight_loader

406 - weight_loader(param, loaded_weight, shard_id)

407 - break

408 - else:

409 - # Skip loading extra bias for GPTQ models.

410 - if name.endswith(".bias") and name not in params_dict:

411 - continue

412 - # Remapping the name of FP8 kv-scale.

413 - name = maybe_remap_kv_scale_name(name, params_dict)

414 - if name is None:

415 - continue

416 - if is_pp_missing_parameter(name, self):

417 - continue

418 - param = params_dict[name]

419 - weight_loader = getattr(param, "weight_loader", default_weight_loader)

420 - weight_loader(param, loaded_weight)

421 - loaded_params.add(name)

422 -

423 - return loaded_params

424 -

364

365 class Gemma3ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

366 + hf_to_vllm_mapper = WeightsMapper(

367 + orig_to_new_stacked={

368 + # weight_name: (param_name, shard_id)

369 + ".q_proj": (".qkv_proj", "q"),

370 + ".k_proj": (".qkv_proj", "k"),

371 + ".v_proj": (".qkv_proj", "v"),

372 + ".gate_proj": (".gate_up_proj", 0),

373 + ".up_proj": (".gate_up_proj", 1),

374 + }

375 + )

376 packed_modules_mapping = {

377 "qkv_proj": [

378 "q_proj",

@@ -491,4 +440,4 @@ class Gemma3ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

440 self,

441 skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),

442 )

494 - return loader.load_weights(weights)

443 + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

vllm/model_executor/models/glm4.py

+11 -73

@@ -39,10 +39,6 @@ from vllm.model_executor.layers.logits_processor import LogitsProcessor

39 from vllm.model_executor.layers.quantization import QuantizationConfig

40 from vllm.model_executor.layers.rotary_embedding import get_rope

41 from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead

42 -from vllm.model_executor.model_loader.weight_utils import (

43 - default_weight_loader,

44 - maybe_remap_kv_scale_name,

45 -)

42 from vllm.sequence import IntermediateTensors

43 from vllm.v1.attention.backend import AttentionType

44

@@ -52,7 +48,6 @@ from .llama import LlamaModel

48 from .utils import (

49 AutoWeightsLoader,

50 PPMissingLayer,

55 - is_pp_missing_parameter,

51 maybe_prefix,

52 )

53

@@ -237,73 +232,11 @@ class Glm4Model(LlamaModel):

232 vllm_config=vllm_config, prefix=prefix, layer_type=Glm4DecoderLayer

233 )

234

240 - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

241 - stacked_params_mapping = [

242 - # (param_name, shard_name, shard_id)

243 - (".qkv_proj", ".q_proj", "q"),

244 - (".qkv_proj", ".k_proj", "k"),

245 - (".qkv_proj", ".v_proj", "v"),

246 - (".gate_up_proj", ".gate_proj", 0),

247 - (".gate_up_proj", ".up_proj", 1),

248 - ]

249 - params_dict = dict(self.named_parameters())

250 - loaded_params: set[str] = set()

251 - for name, loaded_weight in weights:

252 - spec_layer = get_spec_layer_idx_from_weight_name(self.config, name)

253 - if spec_layer is not None:

254 - continue

255 - if "rotary_emb.inv_freq" in name:

256 - continue

257 - if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name:

258 - # Models trained using ColossalAI may include these tensors in

259 - # the checkpoint. Skip them.

260 - continue

261 - if "scale" in name or "zero_point" in name:

262 - # Remapping the name of FP8 kv-scale or zero point.

263 - name = maybe_remap_kv_scale_name(name, params_dict)

264 - if name is None:

265 - continue

266 - for param_name, weight_name, shard_id in stacked_params_mapping:

267 - if weight_name not in name:

268 - continue

269 - name = name.replace(weight_name, param_name)

270 - # Skip loading extra bias for GPTQ models.

271 - if name.endswith(".bias") and name not in params_dict:

272 - continue

273 -

274 - if is_pp_missing_parameter(name, self):

275 - continue

276 -

277 - param = params_dict[name]

278 - weight_loader = param.weight_loader

279 - weight_loader(param, loaded_weight, shard_id)

280 - break

281 - else:

282 - # Skip loading extra bias for GPTQ models.

283 - if name.endswith(".bias") and name not in params_dict:

284 - continue

285 -

286 - if is_pp_missing_parameter(name, self):

287 - continue

288 -

289 - param = params_dict[name]

290 - weight_loader = getattr(param, "weight_loader", default_weight_loader)

291 - weight_loader(param, loaded_weight)

292 - loaded_params.add(name)

293 - return loaded_params

294 -

235

236 class Glm4ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

237 packed_modules_mapping = {

298 - "qkv_proj": [

299 - "q_proj",

300 - "k_proj",

301 - "v_proj",

302 - ],

303 - "gate_up_proj": [

304 - "gate_proj",

305 - "up_proj",

306 - ],

238 + "qkv_proj": ["q_proj", "k_proj", "v_proj"],

239 + "gate_up_proj": ["gate_proj", "up_proj"],

240 }

241

242 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

@@ -360,10 +293,15 @@ class Glm4ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

293 return logits

294

295 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

363 - loader = AutoWeightsLoader(

364 - self,

365 - skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),

366 - )

296 + skip_prefixes = ["lm_head."] if self.config.tie_word_embeddings else []

297 + # Skip the speculative (MTP) layers, which are loaded by the

298 + # draft model instead.

299 + num_nextn_layers = getattr(self.config, "num_nextn_predict_layers", 0)

300 + skip_prefixes += [

301 + f"model.layers.{self.config.num_hidden_layers + i}."

302 + for i in range(num_nextn_layers)

303 + ]

304 + loader = AutoWeightsLoader(self, skip_prefixes=skip_prefixes)

305 return loader.load_weights(weights)

306

307

vllm/model_executor/models/glm4v.py

+13

@@ -61,6 +61,7 @@ from .interfaces import (

61 SupportsMultiModal,

62 SupportsPP,

63 )

64 +from .utils import WeightsMapper

65

66

67 class GLMVImagePixelInputs(TensorSchema):

@@ -376,6 +377,15 @@ class EVA2CLIPModel(nn.Module):

377

378

379 class GLM4VModel(ChatGLMModel):

380 + hf_to_vllm_mapper = ChatGLMModel.hf_to_vllm_mapper | WeightsMapper(

381 + orig_to_new_stacked={

382 + # weight_name: (param_name, shard_id)

383 + # Vision GLU projections

384 + "linear_proj.gate_proj": ("linear_proj.merged_proj", 0),

385 + "linear_proj.dense_h_to_4h": ("linear_proj.merged_proj", 1),

386 + }

387 + )

388 +

389 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

390 super().__init__(vllm_config=vllm_config, prefix=prefix)

391

@@ -507,6 +517,9 @@ class GLM4VMultiModalProcessor(BaseMultiModalProcessor[GLM4VProcessingInfo]):

517 class GLM4VForCausalLM(

518 ChatGLMBaseModel, SupportsMultiModal, SupportsLoRA, SupportsPP, SupportsMRoPE

519 ):

520 + # NOTE: we must bring this to the surface because GLM4VModel.hf_to_vllm_mapper

521 + # contains non-stacking related mappings which LoRA/BnB needs to know about

522 + hf_to_vllm_mapper = GLM4VModel.hf_to_vllm_mapper

523 packed_modules_mapping = {

524 "query_key_value": ["query_key_value"],

525 "dense_h_to_4h": ["dense_h_to_4h"],

vllm/model_executor/models/gpt_j.py

+12 -50

@@ -43,16 +43,12 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (

43 ParallelLMHead,

44 VocabParallelEmbedding,

45 )

46 -from vllm.model_executor.model_loader.weight_utils import (

47 - default_weight_loader,

48 - maybe_remap_kv_scale_name,

49 -)

46 from vllm.sequence import IntermediateTensors

47

48 from .interfaces import SupportsPP

49 from .utils import (

50 AutoWeightsLoader,

55 - is_pp_missing_parameter,

51 + WeightsMapper,

52 make_empty_intermediate_tensors_factory,

53 make_layers,

54 maybe_prefix,

@@ -239,51 +235,17 @@ class GPTJModel(nn.Module):

235 hidden_states = self.ln_f(hidden_states)

236 return hidden_states

237

242 - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

243 - stacked_params_mapping = [

244 - # (param_name, shard_name, shard_id)

245 - ("qkv_proj", "q_proj", "q"),

246 - ("qkv_proj", "k_proj", "k"),

247 - ("qkv_proj", "v_proj", "v"),

248 - ("gate_up_proj", "gate_proj", 0),

249 - ("gate_up_proj", "up_proj", 1),

250 - ]

251 - params_dict = dict(self.named_parameters())

252 - loaded_params: set[str] = set()

253 - for name, loaded_weight in weights:

254 - if "attn.bias" in name or "attn.masked_bias" in name:

255 - continue

256 -

257 - for param_name, weight_name, shard_id in stacked_params_mapping:

258 - if weight_name not in name:

259 - continue

260 - name = name.replace(weight_name, param_name)

261 - # Skip loading extra bias for GPTQ models.

262 - if name.endswith(".bias") and name not in params_dict:

263 - continue

264 - if is_pp_missing_parameter(name, self):

265 - continue

266 - param = params_dict[name]

267 - weight_loader = param.weight_loader

268 - weight_loader(param, loaded_weight, shard_id)

269 - break

270 - else:

271 - name = maybe_remap_kv_scale_name(name, params_dict)

272 - if name is None:

273 - continue

274 - # Skip loading extra bias for GPTQ models.

275 - if name.endswith(".bias") and name not in params_dict:

276 - continue

277 - if is_pp_missing_parameter(name, self):

278 - continue

279 - param = params_dict[name]

280 - weight_loader = getattr(param, "weight_loader", default_weight_loader)

281 - weight_loader(param, loaded_weight)

282 - loaded_params.add(name)

283 - return loaded_params

284 -

238

239 class GPTJForCausalLM(nn.Module, SupportsPP):

240 + hf_to_vllm_mapper = WeightsMapper(

241 + orig_to_new_stacked={

242 + # weight_name: (param_name, shard_id)

243 + ".q_proj": (".qkv_proj", "q"),

244 + ".k_proj": (".qkv_proj", "k"),

245 + ".v_proj": (".qkv_proj", "v"),

246 + }

247 + )

248 +

249 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

250 super().__init__()

251 config = vllm_config.model_config.hf_config

@@ -329,5 +291,5 @@ class GPTJForCausalLM(nn.Module, SupportsPP):

291 return logits

292

293 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

332 - loader = AutoWeightsLoader(self)

333 - return loader.load_weights(weights)

294 + loader = AutoWeightsLoader(self, skip_substrs=["attn.bias", "attn.masked_bias"])

295 + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

vllm/model_executor/models/granite.py

+20 -62

@@ -49,17 +49,13 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (

49 ParallelLMHead,

50 VocabParallelEmbedding,

51 )

52 -from vllm.model_executor.model_loader.weight_utils import (

53 - default_weight_loader,

54 - maybe_remap_kv_scale_name,

55 -)

52 from vllm.sequence import IntermediateTensors

53

58 -from .interfaces import SupportsLoRA, SupportsPP

54 +from .interfaces import SupportsLoRA, SupportsPP, SupportsQuant

55 from .utils import (

56 AutoWeightsLoader,

57 PPMissingLayer,

62 - is_pp_missing_parameter,

58 + WeightsMapper,

59 make_layers,

60 maybe_prefix,

61 )

@@ -252,6 +248,17 @@ class GraniteDecoderLayer(nn.Module):

248

249 @support_torch_compile

250 class GraniteModel(nn.Module):

251 + hf_to_vllm_mapper = WeightsMapper(

252 + orig_to_new_stacked={

253 + # weight_name: (param_name, shard_id)

254 + ".q_proj": (".qkv_proj", "q"),

255 + ".k_proj": (".qkv_proj", "k"),

256 + ".v_proj": (".qkv_proj", "v"),

257 + ".gate_proj": (".gate_up_proj", 0),

258 + ".up_proj": (".gate_up_proj", 1),

259 + }

260 + )

261 +

262 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

263 super().__init__()

264

@@ -322,66 +329,17 @@ class GraniteModel(nn.Module):

329 return hidden_states

330

331 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

325 - stacked_params_mapping = [

326 - # (param_name, shard_name, shard_id)

327 - (".qkv_proj", ".q_proj", "q"),

328 - (".qkv_proj", ".k_proj", "k"),

329 - (".qkv_proj", ".v_proj", "v"),

330 - (".gate_up_proj", ".gate_proj", 0),

331 - (".gate_up_proj", ".up_proj", 1),

332 - ]

333 - params_dict = dict(self.named_parameters())

334 - loaded_params: set[str] = set()

335 - for name, loaded_weight in weights:

336 - for param_name, weight_name, shard_id in stacked_params_mapping:

337 - if weight_name not in name:

338 - continue

339 - name = name.replace(weight_name, param_name)

340 - # Skip loading extra bias for GPTQ models.

341 - if name.endswith(".bias") and name not in params_dict:

342 - continue

343 -

344 - if is_pp_missing_parameter(name, self):

345 - continue

346 -

347 - param = params_dict[name]

348 - weight_loader = param.weight_loader

349 - weight_loader(param, loaded_weight, shard_id)

350 -

351 - break

352 - else:

353 - # Skip loading extra bias for GPTQ models.

354 - if name.endswith(".bias") and name not in params_dict:

355 - continue

356 - # Remapping the name of FP8 kv-scale.

357 - name = maybe_remap_kv_scale_name(name, params_dict)

358 - if name is None:

359 - continue

360 -

361 - if is_pp_missing_parameter(name, self):

362 - continue

332 + loader = AutoWeightsLoader(self)

333 + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

334

364 - param = params_dict[name]

365 - weight_loader = getattr(param, "weight_loader", default_weight_loader)

366 - weight_loader(param, loaded_weight)

367 - loaded_params.add(name)

368 - return loaded_params

335

370 -

371 -class GraniteForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

336 +class GraniteForCausalLM(nn.Module, SupportsLoRA, SupportsPP, SupportsQuant):

337 + hf_to_vllm_mapper = GraniteModel.hf_to_vllm_mapper

338 + # LoRA specific attributes

339 packed_modules_mapping = {

373 - "qkv_proj": [

374 - "q_proj",

375 - "k_proj",

376 - "v_proj",

377 - ],

378 - "gate_up_proj": [

379 - "gate_proj",

380 - "up_proj",

381 - ],

340 + "qkv_proj": ["q_proj", "k_proj", "v_proj"],

341 + "gate_up_proj": ["gate_proj", "up_proj"],

342 }

383 -

384 - # LoRA specific attributes

343 embedding_modules = {

344 "embed_tokens": "input_embeddings",

345 "lm_head": "output_embeddings",

vllm/model_executor/models/hyperclovax.py

+14 -68

@@ -50,10 +50,6 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (

50 ParallelLMHead,

51 VocabParallelEmbedding,

52 )

53 -from vllm.model_executor.model_loader.weight_utils import (

54 - default_weight_loader,

55 - maybe_remap_kv_scale_name,

56 -)

53 from vllm.sequence import IntermediateTensors

54 from vllm.transformers_utils.configs.hyperclovax import HyperCLOVAXConfig

55

@@ -61,7 +57,7 @@ from .interfaces import SupportsLoRA, SupportsPP

57 from .utils import (

58 AutoWeightsLoader,

59 PPMissingLayer,

64 - is_pp_missing_parameter,

60 + WeightsMapper,

61 make_empty_intermediate_tensors_factory,

62 make_layers,

63 maybe_prefix,

@@ -377,71 +373,21 @@ class HyperCLOVAXModel(nn.Module):

373 hidden_states = self.norm(hidden_states)

374 return hidden_states

375

380 - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

381 - stacked_params_mapping = [

382 - # (param_name, shard_name, shard_id)

383 - (".qkv_proj", ".q_proj", "q"),

384 - (".qkv_proj", ".k_proj", "k"),

385 - (".qkv_proj", ".v_proj", "v"),

386 - (".gate_up_proj", ".gate_proj", 0),

387 - (".gate_up_proj", ".up_proj", 1),

388 - ]

389 - params_dict = dict(self.named_parameters())

390 - loaded_params: set[str] = set()

391 - for name, loaded_weight in weights:

392 - if "rotary_emb.inv_freq" in name:

393 - continue

394 - if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name:

395 - # Models trained using ColossalAI may include these tensors in

396 - # the checkpoint. Skip them.

397 - continue

398 - if "scale" in name or "zero_point" in name:

399 - # Remapping the name of FP8 kv-scale or zero point.

400 - remapped_name = maybe_remap_kv_scale_name(name, params_dict)

401 - if remapped_name is None:

402 - continue

403 - name = remapped_name

404 - for param_name, weight_name, shard_id in stacked_params_mapping:

405 - if weight_name not in name:

406 - continue

407 - name = name.replace(weight_name, param_name)

408 - # Skip loading extra bias for GPTQ models.

409 - if name.endswith(".bias") and name not in params_dict:

410 - continue

411 -

412 - if is_pp_missing_parameter(name, self):

413 - continue

414 -

415 - param = params_dict[name]

416 - weight_loader = param.weight_loader # type: ignore[attr-defined]

417 - weight_loader(param, loaded_weight, shard_id)

418 - break

419 - else:

420 - # Skip loading extra bias for GPTQ models.

421 - if name.endswith(".bias") and name not in params_dict:

422 - continue

423 -

424 - if is_pp_missing_parameter(name, self):

425 - continue

426 -

427 - param = params_dict[name]

428 - weight_loader = getattr(param, "weight_loader", default_weight_loader)

429 - weight_loader(param, loaded_weight)

430 - loaded_params.add(name)

431 - return loaded_params

432 -

376

377 class HyperCLOVAXForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

378 + hf_to_vllm_mapper = WeightsMapper(

379 + orig_to_new_stacked={

380 + # weight_name: (param_name, shard_id)

381 + ".q_proj": (".qkv_proj", "q"),

382 + ".k_proj": (".qkv_proj", "k"),

383 + ".v_proj": (".qkv_proj", "v"),

384 + ".gate_proj": (".gate_up_proj", 0),

385 + ".up_proj": (".gate_up_proj", 1),

386 + }

387 + )

388 packed_modules_mapping = {

436 - "qkv_proj": [

437 - "q_proj",

438 - "k_proj",

439 - "v_proj",

440 - ],

441 - "gate_up_proj": [

442 - "gate_proj",

443 - "up_proj",

444 - ],

389 + "qkv_proj": ["q_proj", "k_proj", "v_proj"],

390 + "gate_up_proj": ["gate_proj", "up_proj"],

391 }

392

393 # LoRA specific attributes

@@ -536,4 +482,4 @@ class HyperCLOVAXForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

482 self,

483 skip_prefixes=["lm_head."] if self.config.tie_word_embeddings else None,

484 )

539 - return loader.load_weights(weights)

485 + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

vllm/model_executor/models/interfaces.py

+2 -1

@@ -1033,7 +1033,8 @@ class SupportsQuant:

1033 if self.quant_config is None:

1034 return

1035 if (hf_to_vllm_mapper := self.hf_to_vllm_mapper) is not None:

1036 - self.quant_config.apply_vllm_mapper(hf_to_vllm_mapper)

1036 + unstacked_mapper = hf_to_vllm_mapper.get_unstacked_mapper()

1037 + self.quant_config.apply_vllm_mapper(unstacked_mapper)

1038 if self.packed_modules_mapping is not None:

1039 self.quant_config.packed_modules_mapping.update(self.packed_modules_mapping)

1040

vllm/model_executor/models/internlm2.py

+16 -40

@@ -35,15 +35,14 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (

35 ParallelLMHead,

36 VocabParallelEmbedding,

37 )

38 -from vllm.model_executor.model_loader.weight_utils import default_weight_loader

38 from vllm.sequence import IntermediateTensors

39

41 -from .interfaces import SupportsLoRA, SupportsPP

40 +from .interfaces import SupportsLoRA, SupportsPP, SupportsQuant

41 from .interfaces_base import default_pooling_type

42 from .utils import (

43 AutoWeightsLoader,

44 StageMissingLayer,

46 - is_pp_missing_parameter,

45 + WeightsMapper,

46 make_empty_intermediate_tensors_factory,

47 make_layers,

48 maybe_prefix,

@@ -249,6 +248,14 @@ class InternLMDecoderLayer(nn.Module):

248

249 @support_torch_compile

250 class InternLM2Model(nn.Module):

251 + hf_to_vllm_mapper = WeightsMapper(

252 + orig_to_new_stacked={

253 + # weight_name: (param_name, shard_id)

254 + ".w1": (".gate_up_proj", 0),

255 + ".w3": (".gate_up_proj", 1),

256 + }

257 + )

258 +

259 def __init__(

260 self,

261 *,

@@ -310,43 +317,12 @@ class InternLM2Model(nn.Module):

317 return hidden_states

318

319 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

313 - stacked_params_mapping = [

314 - # (param_name, shard_name, shard_id)

315 - ("gate_up_proj", "w1", 0),

316 - ("gate_up_proj", "w3", 1),

317 - ]

318 - params_dict = dict(self.named_parameters())

319 - loaded_params: set[str] = set()

320 - for name, loaded_weight in weights:

321 - if "rotary_emb.inv_freq" in name:

322 - continue

323 - for param_name, weight_name, shard_id in stacked_params_mapping:

324 - if weight_name not in name:

325 - continue

326 - name = name.replace(weight_name, param_name)

327 - # Skip loading extra bias for GPTQ models.

328 - if name.endswith(".bias") and name not in params_dict:

329 - continue

330 - if is_pp_missing_parameter(name, self):

331 - continue

332 - param = params_dict[name]

333 - weight_loader = param.weight_loader

334 - weight_loader(param, loaded_weight, shard_id)

335 - break

336 - else:

337 - # Skip loading extra bias for GPTQ models.

338 - if name.endswith(".bias") and name not in params_dict:

339 - continue

340 - if is_pp_missing_parameter(name, self):

341 - continue

342 - param = params_dict[name]

343 - weight_loader = getattr(param, "weight_loader", default_weight_loader)

344 - weight_loader(param, loaded_weight)

345 - loaded_params.add(name)

346 - return loaded_params

347 -

348 -

349 -class InternLM2ForCausalLM(nn.Module, SupportsPP, SupportsLoRA):

320 + loader = AutoWeightsLoader(self)

321 + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

322 +

323 +

324 +class InternLM2ForCausalLM(nn.Module, SupportsPP, SupportsLoRA, SupportsQuant):

325 + hf_to_vllm_mapper = InternLM2Model.hf_to_vllm_mapper

326 packed_modules_mapping = {

327 "wqkv": ["wqkv"],

328 "gate_up_proj": ["w1", "w3"],

vllm/model_executor/models/jais2.py

+10 -59

@@ -51,18 +51,14 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (

51 ParallelLMHead,

52 VocabParallelEmbedding,

53 )

54 -from vllm.model_executor.model_loader.weight_utils import (

55 - default_weight_loader,

56 - maybe_remap_kv_scale_name,

57 -)

54 from vllm.sequence import IntermediateTensors

55

56 from .interfaces import SupportsLoRA, SupportsPP

57 from .utils import (

58 AutoWeightsLoader,

59 PPMissingLayer,

60 + WeightsMapper,

61 extract_layer_index,

65 - is_pp_missing_parameter,

62 make_empty_intermediate_tensors_factory,

63 make_layers,

64 maybe_prefix,

@@ -366,61 +362,16 @@ class Jais2Model(nn.Module):

362 hidden_states, _ = self.norm(hidden_states + residual), residual

363 return hidden_states

364

369 - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

370 - stacked_params_mapping = [

371 - # (param_name, shard_name, shard_id)

372 - (".qkv_proj", ".q_proj", "q"),

373 - (".qkv_proj", ".k_proj", "k"),

374 - (".qkv_proj", ".v_proj", "v"),

375 - ]

376 - params_dict = dict(self.named_parameters())

377 - loaded_params: set[str] = set()

378 - for name, loaded_weight in weights:

379 - if "rotary_emb.inv_freq" in name:

380 - continue

381 - if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name:

382 - # Models trained using ColossalAI may include these tensors in

383 - # the checkpoint. Skip them.

384 - continue

385 - if "scale" in name:

386 - name = maybe_remap_kv_scale_name(name, params_dict)

387 - if name is None:

388 - continue

389 - for param_name, weight_name, shard_id in stacked_params_mapping:

390 - if weight_name not in name:

391 - continue

392 - name = name.replace(weight_name, param_name)

393 - # Skip loading extra bias for GPTQ models.

394 - if name.endswith(".bias") and name not in params_dict:

395 - continue

396 -

397 - if is_pp_missing_parameter(name, self):

398 - continue

399 -

400 - param = params_dict[name]

401 - weight_loader = param.weight_loader

402 - weight_loader(param, loaded_weight, shard_id)

403 - break

404 - else:

405 - # Skip loading extra bias for GPTQ models.

406 - if name.endswith(".bias") and name not in params_dict:

407 - continue

408 - # Remapping the name of FP8 kv-scale.

409 - name = maybe_remap_kv_scale_name(name, params_dict)

410 - if name is None:

411 - continue

412 -

413 - if is_pp_missing_parameter(name, self):

414 - continue

415 -

416 - param = params_dict[name]

417 - weight_loader = getattr(param, "weight_loader", default_weight_loader)

418 - weight_loader(param, loaded_weight)

419 - loaded_params.add(name)

420 - return loaded_params

421 -

365

366 class Jais2ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

367 + hf_to_vllm_mapper = WeightsMapper(

368 + orig_to_new_stacked={

369 + # weight_name: (param_name, shard_id)

370 + ".q_proj": (".qkv_proj", "q"),

371 + ".k_proj": (".qkv_proj", "k"),

372 + ".v_proj": (".qkv_proj", "v"),

373 + }

374 + )

375 packed_modules_mapping = {

376 "qkv_proj": ["q_proj", "k_proj", "v_proj"],

377 }

@@ -490,4 +441,4 @@ class Jais2ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

441 self,

442 skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),

443 )

493 - return loader.load_weights(weights)

444 + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

vllm/model_executor/models/jina.py

+3 -2

@@ -254,5 +254,6 @@ class JinaEmbeddingsV5Model(Qwen3ForCausalLM, VllmModelForPooling):

254 tensor = tensor + (lora_B @ lora_A) * scaling

255 yield name, tensor

256

257 - loaded = self.model.load_weights(_merge_weights(weights))

258 - return {f"model.{name}" for name in loaded}

257 + loader = AutoWeightsLoader(self.model, ignore_unexpected_prefixes=["lm_head."])

258 + weights = _merge_weights(weights)

259 + return loader.load_weights(weights, mapper=self.model.hf_to_vllm_mapper)

vllm/model_executor/models/llama.py

+24 -58

@@ -52,10 +52,6 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (

52 ParallelLMHead,

53 VocabParallelEmbedding,

54 )

55 -from vllm.model_executor.model_loader.weight_utils import (

56 - default_weight_loader,

57 - maybe_remap_kv_scale_name,

58 -)

55 from vllm.sequence import IntermediateTensors

56 from vllm.v1.attention.backend import AttentionType

57

@@ -67,12 +63,13 @@ from .interfaces import (

63 SupportsEagle3,

64 SupportsLoRA,

65 SupportsPP,

66 + SupportsQuant,

67 )

68 from .utils import (

69 AutoWeightsLoader,

70 PPMissingLayer,

71 + WeightsMapper,

72 extract_layer_index,

75 - is_pp_missing_parameter,

73 make_empty_intermediate_tensors_factory,

74 make_layers,

75 maybe_prefix,

@@ -345,6 +342,17 @@ class LlamaDecoderLayer(nn.Module):

342 },

343 )

344 class LlamaModel(nn.Module, EagleModelMixin):

345 + hf_to_vllm_mapper = WeightsMapper(

346 + orig_to_new_stacked={

347 + # weight_name: (param_name, shard_id)

348 + ".q_proj": (".qkv_proj", "q"),

349 + ".k_proj": (".qkv_proj", "k"),

350 + ".v_proj": (".qkv_proj", "v"),

351 + ".gate_proj": (".gate_up_proj", 0),

352 + ".up_proj": (".gate_up_proj", 1),

353 + }

354 + )

355 +

356 def __init__(

357 self,

358 *,

@@ -431,67 +439,25 @@ class LlamaModel(nn.Module, EagleModelMixin):

439 return hidden_states

440

441 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

434 - stacked_params_mapping = [

435 - # (param_name, shard_name, shard_id)

436 - (".qkv_proj", ".q_proj", "q"),

437 - (".qkv_proj", ".k_proj", "k"),

438 - (".qkv_proj", ".v_proj", "v"),

439 - (".gate_up_proj", ".gate_proj", 0),

440 - (".gate_up_proj", ".up_proj", 1),

441 - ]

442 - params_dict = dict(self.named_parameters())

443 - loaded_params: set[str] = set()

444 - for name, loaded_weight in weights:

445 - if "rotary_emb.inv_freq" in name:

446 - continue

447 - if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name:

448 - # Models trained using ColossalAI may include these tensors in

449 - # the checkpoint. Skip them.

450 - continue

451 - if "scale" in name or "zero_point" in name:

452 - # Remapping the name of FP8 kv-scale or zero point.

453 - name = maybe_remap_kv_scale_name(name, params_dict)

454 - if name is None:

455 - continue

456 - for param_name, weight_name, shard_id in stacked_params_mapping:

457 - if weight_name not in name:

458 - continue

459 - name = name.replace(weight_name, param_name)

460 - # Skip loading extra bias for GPTQ models.

461 - if name.endswith(".bias") and name not in params_dict:

462 - continue

463 -

464 - if is_pp_missing_parameter(name, self):

465 - continue

466 -

467 - param = params_dict[name]

468 - weight_loader = param.weight_loader

469 - weight_loader(param, loaded_weight, shard_id)

470 - break

471 - else:

472 - # Skip loading extra bias for GPTQ models.

473 - if name.endswith(".bias") and name not in params_dict:

474 - continue

475 -

476 - if is_pp_missing_parameter(name, self):

477 - continue

478 -

479 - param = params_dict[name]

480 - weight_loader = getattr(param, "weight_loader", default_weight_loader)

481 - weight_loader(param, loaded_weight)

482 - loaded_params.add(name)

483 - return loaded_params

442 + loader = AutoWeightsLoader(self)

443 + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

444

445

446 class LlamaForCausalLM(

487 - LocalArgmaxMixin, nn.Module, SupportsLoRA, SupportsPP, SupportsEagle, SupportsEagle3

447 + LocalArgmaxMixin,

448 + nn.Module,

449 + SupportsLoRA,

450 + SupportsPP,

451 + SupportsEagle,

452 + SupportsEagle3,

453 + SupportsQuant,

454 ):

455 + hf_to_vllm_mapper = LlamaModel.hf_to_vllm_mapper

456 + # LoRA specific attributes

457 packed_modules_mapping = {

458 "qkv_proj": ["q_proj", "k_proj", "v_proj"],

459 "gate_up_proj": ["gate_proj", "up_proj"],

460 }

493 -

494 - # LoRA specific attributes

461 embedding_modules = {

462 "embed_tokens": "input_embeddings",

463 "lm_head": "output_embeddings",

vllm/model_executor/models/mamba.py

+4 -21

@@ -26,7 +26,6 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (

26 ParallelLMHead,

27 VocabParallelEmbedding,

28 )

29 -from vllm.model_executor.model_loader.weight_utils import default_weight_loader

29 from vllm.model_executor.models.interfaces import (

30 HasInnerState,

31 IsAttentionFree,

@@ -37,7 +36,7 @@ from vllm.sequence import IntermediateTensors

36

37 from .utils import (

38 AutoWeightsLoader,

40 - is_pp_missing_parameter,

39 + WeightsMapper,

40 make_empty_intermediate_tensors_factory,

41 make_layers,

42 maybe_prefix,

@@ -170,28 +169,12 @@ class MambaModel(nn.Module):

169

170 return hidden_states

171

173 - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

174 - params_dict = dict(self.named_parameters())

175 - loaded_params: set[str] = set()

176 - for name, loaded_weight in weights:

177 - if "A_log" in name:

178 - name = name.replace("A_log", "A")

179 - # Skip loading extra bias for GPTQ models.

180 - if name.endswith(".bias") and name not in params_dict:

181 - continue

182 - if is_pp_missing_parameter(name, self):

183 - continue

184 -

185 - param = params_dict[name]

186 - weight_loader = getattr(param, "weight_loader", default_weight_loader)

187 - weight_loader(param, loaded_weight)

188 - loaded_params.add(name)

189 - return loaded_params

190 -

172

173 class MambaForCausalLM(

174 nn.Module, HasInnerState, IsAttentionFree, SupportsPP, SupportsMambaPrefixCaching

175 ):

176 + hf_to_vllm_mapper = WeightsMapper(orig_to_new_substr={".A_log": ".A"})

177 +

178 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

179 config = vllm_config.model_config.hf_config

180

@@ -279,4 +262,4 @@ class MambaForCausalLM(

262

263 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

264 loader = AutoWeightsLoader(self)

282 - return loader.load_weights(weights)

265 + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

vllm/model_executor/models/mamba2.py

+4 -22

@@ -25,7 +25,6 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (

25 ParallelLMHead,

26 VocabParallelEmbedding,

27 )

28 -from vllm.model_executor.model_loader.weight_utils import default_weight_loader

28 from vllm.model_executor.models.interfaces import (

29 HasInnerState,

30 IsAttentionFree,

@@ -35,7 +34,7 @@ from vllm.sequence import IntermediateTensors

34

35 from .utils import (

36 AutoWeightsLoader,

38 - is_pp_missing_parameter,

37 + WeightsMapper,

38 make_empty_intermediate_tensors_factory,

39 make_layers,

40 maybe_prefix,

@@ -167,29 +166,12 @@ class Mamba2Model(nn.Module):

166

167 return hidden_states

168

170 - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

171 - params_dict = dict(self.named_parameters())

172 - loaded_params: set[str] = set()

173 - for name, loaded_weight in weights:

174 - if "A_log" in name:

175 - name = name.replace("A_log", "A")

176 -

177 - # Skip loading extra bias for GPTQ models.

178 - if name.endswith(".bias") and name not in params_dict:

179 - continue

180 - if is_pp_missing_parameter(name, self):

181 - continue

182 -

183 - param = params_dict[name]

184 - weight_loader = getattr(param, "weight_loader", default_weight_loader)

185 - weight_loader(param, loaded_weight)

186 - loaded_params.add(name)

187 - return loaded_params

188 -

169

170 class Mamba2ForCausalLM(

171 nn.Module, HasInnerState, IsAttentionFree, SupportsMambaPrefixCaching

172 ):

173 + hf_to_vllm_mapper = WeightsMapper(orig_to_new_substr={".A_log": ".A"})

174 +

175 @classmethod

176 def get_mamba_state_dtype_from_config(

177 cls,

@@ -292,4 +274,4 @@ class Mamba2ForCausalLM(

274

275 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

276 loader = AutoWeightsLoader(self)

295 - return loader.load_weights(weights)

277 + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

vllm/model_executor/models/mimo.py

+8 -49

@@ -38,14 +38,10 @@ from vllm.distributed import get_pp_group

38 from vllm.logger import init_logger

39 from vllm.model_executor.layers.logits_processor import LogitsProcessor

40 from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead

41 -from vllm.model_executor.model_loader.weight_utils import (

42 - default_weight_loader,

43 - maybe_remap_kv_scale_name,

44 -)

41 from vllm.model_executor.models.qwen2 import Qwen2ForCausalLM, Qwen2Model

42 from vllm.sequence import IntermediateTensors

43

48 -from .utils import PPMissingLayer, is_pp_missing_parameter, maybe_prefix

44 +from .utils import AutoWeightsLoader, PPMissingLayer, maybe_prefix

45

46 logger = init_logger(__name__)

47

@@ -89,50 +85,6 @@ class MiMoModel(Qwen2Model):

85 hidden_states = hidden_states + residual

86 return hidden_states

87

92 - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

93 - stacked_params_mapping = [

94 - ("qkv_proj", "q_proj", "q"),

95 - ("qkv_proj", "k_proj", "k"),

96 - ("qkv_proj", "v_proj", "v"),

97 - ("gate_up_proj", "gate_proj", 0),

98 - ("gate_up_proj", "up_proj", 1),

99 - ]

100 - params_dict = dict(self.named_parameters(remove_duplicate=False))

101 - loaded_params: set[str] = set()

102 - for name, loaded_weight in weights:

103 - if "mtp_layers" in name:

104 - continue

105 - if "rotary_emb.inv_freq" in name:

106 - continue

107 - for param_name, weight_name, shard_id in stacked_params_mapping:

108 - if weight_name not in name:

109 - continue

110 - name = name.replace(weight_name, param_name)

111 - # Skip loading extra bias for GPTQ models.

112 - if name.endswith(".bias") and name not in params_dict:

113 - continue

114 - if is_pp_missing_parameter(name, self):

115 - continue

116 - param = params_dict[name]

117 - weight_loader = param.weight_loader

118 - weight_loader(param, loaded_weight, shard_id)

119 - break

120 - else:

121 - # Skip loading extra bias for GPTQ models.

122 - if name.endswith(".bias") and name not in params_dict:

123 - continue

124 - # Remapping the name of FP8 kv-scale.

125 - name = maybe_remap_kv_scale_name(name, params_dict)

126 - if name is None:

127 - continue

128 - if is_pp_missing_parameter(name, self):

129 - continue

130 - param = params_dict[name]

131 - weight_loader = getattr(param, "weight_loader", default_weight_loader)

132 - weight_loader(param, loaded_weight)

133 - loaded_params.add(name)

134 - return loaded_params

135 -

88

89 class MiMoForCausalLM(Qwen2ForCausalLM, nn.Module):

90 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

@@ -167,6 +119,13 @@ class MiMoForCausalLM(Qwen2ForCausalLM, nn.Module):

119 self.model.make_empty_intermediate_tensors

120 )

121

122 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

123 + skip_prefixes = ["lm_head."] if self.config.tie_word_embeddings else []

124 + # MTP layers are loaded by the draft model, not the main model.

125 + skip_prefixes.append("model.mtp_layers.")

126 + loader = AutoWeightsLoader(self, skip_prefixes=skip_prefixes)

127 + return loader.load_weights(weights)

128 +

129 def compute_logits(

130 self,

131 hidden_states: torch.Tensor,

vllm/model_executor/models/mistral_eagle.py

+5 -5

@@ -108,11 +108,6 @@ class EagleMistralModel(MistralModel):

108 hidden_states, _ = self.norm(hidden_states, residual)

109 return hidden_states, hidden_states

110

111 - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

112 - # Pretend embed_tokens is loaded; the actual weight is shared

113 - # from the target model at runtime by `load_eagle_model`.

114 - return super().load_weights(weights) | {"embed_tokens.weight"}

115 -

111

112 class EagleMistralForCausalLM(MistralForCausalLM):

113 mistral_mapping = MistralForCausalLM.mistral_mapping | {

@@ -166,3 +161,8 @@ class EagleMistralForCausalLM(MistralForCausalLM):

161 multimodal_embeddings=multimodal_embeddings,

162 is_multimodal=is_multimodal,

163 )

164 +

165 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

166 + # Pretend embed_tokens is loaded; the actual weight is shared

167 + # from the target model at runtime by `load_eagle_model`.

168 + return super().load_weights(weights) | {"model.embed_tokens.weight"}

vllm/model_executor/models/mpt.py

-17

@@ -27,13 +27,11 @@ from vllm.model_executor.layers.linear import (

27 from vllm.model_executor.layers.logits_processor import LogitsProcessor

28 from vllm.model_executor.layers.quantization import QuantizationConfig

29 from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding

30 -from vllm.model_executor.model_loader.weight_utils import default_weight_loader

30 from vllm.sequence import IntermediateTensors

31

32 from .interfaces import SupportsPP

33 from .utils import (

34 AutoWeightsLoader,

36 - is_pp_missing_parameter,

35 make_empty_intermediate_tensors_factory,

36 make_layers,

37 maybe_prefix,

@@ -274,21 +272,6 @@ class MPTModel(nn.Module):

272 hidden_states = self.norm_f(hidden_states)

273 return hidden_states

274

277 - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

278 - params_dict = dict(self.named_parameters(remove_duplicate=False))

279 - loaded_params: set[str] = set()

280 - for name, loaded_weight in weights:

281 - # Skip loading extra bias for GPTQ models.

282 - if name.endswith(".bias") and name not in params_dict:

283 - continue

284 - if is_pp_missing_parameter(name, self):

285 - continue

286 - param = params_dict[name]

287 - weight_loader = getattr(param, "weight_loader", default_weight_loader)

288 - weight_loader(param, loaded_weight)

289 - loaded_params.add(name)

290 - return loaded_params

291 -

275

276 class MPTForCausalLM(nn.Module, SupportsPP):

277 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

vllm/model_executor/models/nemotron.py

+11 -55

@@ -47,10 +47,6 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (

47 ParallelLMHead,

48 VocabParallelEmbedding,

49 )

50 -from vllm.model_executor.model_loader.weight_utils import (

51 - default_weight_loader,

52 - maybe_remap_kv_scale_name,

53 -)

50 from vllm.sequence import IntermediateTensors

51 from vllm.transformers_utils.configs.nemotron import NemotronConfig

52

@@ -58,7 +54,7 @@ from .interfaces import SupportsLoRA, SupportsPP

54 from .utils import (

55 AutoWeightsLoader,

56 PPMissingLayer,

61 - is_pp_missing_parameter,

57 + WeightsMapper,

58 make_empty_intermediate_tensors_factory,

59 make_layers,

60 maybe_prefix,

@@ -366,58 +362,18 @@ class NemotronModel(nn.Module):

362 hidden_states, _ = self.norm(hidden_states, residual)

363 return hidden_states

364

369 - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

370 - stacked_params_mapping = [

371 - # (param_name, shard_name, shard_id)

372 - (".qkv_proj", ".q_proj", "q"),

373 - (".qkv_proj", ".k_proj", "k"),

374 - (".qkv_proj", ".v_proj", "v"),

375 - ]

376 - params_dict = dict(self.named_parameters())

377 - loaded_params: set[str] = set()

378 - for name, loaded_weight in weights:

379 - for param_name, weight_name, shard_id in stacked_params_mapping:

380 - if weight_name not in name:

381 - continue

382 - name = name.replace(weight_name, param_name)

383 - # Skip loading extra bias for GPTQ models.

384 - if name.endswith(".bias") and name not in params_dict:

385 - continue

386 -

387 - if is_pp_missing_parameter(name, self):

388 - continue

389 -

390 - param = params_dict[name]

391 - weight_loader = param.weight_loader

392 - weight_loader(param, loaded_weight, shard_id)

393 -

394 - break

395 - else:

396 - # Skip loading extra bias for GPTQ models.

397 - if name.endswith(".bias") and name not in params_dict:

398 - continue

399 - # Remapping the name of FP8 kv-scale.

400 - name = maybe_remap_kv_scale_name(name, params_dict)

401 - if name is None:

402 - continue

403 -

404 - if is_pp_missing_parameter(name, self):

405 - continue

406 -

407 - param = params_dict[name]

408 - weight_loader = getattr(param, "weight_loader", default_weight_loader)

409 - weight_loader(param, loaded_weight)

410 - loaded_params.add(name)

411 - return loaded_params

412 -

365

366 class NemotronForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

367 + hf_to_vllm_mapper = WeightsMapper(

368 + orig_to_new_stacked={

369 + # weight_name: (param_name, shard_id)

370 + ".q_proj": (".qkv_proj", "q"),

371 + ".k_proj": (".qkv_proj", "k"),

372 + ".v_proj": (".qkv_proj", "v"),

373 + }

374 + )

375 packed_modules_mapping = {

416 - "qkv_proj": [

417 - "q_proj",

418 - "k_proj",

419 - "v_proj",

420 - ],

376 + "qkv_proj": ["q_proj", "k_proj", "v_proj"],

377 }

378

379 # LoRA specific attributes

@@ -485,4 +441,4 @@ class NemotronForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

441

442 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

443 loader = AutoWeightsLoader(self)

488 - return loader.load_weights(weights)

444 + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

vllm/model_executor/models/nemotron_nas.py

+12 -58

@@ -42,10 +42,6 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (

42 ParallelLMHead,

43 VocabParallelEmbedding,

44 )

45 -from vllm.model_executor.model_loader.weight_utils import (

46 - default_weight_loader,

47 - maybe_remap_kv_scale_name,

48 -)

45 from vllm.model_executor.models.llama import LlamaAttention, LlamaMLP

46 from vllm.sequence import IntermediateTensors

47 from vllm.v1.attention.backend import AttentionType

@@ -54,7 +50,7 @@ from .interfaces import HasNoOps, SupportsLoRA, SupportsPP

50 from .utils import (

51 AutoWeightsLoader,

52 PPMissingLayer,

57 - is_pp_missing_parameter,

53 + WeightsMapper,

54 make_empty_intermediate_tensors_factory,

55 make_layers,

56 maybe_prefix,

@@ -315,60 +311,18 @@ class DeciModel(nn.Module):

311 hidden_states, _ = self.norm(hidden_states, residual)

312 return hidden_states

313

318 - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

319 - stacked_params_mapping = [

320 - # (param_name, shard_name, shard_id)

321 - (".qkv_proj", ".q_proj", "q"),

322 - (".qkv_proj", ".k_proj", "k"),

323 - (".qkv_proj", ".v_proj", "v"),

324 - (".gate_up_proj", ".gate_proj", 0),

325 - (".gate_up_proj", ".up_proj", 1),

326 - ]

327 - params_dict = dict(self.named_parameters())

328 - loaded_params: set[str] = set()

329 - for name, loaded_weight in weights:

330 - if "rotary_emb.inv_freq" in name:

331 - continue

332 - if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name:

333 - # Models trained using ColossalAI may include these tensors in

334 - # the checkpoint. Skip them.

335 - continue

336 - if "scale" in name or "zero_point" in name:

337 - # Remapping the name of FP8 kv-scale.

338 - name = maybe_remap_kv_scale_name(name, params_dict)

339 - if name is None:

340 - continue

341 - for param_name, weight_name, shard_id in stacked_params_mapping:

342 - if weight_name not in name:

343 - continue

344 - name = name.replace(weight_name, param_name)

345 - # Skip loading extra bias for GPTQ models.

346 - if name.endswith(".bias") and name not in params_dict:

347 - continue

348 -

349 - if is_pp_missing_parameter(name, self):

350 - continue

351 -

352 - param = params_dict[name]

353 - weight_loader = param.weight_loader

354 - weight_loader(param, loaded_weight, shard_id)

355 - break

356 - else:

357 - # Skip loading extra bias for GPTQ models.

358 - if name.endswith(".bias") and name not in params_dict:

359 - continue

360 -

361 - if is_pp_missing_parameter(name, self):

362 - continue

363 -

364 - param = params_dict[name]

365 - weight_loader = getattr(param, "weight_loader", default_weight_loader)

366 - weight_loader(param, loaded_weight)

367 - loaded_params.add(name)

368 - return loaded_params

369 -

314

315 class DeciLMForCausalLM(nn.Module, SupportsLoRA, SupportsPP, HasNoOps):

316 + hf_to_vllm_mapper = WeightsMapper(

317 + orig_to_new_stacked={

318 + # weight_name: (param_name, shard_id)

319 + ".q_proj": (".qkv_proj", "q"),

320 + ".k_proj": (".qkv_proj", "k"),

321 + ".v_proj": (".qkv_proj", "v"),

322 + ".gate_proj": (".gate_up_proj", 0),

323 + ".up_proj": (".gate_up_proj", 1),

324 + }

325 + )

326 packed_modules_mapping = {

327 "qkv_proj": ["q_proj", "k_proj", "v_proj"],

328 "gate_up_proj": ["gate_proj", "up_proj"],

@@ -462,4 +416,4 @@ class DeciLMForCausalLM(nn.Module, SupportsLoRA, SupportsPP, HasNoOps):

416 self,

417 skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),

418 )

465 - return loader.load_weights(weights)

419 + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

vllm/model_executor/models/olmo.py

+14 -49

@@ -48,13 +48,12 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (

48 ParallelLMHead,

49 VocabParallelEmbedding,

50 )

51 -from vllm.model_executor.model_loader.weight_utils import default_weight_loader

51 from vllm.sequence import IntermediateTensors

52

53 from .interfaces import SupportsLoRA, SupportsPP

54 from .utils import (

55 AutoWeightsLoader,

57 - is_pp_missing_parameter,

56 + WeightsMapper,

57 make_empty_intermediate_tensors_factory,

58 make_layers,

59 maybe_prefix,

@@ -301,59 +300,25 @@ class OlmoModel(nn.Module):

300 hidden_states = self.norm(hidden_states)

301 return hidden_states

302

304 - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

305 - stacked_params_mapping = [

306 - # (param_name, shard_name, shard_id)

307 - ("qkv_proj", "q_proj", "q"),

308 - ("qkv_proj", "k_proj", "k"),

309 - ("qkv_proj", "v_proj", "v"),

310 - ("gate_up_proj", "gate_proj", 0),

311 - ("gate_up_proj", "up_proj", 1),

312 - ]

313 - params_dict = dict(self.named_parameters(remove_duplicate=False))

314 - loaded_params: set[str] = set()

315 - for name, loaded_weight in weights:

316 - for param_name, weight_name, shard_id in stacked_params_mapping:

317 - if weight_name not in name:

318 - continue

319 - name = name.replace(weight_name, param_name)

320 - # Skip loading extra bias for GPTQ models.

321 - if name.endswith(".bias") and name not in params_dict:

322 - continue

323 - if is_pp_missing_parameter(name, self):

324 - continue

325 - param = params_dict[name]

326 - weight_loader = param.weight_loader

327 - weight_loader(param, loaded_weight, shard_id)

328 - break

329 - else:

330 - # Skip loading extra bias for GPTQ models.

331 - if name.endswith(".bias") and name not in params_dict:

332 - continue

333 - if is_pp_missing_parameter(name, self):

334 - continue

335 - param = params_dict[name]

336 - weight_loader = getattr(param, "weight_loader", default_weight_loader)

337 - weight_loader(param, loaded_weight)

338 - loaded_params.add(name)

339 - return loaded_params

340 -

303

304 class OlmoForCausalLM(nn.Module, SupportsPP, SupportsLoRA):

305 """

306 Extremely barebones HF model wrapper.

307 """

308

309 + hf_to_vllm_mapper = WeightsMapper(

310 + orig_to_new_stacked={

311 + # weight_name: (param_name, shard_id)

312 + ".q_proj": (".qkv_proj", "q"),

313 + ".k_proj": (".qkv_proj", "k"),

314 + ".v_proj": (".qkv_proj", "v"),

315 + ".gate_proj": (".gate_up_proj", 0),

316 + ".up_proj": (".gate_up_proj", 1),

317 + }

318 + )

319 packed_modules_mapping = {

348 - "qkv_proj": [

349 - "q_proj",

350 - "k_proj",

351 - "v_proj",

352 - ],

353 - "gate_up_proj": [

354 - "gate_proj",

355 - "up_proj",

356 - ],

320 + "qkv_proj": ["q_proj", "k_proj", "v_proj"],

321 + "gate_up_proj": ["gate_proj", "up_proj"],

322 }

323

324 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

@@ -410,4 +375,4 @@ class OlmoForCausalLM(nn.Module, SupportsPP, SupportsLoRA):

375 ["lm_head.weight"] if self.config.tie_word_embeddings else None

376 ),

377 )

413 - return loader.load_weights(weights)

378 + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

vllm/model_executor/models/olmo2.py

+14 -48

@@ -52,12 +52,11 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (

52 ParallelLMHead,

53 VocabParallelEmbedding,

54 )

55 -from vllm.model_executor.model_loader.weight_utils import default_weight_loader

55 from vllm.model_executor.models.interfaces import SupportsLoRA, SupportsPP

56 from vllm.model_executor.models.utils import (

57 AutoWeightsLoader,

58 + WeightsMapper,

59 extract_layer_index,

60 - is_pp_missing_parameter,

60 make_empty_intermediate_tensors_factory,

61 make_layers,

62 maybe_prefix,

@@ -343,58 +342,25 @@ class Olmo2Model(nn.Module):

342 hidden_states = self.norm(hidden_states)

343 return hidden_states

344

346 - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

347 - stacked_params_mapping = [

348 - # (param_name, shard_name, shard_id)

349 - ("qkv_proj", "q_proj", "q"),

350 - ("qkv_proj", "k_proj", "k"),

351 - ("qkv_proj", "v_proj", "v"),

352 - ("gate_up_proj", "gate_proj", 0),

353 - ("gate_up_proj", "up_proj", 1),

354 - ]

355 -

356 - params_dict = dict(self.named_parameters(remove_duplicate=False))

357 - loaded_params: set[str] = set()

358 - for name, loaded_weight in weights:

359 - if is_pp_missing_parameter(name, self):

360 - continue

361 - for param_name, weight_name, shard_id in stacked_params_mapping:

362 - if weight_name not in name:

363 - continue

364 - name = name.replace(weight_name, param_name)

365 - # Skip loading extra bias for GPTQ models.

366 - if name.endswith(".bias") and name not in params_dict:

367 - continue

368 - param = params_dict[name]

369 - weight_loader = param.weight_loader # type: ignore

370 - weight_loader(param, loaded_weight, shard_id)

371 - break

372 - else:

373 - # Skip loading extra bias for GPTQ models.

374 - if name.endswith(".bias") and name not in params_dict:

375 - continue

376 - param = params_dict[name]

377 - weight_loader = getattr(param, "weight_loader", default_weight_loader)

378 - weight_loader(param, loaded_weight)

379 - loaded_params.add(name)

380 - return loaded_params

381 -

345

346 class Olmo2ForCausalLM(nn.Module, SupportsPP, SupportsLoRA):

347 """

348 Extremely barebones HF model wrapper.

349 """

350

351 + hf_to_vllm_mapper = WeightsMapper(

352 + orig_to_new_stacked={

353 + # weight_name: (param_name, shard_id)

354 + ".q_proj": (".qkv_proj", "q"),

355 + ".k_proj": (".qkv_proj", "k"),

356 + ".v_proj": (".qkv_proj", "v"),

357 + ".gate_proj": (".gate_up_proj", 0),

358 + ".up_proj": (".gate_up_proj", 1),

359 + }

360 + )

361 packed_modules_mapping = {

389 - "qkv_proj": [

390 - "q_proj",

391 - "k_proj",

392 - "v_proj",

393 - ],

394 - "gate_up_proj": [

395 - "gate_proj",

396 - "up_proj",

397 - ],

362 + "qkv_proj": ["q_proj", "k_proj", "v_proj"],

363 + "gate_up_proj": ["gate_proj", "up_proj"],

364 }

365

366 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

@@ -451,4 +417,4 @@ class Olmo2ForCausalLM(nn.Module, SupportsPP, SupportsLoRA):

417 ["lm_head.weight"] if self.config.tie_word_embeddings else None

418 ),

419 )

454 - return loader.load_weights(weights)

420 + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

vllm/model_executor/models/opt.py

+10 -42

@@ -44,14 +44,12 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (

44 ParallelLMHead,

45 VocabParallelEmbedding,

46 )

47 -from vllm.model_executor.model_loader.weight_utils import default_weight_loader

47 from vllm.sequence import IntermediateTensors

48

49 from .interfaces import SupportsLoRA, SupportsPP

50 from .utils import (

51 AutoWeightsLoader,

52 WeightsMapper,

54 - is_pp_missing_parameter,

53 make_empty_intermediate_tensors_factory,

54 make_layers,

55 maybe_prefix,

@@ -325,52 +323,22 @@ class OPTModel(nn.Module):

323 input_ids, positions, intermediate_tensors, inputs_embeds=inputs_embeds

324 )

325

328 - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

329 - stacked_params_mapping = [

330 - # (param_name, shard_name, shard_id)

331 - ("qkv_proj", "q_proj", "q"),

332 - ("qkv_proj", "k_proj", "k"),

333 - ("qkv_proj", "v_proj", "v"),

334 - ]

335 - params_dict = dict(self.named_parameters(remove_duplicate=False))

336 - loaded_params: set[str] = set()

337 - for name, loaded_weight in weights:

338 - for param_name, weight_name, shard_id in stacked_params_mapping:

339 - if weight_name not in name:

340 - continue

341 - name = name.replace(weight_name, param_name)

342 - # Skip loading extra bias for GPTQ models.

343 - if name.endswith(".bias") and name not in params_dict:

344 - continue

345 - if is_pp_missing_parameter(name, self):

346 - continue

347 - param = params_dict[name]

348 - weight_loader = param.weight_loader

349 - weight_loader(param, loaded_weight, shard_id)

350 - break

351 - else:

352 - # Skip loading extra bias for GPTQ models.

353 - if name.endswith(".bias") and name not in params_dict:

354 - continue

355 - if is_pp_missing_parameter(name, self):

356 - continue

357 - param = params_dict[name]

358 - weight_loader = getattr(param, "weight_loader", default_weight_loader)

359 - weight_loader(param, loaded_weight)

360 - loaded_params.add(name)

361 - return loaded_params

362 -

326

327 class OPTForCausalLM(nn.Module, SupportsPP, SupportsLoRA):

365 - packed_modules_mapping = {

366 - "qkv_proj": ["q_proj", "k_proj", "v_proj"],

367 - }

368 -

328 hf_to_vllm_mapper = WeightsMapper(

329 + orig_to_new_stacked={

330 + # weight_name: (param_name, shard_id)

331 + ".q_proj": (".qkv_proj", "q"),

332 + ".k_proj": (".qkv_proj", "k"),

333 + ".v_proj": (".qkv_proj", "v"),

334 + },

335 orig_to_new_prefix={

336 "decoder.": "model.decoder.",

372 - }

337 + },

338 )

339 + packed_modules_mapping = {

340 + "qkv_proj": ["q_proj", "k_proj", "v_proj"],

341 + }

342

343 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

344 super().__init__()

vllm/model_executor/models/orion.py

+13 -40

@@ -32,13 +32,12 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (

32 ParallelLMHead,

33 VocabParallelEmbedding,

34 )

35 -from vllm.model_executor.model_loader.weight_utils import default_weight_loader

35 from vllm.sequence import IntermediateTensors

36

37 from .interfaces import SupportsPP

38 from .utils import (

39 AutoWeightsLoader,

41 - is_pp_missing_parameter,

40 + WeightsMapper,

41 make_empty_intermediate_tensors_factory,

42 make_layers,

43 maybe_prefix,

@@ -277,45 +276,19 @@ class OrionModel(nn.Module):

276 hidden_states = self.norm(hidden_states)

277 return hidden_states

278

280 - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

281 - stacked_params_mapping = [

282 - # (param_name, shard_name, shard_id)

283 - ("qkv_proj", "q_proj", "q"),

284 - ("qkv_proj", "k_proj", "k"),

285 - ("qkv_proj", "v_proj", "v"),

286 - ("gate_up_proj", "gate_proj", 0),

287 - ("gate_up_proj", "up_proj", 1),

288 - ]

289 - params_dict = dict(self.named_parameters())

290 - loaded_params: set[str] = set()

291 - for name, loaded_weight in weights:

292 - for param_name, weight_name, shard_id in stacked_params_mapping:

293 - if weight_name not in name:

294 - continue

295 - name = name.replace(weight_name, param_name)

296 - # Skip loading extra bias for GPTQ models.

297 - if name.endswith(".bias") and name not in params_dict:

298 - continue

299 - if is_pp_missing_parameter(name, self):

300 - continue

301 - param = params_dict[name]

302 - weight_loader = param.weight_loader

303 - weight_loader(param, loaded_weight, shard_id)

304 - break

305 - else:

306 - # Skip loading extra bias for GPTQ models.

307 - if name.endswith(".bias") and name not in params_dict:

308 - continue

309 - if is_pp_missing_parameter(name, self):

310 - continue

311 - param = params_dict[name]

312 - weight_loader = getattr(param, "weight_loader", default_weight_loader)

313 - weight_loader(param, loaded_weight)

314 - loaded_params.add(name)

315 - return loaded_params

316 -

279

280 class OrionForCausalLM(nn.Module, SupportsPP):

281 + hf_to_vllm_mapper = WeightsMapper(

282 + orig_to_new_stacked={

283 + # weight_name: (param_name, shard_id)

284 + ".q_proj": (".qkv_proj", "q"),

285 + ".k_proj": (".qkv_proj", "k"),

286 + ".v_proj": (".qkv_proj", "v"),

287 + ".gate_proj": (".gate_up_proj", 0),

288 + ".up_proj": (".gate_up_proj", 1),

289 + }

290 + )

291 +

292 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

293 super().__init__()

294 config = vllm_config.model_config.hf_config

@@ -362,4 +335,4 @@ class OrionForCausalLM(nn.Module, SupportsPP):

335

336 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

337 loader = AutoWeightsLoader(self)

365 - return loader.load_weights(weights)

338 + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

vllm/model_executor/models/ouro.py

+14 -61

@@ -51,16 +51,13 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (

51 ParallelLMHead,

52 VocabParallelEmbedding,

53 )

54 -from vllm.model_executor.model_loader.weight_utils import (

55 - default_weight_loader,

56 - maybe_remap_kv_scale_name,

57 -)

54 from vllm.sequence import IntermediateTensors

55 from vllm.v1.attention.backend import AttentionType

56

57 from .interfaces import SupportsLoRA

58 from .utils import (

59 AutoWeightsLoader,

60 + WeightsMapper,

61 extract_layer_index,

62 make_empty_intermediate_tensors_factory,

63 make_layers,

@@ -376,65 +373,21 @@ class OuroModel(nn.Module):

373 hidden_states, _ = self.norm(hidden_states, residual)

374 return hidden_states

375

379 - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

380 - stacked_params_mapping = [

381 - # (param_name, shard_name, shard_id)

382 - ("qkv_proj", "q_proj", "q"),

383 - ("qkv_proj", "k_proj", "k"),

384 - ("qkv_proj", "v_proj", "v"),

385 - ("gate_up_proj", "gate_proj", 0),

386 - ("gate_up_proj", "up_proj", 1),

387 - ]

388 - params_dict = dict(self.named_parameters(remove_duplicate=False))

389 - loaded_params: set[str] = set()

390 - for name, loaded_weight in weights:

391 - if "rotary_emb.inv_freq" in name:

392 - continue

393 - for param_name, weight_name, shard_id in stacked_params_mapping:

394 - if weight_name not in name:

395 - continue

396 - name = name.replace(weight_name, param_name)

397 - # Skip loading extra bias for GPTQ models.

398 - if name.endswith(".bias") and name not in params_dict:

399 - continue

400 - if name.endswith("scale"):

401 - # Remapping the name of FP8 kv-scale.

402 - name = maybe_remap_kv_scale_name(name, params_dict)

403 - if name is None:

404 - continue

405 - param = params_dict[name]

406 - weight_loader = getattr(param, "weight_loader", default_weight_loader)

407 - if weight_loader == default_weight_loader:

408 - weight_loader(param, loaded_weight)

409 - else:

410 - weight_loader(param, loaded_weight, shard_id)

411 - break

412 - else:

413 - # Skip loading extra bias for GPTQ models.

414 - if name.endswith(".bias") and name not in params_dict:

415 - continue

416 - # Remapping the name of FP8 kv-scale.

417 - name = maybe_remap_kv_scale_name(name, params_dict)

418 - if name is None:

419 - continue

420 - param = params_dict[name]

421 - weight_loader = getattr(param, "weight_loader", default_weight_loader)

422 - weight_loader(param, loaded_weight)

423 - loaded_params.add(name)

424 - return loaded_params

425 -

376

377 class OuroForCausalLM(nn.Module, SupportsLoRA):

378 + hf_to_vllm_mapper = WeightsMapper(

379 + orig_to_new_stacked={

380 + # weight_name: (param_name, shard_id)

381 + ".q_proj": (".qkv_proj", "q"),

382 + ".k_proj": (".qkv_proj", "k"),

383 + ".v_proj": (".qkv_proj", "v"),

384 + ".gate_proj": (".gate_up_proj", 0),

385 + ".up_proj": (".gate_up_proj", 1),

386 + }

387 + )

388 packed_modules_mapping = {

429 - "qkv_proj": [

430 - "q_proj",

431 - "k_proj",

432 - "v_proj",

433 - ],

434 - "gate_up_proj": [

435 - "gate_proj",

436 - "up_proj",

437 - ],

389 + "qkv_proj": ["q_proj", "k_proj", "v_proj"],

390 + "gate_up_proj": ["gate_proj", "up_proj"],

391 }

392

393 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

@@ -492,4 +445,4 @@ class OuroForCausalLM(nn.Module, SupportsLoRA):

445 self,

446 skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),

447 )

495 - return loader.load_weights(weights)

448 + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

vllm/model_executor/models/phi.py

+11 -49

@@ -62,13 +62,12 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (

62 ParallelLMHead,

63 VocabParallelEmbedding,

64 )

65 -from vllm.model_executor.model_loader.weight_utils import default_weight_loader

65 from vllm.sequence import IntermediateTensors

66

67 from .interfaces import SupportsLoRA, SupportsPP

68 from .utils import (

69 AutoWeightsLoader,

71 - is_pp_missing_parameter,

70 + WeightsMapper,

71 make_empty_intermediate_tensors_factory,

72 make_layers,

73 maybe_prefix,

@@ -257,55 +256,18 @@ class PhiModel(nn.Module):

256

257 return hidden_states

258

260 - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

261 - stacked_params_mapping = [

262 - # (param_name, shard_name, shard_id)

263 - ("qkv_proj", "q_proj", "q"),

264 - ("qkv_proj", "k_proj", "k"),

265 - ("qkv_proj", "v_proj", "v"),

266 - ]

267 - params_dict = dict(self.named_parameters())

268 - loaded_params: set[str] = set()

269 -

270 - for name, loaded_weight in weights:

271 - if "rotary_emb.inv_freq" in name:

272 - continue

273 -

274 - for param_name, weight_name, shard_id in stacked_params_mapping:

275 - if weight_name not in name:

276 - continue

277 - name = name.replace(weight_name, param_name)

278 - # Skip loading extra bias for GPTQ models.

279 - if name.endswith(".bias") and name not in params_dict:

280 - continue

281 - if is_pp_missing_parameter(name, self):

282 - continue

283 - param = params_dict[name]

284 - weight_loader = param.weight_loader

285 - weight_loader(param, loaded_weight, shard_id)

286 - break

287 - else:

288 - # Skip loading extra bias for GPTQ models.

289 - if name.endswith(".bias") and name not in params_dict:

290 - continue

291 - # pylint: disable=E1136

292 -

293 - if is_pp_missing_parameter(name, self):

294 - continue

295 - param = params_dict[name]

296 - weight_loader = getattr(param, "weight_loader", default_weight_loader)

297 - weight_loader(param, loaded_weight)

298 - loaded_params.add(name)

299 - return loaded_params

300 -

259

260 class PhiForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

261 + hf_to_vllm_mapper = WeightsMapper(

262 + orig_to_new_stacked={

263 + # weight_name: (param_name, shard_id)

264 + ".q_proj": (".qkv_proj", "q"),

265 + ".k_proj": (".qkv_proj", "k"),

266 + ".v_proj": (".qkv_proj", "v"),

267 + }

268 + )

269 packed_modules_mapping = {

304 - "qkv_proj": [

305 - "q_proj",

306 - "k_proj",

307 - "v_proj",

308 - ]

270 + "qkv_proj": ["q_proj", "k_proj", "v_proj"],

271 }

272

273 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

@@ -360,4 +322,4 @@ class PhiForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

322

323 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

324 loader = AutoWeightsLoader(self)

363 - return loader.load_weights(weights)

325 + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

vllm/model_executor/models/qwen2.py

+19 -66

@@ -54,10 +54,6 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (

54 ParallelLMHead,

55 VocabParallelEmbedding,

56 )

57 -from vllm.model_executor.model_loader.weight_utils import (

58 - default_weight_loader,

59 - maybe_remap_kv_scale_name,

60 -)

57 from vllm.sequence import IntermediateTensors

58 from vllm.transformers_utils.config import is_interleaved, set_default_rope_theta

59 from vllm.v1.attention.backend import AttentionType

@@ -68,12 +64,13 @@ from .interfaces import (

64 SupportsEagle3,

65 SupportsLoRA,

66 SupportsPP,

67 + SupportsQuant,

68 )

69 from .utils import (

70 AutoWeightsLoader,

71 PPMissingLayer,

72 + WeightsMapper,

73 extract_layer_index,

76 - is_pp_missing_parameter,

74 make_empty_intermediate_tensors_factory,

75 make_layers,

76 maybe_prefix,

@@ -323,6 +320,17 @@ class Qwen2DecoderLayer(nn.Module):

320 }

321 )

322 class Qwen2Model(nn.Module, EagleModelMixin):

323 + hf_to_vllm_mapper = WeightsMapper(

324 + orig_to_new_stacked={

325 + # weight_name: (param_name, shard_id)

326 + ".q_proj": (".qkv_proj", "q"),

327 + ".k_proj": (".qkv_proj", "k"),

328 + ".v_proj": (".qkv_proj", "v"),

329 + ".gate_proj": (".gate_up_proj", 0),

330 + ".up_proj": (".gate_up_proj", 1),

331 + }

332 + )

333 +

334 def __init__(

335 self,

336 *,

@@ -426,72 +434,17 @@ class Qwen2Model(nn.Module, EagleModelMixin):

434 return hidden_states

435

436 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

429 - stacked_params_mapping = [

430 - # (param_name, shard_name, shard_id)

431 - ("qkv_proj", "q_proj", "q"),

432 - ("qkv_proj", "k_proj", "k"),

433 - ("qkv_proj", "v_proj", "v"),

434 - ("gate_up_proj", "gate_proj", 0),

435 - ("gate_up_proj", "up_proj", 1),

436 - ]

437 - params_dict = dict(self.named_parameters(remove_duplicate=False))

438 - loaded_params: set[str] = set()

439 - for name, loaded_weight in weights:

440 - if "rotary_emb.inv_freq" in name:

441 - continue

442 - for param_name, weight_name, shard_id in stacked_params_mapping:

443 - if weight_name not in name:

444 - continue

445 - name = name.replace(weight_name, param_name)

446 - # Skip loading extra bias for GPTQ models.

447 - if name.endswith(".bias") and name not in params_dict:

448 - continue

449 - if is_pp_missing_parameter(name, self):

450 - continue

451 - if name.endswith("scale"):

452 - # Remapping the name of FP8 kv-scale.

453 - name = maybe_remap_kv_scale_name(name, params_dict)

454 - if name is None:

455 - continue

456 - param = params_dict[name]

457 - weight_loader = getattr(param, "weight_loader", default_weight_loader)

458 - if weight_loader == default_weight_loader:

459 - weight_loader(param, loaded_weight)

460 - else:

461 - weight_loader(param, loaded_weight, shard_id)

462 - break

463 - else:

464 - # Skip loading extra bias for GPTQ models.

465 - if name.endswith(".bias") and name not in params_dict:

466 - continue

467 - # Remapping the name of FP8 kv-scale.

468 - name = maybe_remap_kv_scale_name(name, params_dict)

469 - if name is None:

470 - continue

471 - if is_pp_missing_parameter(name, self):

472 - continue

473 - if name not in params_dict:

474 - continue

475 - param = params_dict[name]

476 - weight_loader = getattr(param, "weight_loader", default_weight_loader)

477 - weight_loader(param, loaded_weight)

478 - loaded_params.add(name)

479 - return loaded_params

437 + loader = AutoWeightsLoader(self)

438 + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

439

440

441 class Qwen2ForCausalLM(

483 - nn.Module, SupportsLoRA, SupportsPP, SupportsEagle, SupportsEagle3

442 + nn.Module, SupportsLoRA, SupportsPP, SupportsEagle, SupportsEagle3, SupportsQuant

443 ):

444 + hf_to_vllm_mapper = Qwen2Model.hf_to_vllm_mapper

445 packed_modules_mapping = {

486 - "qkv_proj": [

487 - "q_proj",

488 - "k_proj",

489 - "v_proj",

490 - ],

491 - "gate_up_proj": [

492 - "gate_proj",

493 - "up_proj",

494 - ],

446 + "qkv_proj": ["q_proj", "k_proj", "v_proj"],

447 + "gate_up_proj": ["gate_proj", "up_proj"],

448 }

449

450 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

vllm/model_executor/models/qwen2_rm.py

+3 -9

@@ -28,16 +28,10 @@ class Qwen2RewardBaseModel(nn.Module, SupportsLoRA, SupportsPP):

28 is_pooling_model = True

29 pooler: Pooler

30

31 + hf_to_vllm_mapper = Qwen2Model.hf_to_vllm_mapper

32 packed_modules_mapping = {

32 - "qkv_proj": [

33 - "q_proj",

34 - "k_proj",

35 - "v_proj",

36 - ],

37 - "gate_up_proj": [

38 - "gate_proj",

39 - "up_proj",

40 - ],

33 + "qkv_proj": ["q_proj", "k_proj", "v_proj"],

34 + "gate_up_proj": ["gate_proj", "up_proj"],

35 }

36

37 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

vllm/model_executor/models/qwen3.py

+3 -10

@@ -267,18 +267,11 @@ class Qwen3Model(Qwen2Model):

267 class Qwen3ForCausalLM(

268 LocalArgmaxMixin, nn.Module, SupportsLoRA, SupportsPP, SupportsEagle, SupportsEagle3

269 ):

270 + hf_to_vllm_mapper = Qwen3Model.hf_to_vllm_mapper

271 packed_modules_mapping = {

271 - "qkv_proj": [

272 - "q_proj",

273 - "k_proj",

274 - "v_proj",

275 - ],

276 - "gate_up_proj": [

277 - "gate_proj",

278 - "up_proj",

279 - ],

272 + "qkv_proj": ["q_proj", "k_proj", "v_proj"],

273 + "gate_up_proj": ["gate_proj", "up_proj"],

274 }

281 -

275 embedding_modules = {

276 "embed_tokens": "input_embeddings",

277 "lm_head": "output_embeddings",

vllm/model_executor/models/rnj1.py

+14 -72

@@ -30,18 +30,14 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (

30 ParallelLMHead,

31 VocabParallelEmbedding,

32 )

33 -from vllm.model_executor.model_loader.weight_utils import (

34 - default_weight_loader,

35 - maybe_remap_kv_scale_name,

36 -)

33 from vllm.sequence import IntermediateTensors

34 from vllm.v1.attention.backend import AttentionType

35

36 from .interfaces import SupportsLoRA, SupportsPP

37 from .utils import (

38 AutoWeightsLoader,

39 + WeightsMapper,

40 extract_layer_index,

44 - is_pp_missing_parameter,

41 make_empty_intermediate_tensors_factory,

42 make_layers,

43 maybe_prefix,

@@ -331,75 +327,21 @@ class Rnj1Model(nn.Module):

327 hidden_states, _ = self.norm(hidden_states, residual)

328 return hidden_states

329

334 - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

335 - stacked_params_mapping = [

336 - # (param_name, shard_name, shard_id)

337 - ("qkv_proj", "q_proj", "q"),

338 - ("qkv_proj", "k_proj", "k"),

339 - ("qkv_proj", "v_proj", "v"),

340 - ("gate_up_proj", "gate_proj", 0),

341 - ("gate_up_proj", "up_proj", 1),

342 - ]

343 - params_dict = dict(self.named_parameters())

344 - loaded_params: set[str] = set()

345 - for name, loaded_weight in weights:

346 - if (

347 - self.quant_config

348 - and self.quant_config.get_name() == "gguf"

349 - and name.endswith("norm.weight")

350 - ):

351 - loaded_weight -= 1

352 -

353 - if name.endswith((".k_scale", ".v_scale", ".q_scale", ".prob_scale")):

354 - remapped_name = maybe_remap_kv_scale_name(name, params_dict)

355 - if remapped_name is not None and remapped_name in params_dict:

356 - param = params_dict[remapped_name]

357 - weight_loader = getattr(

358 - param, "weight_loader", default_weight_loader

359 - )

360 - weight_loader(param, loaded_weight)

361 - loaded_params.add(remapped_name)

362 - continue

363 -

364 - for param_name, shard_name, shard_id in stacked_params_mapping:

365 - if shard_name not in name:

366 - continue

367 - name = name.replace(shard_name, param_name)

368 - if name.endswith(".bias") and name not in params_dict:

369 - continue

370 - if is_pp_missing_parameter(name, self):

371 - continue

372 - param = params_dict[name]

373 - weight_loader = param.weight_loader

374 - weight_loader(param, loaded_weight, shard_id)

375 - break

376 - else:

377 - if name.endswith(".bias") and name not in params_dict:

378 - continue

379 - name = maybe_remap_kv_scale_name(name, params_dict)

380 - if name is None:

381 - continue

382 - if is_pp_missing_parameter(name, self):

383 - continue

384 - param = params_dict[name]

385 - weight_loader = getattr(param, "weight_loader", default_weight_loader)

386 - weight_loader(param, loaded_weight)

387 - loaded_params.add(name)

388 -

389 - return loaded_params

390 -

330

331 class Rnj1ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

332 + hf_to_vllm_mapper = WeightsMapper(

333 + orig_to_new_stacked={

334 + # weight_name: (param_name, shard_id)

335 + ".q_proj": (".qkv_proj", "q"),

336 + ".k_proj": (".qkv_proj", "k"),

337 + ".v_proj": (".qkv_proj", "v"),

338 + ".gate_proj": (".gate_up_proj", 0),

339 + ".up_proj": (".gate_up_proj", 1),

340 + }

341 + )

342 packed_modules_mapping = {

394 - "qkv_proj": [

395 - "q_proj",

396 - "k_proj",

397 - "v_proj",

398 - ],

399 - "gate_up_proj": [

400 - "gate_proj",

401 - "up_proj",

402 - ],

343 + "qkv_proj": ["q_proj", "k_proj", "v_proj"],

344 + "gate_up_proj": ["gate_proj", "up_proj"],

345 }

346

347 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

@@ -457,4 +399,4 @@ class Rnj1ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

399 self,

400 skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),

401 )

460 - return loader.load_weights(weights)

402 + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

vllm/model_executor/models/seed_oss.py

+14 -58

@@ -49,10 +49,6 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (

49 ParallelLMHead,

50 VocabParallelEmbedding,

51 )

52 -from vllm.model_executor.model_loader.weight_utils import (

53 - default_weight_loader,

54 - maybe_remap_kv_scale_name,

55 -)

52 from vllm.sequence import IntermediateTensors

53 from vllm.transformers_utils.config import set_default_rope_theta

54 from vllm.v1.attention.backend import AttentionType

@@ -61,7 +57,7 @@ from .interfaces import SupportsLoRA, SupportsPP

57 from .utils import (

58 AutoWeightsLoader,

59 PPMissingLayer,

64 - is_pp_missing_parameter,

60 + WeightsMapper,

61 make_empty_intermediate_tensors_factory,

62 make_layers,

63 maybe_prefix,

@@ -362,61 +358,21 @@ class SeedOssModel(nn.Module):

358 hidden_states, _ = self.norm(hidden_states, residual)

359 return hidden_states

360

365 - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

366 - stacked_params_mapping = [

367 - # (param_name, shard_name, shard_id)

368 - ("qkv_proj", "q_proj", "q"),

369 - ("qkv_proj", "k_proj", "k"),

370 - ("qkv_proj", "v_proj", "v"),

371 - ("gate_up_proj", "gate_proj", 0),

372 - ("gate_up_proj", "up_proj", 1),

373 - ]

374 - params_dict = dict(self.named_parameters(remove_duplicate=False))

375 - loaded_params: set[str] = set()

376 - for name, loaded_weight in weights:

377 - if "rotary_emb.inv_freq" in name:

378 - continue

379 - for param_name, weight_name, shard_id in stacked_params_mapping:

380 - if weight_name not in name:

381 - continue

382 - name = name.replace(weight_name, param_name)

383 - # Skip loading extra bias for GPTQ models.

384 - if name.endswith(".bias") and name not in params_dict:

385 - continue

386 - if is_pp_missing_parameter(name, self):

387 - continue

388 - param = params_dict[name]

389 - weight_loader = param.weight_loader

390 - weight_loader(param, loaded_weight, shard_id)

391 - break

392 - else:

393 - # Skip loading extra bias for GPTQ models.

394 - if name.endswith(".bias") and name not in params_dict:

395 - continue

396 - # Remapping the name of FP8 kv-scale.

397 - name = maybe_remap_kv_scale_name(name, params_dict)

398 - if name is None:

399 - continue

400 - if is_pp_missing_parameter(name, self):

401 - continue

402 - param = params_dict[name]

403 - weight_loader = getattr(param, "weight_loader", default_weight_loader)

404 - weight_loader(param, loaded_weight)

405 - loaded_params.add(name)

406 - return loaded_params

407 -

361

362 class SeedOssForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

363 + hf_to_vllm_mapper = WeightsMapper(

364 + orig_to_new_stacked={

365 + # weight_name: (param_name, shard_id)

366 + ".q_proj": (".qkv_proj", "q"),

367 + ".k_proj": (".qkv_proj", "k"),

368 + ".v_proj": (".qkv_proj", "v"),

369 + ".gate_proj": (".gate_up_proj", 0),

370 + ".up_proj": (".gate_up_proj", 1),

371 + }

372 + )

373 packed_modules_mapping = {

411 - "qkv_proj": [

412 - "q_proj",

413 - "k_proj",

414 - "v_proj",

415 - ],

416 - "gate_up_proj": [

417 - "gate_proj",

418 - "up_proj",

419 - ],

374 + "qkv_proj": ["q_proj", "k_proj", "v_proj"],

375 + "gate_up_proj": ["gate_proj", "up_proj"],

376 }

377

378 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

@@ -477,4 +433,4 @@ class SeedOssForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

433 self,

434 skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),

435 )

480 - return loader.load_weights(weights)

436 + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

vllm/model_executor/models/solar.py

+14 -62

@@ -48,17 +48,13 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (

48 ParallelLMHead,

49 VocabParallelEmbedding,

50 )

51 -from vllm.model_executor.model_loader.weight_utils import (

52 - default_weight_loader,

53 - maybe_remap_kv_scale_name,

54 -)

51 from vllm.sequence import IntermediateTensors

52

53 from .interfaces import SupportsLoRA, SupportsPP

54 from .utils import (

55 AutoWeightsLoader,

56 PPMissingLayer,

61 - is_pp_missing_parameter,

57 + WeightsMapper,

58 make_empty_intermediate_tensors_factory,

59 make_layers,

60 maybe_prefix,

@@ -347,66 +343,22 @@ class SolarModel(nn.Module):

343 hidden_states, _ = self.norm(hidden_states, residual)

344 return hidden_states

345

350 - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

351 - stacked_params_mapping = [

352 - # (param_name, shard_name, shard_id)

353 - (".qkv_proj", ".q_proj", "q"),

354 - (".qkv_proj", ".k_proj", "k"),

355 - (".qkv_proj", ".v_proj", "v"),

356 - (".gate_up_proj", ".gate_proj", 0),

357 - (".gate_up_proj", ".up_proj", 1),

358 - ]

359 - params_dict = dict(self.named_parameters())

360 - loaded_params: set[str] = set()

361 - for name, loaded_weight in weights:

362 - for param_name, weight_name, shard_id in stacked_params_mapping:

363 - if weight_name not in name:

364 - continue

365 - name = name.replace(weight_name, param_name)

366 - # Skip loading extra bias for GPTQ models.

367 - if name.endswith(".bias") and name not in params_dict:

368 - continue

369 -

370 - if is_pp_missing_parameter(name, self):

371 - continue

372 -

373 - param = params_dict[name]

374 - weight_loader = param.weight_loader

375 - weight_loader(param, loaded_weight, shard_id)

376 -

377 - break

378 - else:

379 - # Skip loading extra bias for GPTQ models.

380 - if name.endswith(".bias") and name not in params_dict:

381 - continue

382 - # Remapping the name of FP8 kv-scale.

383 - name = maybe_remap_kv_scale_name(name, params_dict)

384 - if name is None:

385 - continue

386 -

387 - if is_pp_missing_parameter(name, self):

388 - continue

389 -

390 - param = params_dict[name]

391 - weight_loader = getattr(param, "weight_loader", default_weight_loader)

392 - weight_loader(param, loaded_weight)

393 - loaded_params.add(name)

394 - return loaded_params

395 -

346

347 class SolarForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

348 + hf_to_vllm_mapper = WeightsMapper(

349 + orig_to_new_stacked={

350 + # weight_name: (param_name, shard_id)

351 + ".q_proj": (".qkv_proj", "q"),

352 + ".k_proj": (".qkv_proj", "k"),

353 + ".v_proj": (".qkv_proj", "v"),

354 + ".gate_proj": (".gate_up_proj", 0),

355 + ".up_proj": (".gate_up_proj", 1),

356 + }

357 + )

358 packed_modules_mapping = {

399 - "qkv_proj": [

400 - "q_proj",

401 - "k_proj",

402 - "v_proj",

403 - ],

404 - "gate_up_proj": [

405 - "gate_proj",

406 - "up_proj",

407 - ],

359 + "qkv_proj": ["q_proj", "k_proj", "v_proj"],

360 + "gate_up_proj": ["gate_proj", "up_proj"],

361 }

409 -

362 # LoRA specific attributes

363 embedding_modules = {

364 "embed_tokens": "input_embeddings",

@@ -468,4 +420,4 @@ class SolarForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

420

421 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

422 loader = AutoWeightsLoader(self)

471 - return loader.load_weights(weights)

423 + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

vllm/model_executor/models/stablelm.py

+13 -40

@@ -45,13 +45,12 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (

45 ParallelLMHead,

46 VocabParallelEmbedding,

47 )

48 -from vllm.model_executor.model_loader.weight_utils import default_weight_loader

48 from vllm.sequence import IntermediateTensors

49

50 from .interfaces import SupportsPP

51 from .utils import (

52 AutoWeightsLoader,

54 - is_pp_missing_parameter,

53 + WeightsMapper,

54 make_empty_intermediate_tensors_factory,

55 make_layers,

56 maybe_prefix,

@@ -266,45 +265,19 @@ class StableLMEpochModel(nn.Module):

265 hidden_states = self.norm(hidden_states)

266 return hidden_states

267

269 - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

270 - stacked_params_mapping = [

271 - # (param_name, shard_name, shard_id)

272 - ("qkv_proj", "q_proj", "q"),

273 - ("qkv_proj", "k_proj", "k"),

274 - ("qkv_proj", "v_proj", "v"),

275 - ("gate_up_proj", "gate_proj", 0),

276 - ("gate_up_proj", "up_proj", 1),

277 - ]

278 - params_dict = dict(self.named_parameters())

279 - loaded_params: set[str] = set()

280 - for name, loaded_weight in weights:

281 - for param_name, weight_name, shard_id in stacked_params_mapping:

282 - if weight_name not in name:

283 - continue

284 - name = name.replace(weight_name, param_name)

285 - # Skip loading extra bias for GPTQ models.

286 - if name.endswith(".bias") and name not in params_dict:

287 - continue

288 - if is_pp_missing_parameter(name, self):

289 - continue

290 - param = params_dict[name]

291 - weight_loader = param.weight_loader

292 - weight_loader(param, loaded_weight, shard_id)

293 - break

294 - else:

295 - # Skip loading extra bias for GPTQ models.

296 - if name.endswith(".bias") and name not in params_dict:

297 - continue

298 - if is_pp_missing_parameter(name, self):

299 - continue

300 - param = params_dict[name]

301 - weight_loader = getattr(param, "weight_loader", default_weight_loader)

302 - weight_loader(param, loaded_weight)

303 - loaded_params.add(name)

304 - return loaded_params

305 -

268

269 class StablelmForCausalLM(nn.Module, SupportsPP):

270 + hf_to_vllm_mapper = WeightsMapper(

271 + orig_to_new_stacked={

272 + # weight_name: (param_name, shard_id)

273 + ".q_proj": (".qkv_proj", "q"),

274 + ".k_proj": (".qkv_proj", "k"),

275 + ".v_proj": (".qkv_proj", "v"),

276 + ".gate_proj": (".gate_up_proj", 0),

277 + ".up_proj": (".gate_up_proj", 1),

278 + }

279 + )

280 +

281 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

282 super().__init__()

283 config = vllm_config.model_config.hf_config

@@ -351,4 +324,4 @@ class StablelmForCausalLM(nn.Module, SupportsPP):

324

325 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

326 loader = AutoWeightsLoader(self)

354 - return loader.load_weights(weights)

327 + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

vllm/model_executor/models/starcoder2.py

+11 -39

@@ -45,16 +45,12 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (

45 ParallelLMHead,

46 VocabParallelEmbedding,

47 )

48 -from vllm.model_executor.model_loader.weight_utils import (

49 - default_weight_loader,

50 - maybe_remap_kv_scale_name,

51 -)

48 from vllm.sequence import IntermediateTensors

49

50 from .interfaces import SupportsPP

51 from .utils import (

52 AutoWeightsLoader,

57 - is_pp_missing_parameter,

53 + WeightsMapper,

54 make_empty_intermediate_tensors_factory,

55 make_layers,

56 maybe_prefix,

@@ -272,41 +268,17 @@ class Starcoder2Model(nn.Module):

268 hidden_states = self.norm(hidden_states)

269 return hidden_states

270

275 - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

276 - stacked_params_mapping = [

277 - # (param_name, shard_name, shard_id)

278 - ("qkv_proj", "q_proj", "q"),

279 - ("qkv_proj", "k_proj", "k"),

280 - ("qkv_proj", "v_proj", "v"),

281 - ]

282 -

283 - params_dict = dict(self.named_parameters(remove_duplicate=False))

284 - loaded_params: set[str] = set()

285 - for name, loaded_weight in weights:

286 - for param_name, weight_name, shard_id in stacked_params_mapping:

287 - if weight_name not in name:

288 - continue

289 - name = name.replace(weight_name, param_name)

290 - if is_pp_missing_parameter(name, self):

291 - continue

292 - param = params_dict[name]

293 - weight_loader = param.weight_loader

294 - weight_loader(param, loaded_weight, shard_id)

295 - break

296 - else:

297 - name = maybe_remap_kv_scale_name(name, params_dict)

298 - if name is None:

299 - continue

300 - if is_pp_missing_parameter(name, self):

301 - continue

302 - param = params_dict[name]

303 - weight_loader = getattr(param, "weight_loader", default_weight_loader)

304 - weight_loader(param, loaded_weight)

305 - loaded_params.add(name)

306 - return loaded_params

307 -

271

272 class Starcoder2ForCausalLM(nn.Module, SupportsPP):

273 + hf_to_vllm_mapper = WeightsMapper(

274 + orig_to_new_stacked={

275 + # weight_name: (param_name, shard_id)

276 + ".q_proj": (".qkv_proj", "q"),

277 + ".k_proj": (".qkv_proj", "k"),

278 + ".v_proj": (".qkv_proj", "v"),

279 + }

280 + )

281 +

282 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

283 super().__init__()

284 config = vllm_config.model_config.hf_config

@@ -362,4 +334,4 @@ class Starcoder2ForCausalLM(nn.Module, SupportsPP):

334 ["lm_head.weight"] if self.config.tie_word_embeddings else None

335 ),

336 )

365 - return loader.load_weights(weights)

337 + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

vllm/model_executor/models/step1.py

+16 -45

@@ -30,7 +30,6 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (

30 ParallelLMHead,

31 VocabParallelEmbedding,

32 )

33 -from vllm.model_executor.model_loader.weight_utils import default_weight_loader

33 from vllm.model_executor.models.interfaces import (

34 EagleModelMixin,

35 SupportsEagle,

@@ -40,7 +39,7 @@ from vllm.model_executor.models.interfaces import (

39 from vllm.model_executor.models.utils import (

40 AutoWeightsLoader,

41 PPMissingLayer,

43 - is_pp_missing_parameter,

42 + WeightsMapper,

43 make_empty_intermediate_tensors_factory,

44 make_layers,

45 maybe_prefix,

@@ -48,11 +47,6 @@ from vllm.model_executor.models.utils import (

47 from vllm.sequence import IntermediateTensors

48 from vllm.v1.attention.backend import AttentionType

49

51 -STEP_PACKED_MODULES_MAPPING = {

52 - "qkv_proj": ["q_proj", "k_proj", "v_proj"],

53 - "gate_up_proj": ["gate_proj", "up_proj"],

54 -}

55 -

50

51 def _get_step_alibi_slopes(total_num_heads: int) -> torch.Tensor:

52 """Reference ALiBi slopes used by Step models."""

@@ -242,42 +236,6 @@ class StepDecoderLayer(nn.Module):

236 hidden_states = self.mlp(hidden_states)

237 return hidden_states, residual

238

245 - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

246 - stacked_params_mapping = [

247 - (".qkv_proj", ".q_proj", "q"),

248 - (".qkv_proj", ".k_proj", "k"),

249 - (".qkv_proj", ".v_proj", "v"),

250 - (".gate_up_proj", ".gate_proj", 0),

251 - (".gate_up_proj", ".up_proj", 1),

252 - ]

253 - params_dict = dict(self.named_parameters())

254 - loaded_params: set[str] = set()

255 - for name, loaded_weight in weights:

256 - for param_name, weight_name, shard_id in stacked_params_mapping:

257 - if weight_name not in name:

258 - continue

259 - name = name.replace(weight_name, param_name)

260 - # Skip loading extra bias for GPTQ models.

261 - if name.endswith(".bias") and name not in params_dict:

262 - continue

263 - if is_pp_missing_parameter(name, self):

264 - continue

265 - param = params_dict[name]

266 - weight_loader = param.weight_loader

267 - weight_loader(param, loaded_weight, shard_id)

268 - break

269 - else:

270 - # Skip loading extra bias for GPTQ models.

271 - if name.endswith(".bias") and name not in params_dict:

272 - continue

273 - if is_pp_missing_parameter(name, self):

274 - continue

275 - param = params_dict[name]

276 - weight_loader = getattr(param, "weight_loader", default_weight_loader) # type: ignore[name-defined]

277 - weight_loader(param, loaded_weight)

278 - loaded_params.add(name)

279 - return loaded_params

280 -

239

240 class StepDecoderModel(nn.Module, EagleModelMixin):

241 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

@@ -354,7 +312,20 @@ class StepDecoderModel(nn.Module, EagleModelMixin):

312

313

314 class Step1ForCausalLM(nn.Module, SupportsPP, SupportsEagle, SupportsEagle3):

357 - packed_modules_mapping = STEP_PACKED_MODULES_MAPPING

315 + hf_to_vllm_mapper = WeightsMapper(

316 + orig_to_new_stacked={

317 + # weight_name: (param_name, shard_id)

318 + ".q_proj": (".qkv_proj", "q"),

319 + ".k_proj": (".qkv_proj", "k"),

320 + ".v_proj": (".qkv_proj", "v"),

321 + ".gate_proj": (".gate_up_proj", 0),

322 + ".up_proj": (".gate_up_proj", 1),

323 + }

324 + )

325 + packed_modules_mapping = {

326 + "qkv_proj": ["q_proj", "k_proj", "v_proj"],

327 + "gate_up_proj": ["gate_proj", "up_proj"],

328 + }

329

330 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

331 super().__init__()

@@ -413,4 +384,4 @@ class Step1ForCausalLM(nn.Module, SupportsPP, SupportsEagle, SupportsEagle3):

384

385 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

386 loader = AutoWeightsLoader(self)

416 - return loader.load_weights(weights)

387 + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

vllm/model_executor/models/transformers/base.py

-3

@@ -158,9 +158,6 @@ class Base(

158 "Transformers modeling backend does "

159 "not support MXFP4 quantization yet."

160 )

161 - # Skip loading extra bias for GPTQ models.

162 - if "gptq" in quant_method_name:

163 - self.ignore_unexpected_suffixes.append(".bias")

161

162 self._patch_config()

163 from_config_kwargs = dict(

vllm/model_executor/models/utils.py

+67 -24

@@ -4,8 +4,8 @@

4 import itertools

5 from collections.abc import Callable, Iterable, Mapping

6 from contextlib import contextmanager

7 -from dataclasses import dataclass, field

8 -from typing import Any, Literal, Protocol, overload

7 +from dataclasses import dataclass, field, replace

8 +from typing import TYPE_CHECKING, Any, Literal, Protocol, TypeAlias, overload

9

10 import regex as re

11 import torch

@@ -19,9 +19,6 @@ from vllm.distributed import (

19 get_tensor_model_parallel_world_size,

20 )

21 from vllm.logger import init_logger

22 -from vllm.model_executor.layers.quantization.base_config import (

23 - QuantizationConfig,

24 -)

22 from vllm.model_executor.model_loader.reload import (

23 support_quantized_model_reload_from_hp_weights,

24 )

@@ -35,8 +32,13 @@ from vllm.utils.torch_utils import (

32 direct_register_custom_op,

33 )

34

35 +if TYPE_CHECKING:

36 + from vllm.model_executor.layers.quantization import QuantizationConfig

37 +

38 logger = init_logger(__name__)

39

40 +ShardId: TypeAlias = str | int | tuple[int, ...]

41 +

42

43 @dataclass

44 class WeightsMapper:

@@ -47,6 +49,7 @@ class WeightsMapper:

49 orig_to_new_renamings: list[Any] = field(default_factory=list)

50 orig_to_new_regex: Mapping[re.Pattern, str | None] = field(default_factory=dict)

51 orig_to_new_substr: Mapping[str, str | None] = field(default_factory=dict)

52 + orig_to_new_stacked: Mapping[str, tuple[str, ShardId]] = field(default_factory=dict)

53 orig_to_new_prefix: Mapping[str, str | None] = field(default_factory=dict)

54 orig_to_new_suffix: Mapping[str, str | None] = field(default_factory=dict)

55

@@ -59,11 +62,36 @@ class WeightsMapper:

62 ],

63 orig_to_new_regex={**self.orig_to_new_regex, **other.orig_to_new_regex},

64 orig_to_new_substr={**self.orig_to_new_substr, **other.orig_to_new_substr},

65 + orig_to_new_stacked={

66 + **self.orig_to_new_stacked,

67 + **other.orig_to_new_stacked,

68 + },

69 orig_to_new_prefix={**self.orig_to_new_prefix, **other.orig_to_new_prefix},

70 orig_to_new_suffix={**self.orig_to_new_suffix, **other.orig_to_new_suffix},

71 )

72

73 def _map_name(self, key: str) -> str | None:

74 + """Map a weight name (backward-compatible wrapper that discards shard_id)."""

75 + result = self._map_name_with_shard(key)

76 + return result[0] if result is not None else None

77 +

78 + def _map_name_with_shard(self, key: str) -> tuple[str, ShardId | None] | None:

79 + """Map a weight name and extract any shard_id metadata.

80 +

81 + Returns:

82 + (mapped_name, shard_id) if the name should be kept.

83 + None if the name should be dropped.

84 + """

85 + # Deprecation warnings

86 + if key.endswith(".kv_scale"):

87 + logger.warning_once(

88 + "DEPRECATED. Found kv_scale in the checkpoint. "

89 + "This format is deprecated in favor of separate k_scale and "

90 + "v_scale tensors and will be removed in a future release. "

91 + "Functionally, we will remap kv_scale to k_scale and duplicate "

92 + "k_scale to v_scale"

93 + )

94 +

95 for renaming in self.orig_to_new_renamings:

96 key, _ = renaming.rename_source_key(key)

97

@@ -81,6 +109,12 @@ class WeightsMapper:

109

110 key = key.replace(substr, new_key, 1)

111

112 + shard_id: ShardId | None = None

113 + for substr, (new_key, new_shard_id) in self.orig_to_new_stacked.items():

114 + if substr in key:

115 + key = key.replace(substr, new_key, 1)

116 + shard_id = new_shard_id

117 +

118 for prefix, new_key in self.orig_to_new_prefix.items():

119 if key.startswith(prefix):

120 if new_key is None:

@@ -95,16 +129,19 @@ class WeightsMapper:

129

130 key = new_key.join(key.rsplit(suffix, 1))

131

98 - return key

132 + return key, shard_id

133

134 def apply(

135 self, weights: Iterable[tuple[str, torch.Tensor]]

136 ) -> Iterable[tuple[str, torch.Tensor]]:

103 - return (

104 - (out_name, data)

105 - for name, data in weights

106 - if (out_name := self._map_name(name)) is not None

107 - )

137 + for name, data in weights:

138 + result = self._map_name_with_shard(name)

139 + if result is None:

140 + continue

141 + out_name, shard_id = result

142 + if shard_id is not None:

143 + data.shard_id = shard_id

144 + yield out_name, data

145

146 def apply_list(self, values: list[str]) -> list[str]:

147 return [

@@ -120,6 +157,15 @@ class WeightsMapper:

157 if (out_name := self._map_name(name)) is not None

158 }

159

160 + def get_unstacked_mapper(self) -> "WeightsMapper":

161 + """Mapper variant that drops stacked maps, keeping all genuine renames/prefixes.

162 +

163 + Consumers that reference the checkpoint's *unstacked* module names (LoRA name

164 + parsing and the quantization config's layer lists) need the constituent names

165 + (e.g. `q_proj`) to survive rather than being rewritten to the stacked vLLM name

166 + (`qkv_proj`)."""

167 + return replace(self, orig_to_new_stacked={})

168 +

169

170 class AutoWeightsLoader:

171 """

@@ -352,20 +398,19 @@ class AutoWeightsLoader:

398 *,

399 mapper: WeightsMapper | None = None,

400 ) -> set[str]:

401 + # Ignore unexpected biases (typically from GPTQ models)

402 + self.ignore_unexpected_suffixes.append(".bias")

403 +

404 # Many models store quant_config in the base model instead of the causal model.

405 # We look at the causal model's direct children for this reason.

406 modules = (self.module, *self.module.children())

407 iterator = (m.quant_config for m in modules if hasattr(m, "quant_config"))

359 - quant_config = next(iterator, None)

360 - cache_scale_mapper = (

361 - quant_config.get_cache_scale_mapper() if quant_config is not None else None

362 - )

363 - if cache_scale_mapper is not None:

364 - mapper = (

365 - mapper | cache_scale_mapper

366 - if mapper is not None

367 - else cache_scale_mapper

368 - )

408 + if quant_config := next(iterator, None):

409 + # Get mappings and ignore prefixes for KV cache quantization scales

410 + mapper = mapper or WeightsMapper()

411 + mapper |= quant_config.get_cache_scale_mapper()

412 + ignore_unexpected_suffixes = quant_config._ignore_unexpected_suffixes

413 + self.ignore_unexpected_suffixes.extend(ignore_unexpected_suffixes)

414 if mapper is not None:

415 weights = mapper.apply(weights)

416 # filter out weights with first-prefix/substr to skip in name

@@ -734,9 +779,7 @@ def maybe_prefix(prefix: str, name: str) -> str:

779 return name if not prefix else f"{prefix}.{name}"

780

781

737 -def get_draft_quant_config(

738 - vllm_config: VllmConfig,

739 -) -> QuantizationConfig | None:

782 +def get_draft_quant_config(vllm_config: VllmConfig) -> "QuantizationConfig | None":

783 """Get quantization config for Draft models.

784

785 Draft models should use their own quantization config instead of the verifier/target

vllm/model_executor/models/whisper.py

+9 -38

@@ -44,7 +44,6 @@ from vllm.model_executor.layers.linear import (

44 from vllm.model_executor.layers.logits_processor import LogitsProcessor

45 from vllm.model_executor.layers.quantization import QuantizationConfig

46 from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead

47 -from vllm.model_executor.model_loader.weight_utils import default_weight_loader

47 from vllm.model_executor.models.whisper_utils import (

48 ISO639_1_SUPPORTED_LANGS,

49 )

@@ -617,42 +616,6 @@ class WhisperModel(nn.Module):

616 return None

617 return self.encoder(input_features)

618

620 - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

621 - stacked_params_mapping = [

622 - # (param_name, shard_name, shard_id)

623 - (".self_attn.qkv_proj", ".self_attn.q_proj", "q"),

624 - (".self_attn.qkv_proj", ".self_attn.k_proj", "k"),

625 - (".self_attn.qkv_proj", ".self_attn.v_proj", "v"),

626 - # MergedColumnParallelLinear uses integer indices (0, 1)

627 - (".encoder_attn.kv_proj", ".encoder_attn.k_proj", 0),

628 - (".encoder_attn.kv_proj", ".encoder_attn.v_proj", 1),

629 - ]

630 - params_dict = dict(self.named_parameters())

631 - loaded_params: set[str] = set()

632 - for name, loaded_weight in weights:

633 - for param_name, weight_name, shard_id in stacked_params_mapping:

634 - if weight_name not in name:

635 - continue

636 - name = name.replace(weight_name, param_name)

637 - # Skip loading extra bias for GPTQ models.

638 - if name.endswith(".bias") and name not in params_dict:

639 - continue

640 -

641 - param = params_dict[name]

642 - weight_loader = param.weight_loader

643 - weight_loader(param, loaded_weight, shard_id)

644 - break

645 - else:

646 - # Skip loading extra bias for GPTQ models.

647 - if name.endswith(".bias") and name not in params_dict:

648 - continue

649 -

650 - param = params_dict[name]

651 - weight_loader = getattr(param, "weight_loader", default_weight_loader)

652 - weight_loader(param, loaded_weight)

653 - loaded_params.add(name)

654 - return loaded_params

655 -

619

620 class WhisperProcessingInfo(BaseProcessingInfo):

621 def get_hf_config(self) -> WhisperConfig:

@@ -808,7 +771,15 @@ class WhisperForConditionalGeneration(

771 }

772

773 hf_to_vllm_mapper = WeightsMapper(

811 - orig_to_new_substr={".fc1.": ".mlp.fc1.", ".fc2.": ".mlp.fc2."}

774 + orig_to_new_substr={".fc1.": ".mlp.fc1.", ".fc2.": ".mlp.fc2."},

775 + orig_to_new_stacked={

776 + # weight_name: (param_name, shard_id)

777 + ".self_attn.q_proj": (".self_attn.qkv_proj", "q"),

778 + ".self_attn.k_proj": (".self_attn.qkv_proj", "k"),

779 + ".self_attn.v_proj": (".self_attn.qkv_proj", "v"),

780 + ".encoder_attn.k_proj": (".encoder_attn.kv_proj", 0),

781 + ".encoder_attn.v_proj": (".encoder_attn.kv_proj", 1),

782 + },

783 )

784

785 # Whisper only supports audio-conditioned generation.

Remove unnecessary `load_weights` methods (#44589)