From 3989a798249bfa24b6dd22aff599796fcf92dce9 Mon Sep 17 00:00:00 2001 From: Dipika Sikka Date: Fri, 13 Dec 2024 00:07:20 -0500 Subject: [PATCH] [Bugfix] Update starcoder2 to remap k/v scale names for kv_cache quantization (#11148) --- vllm/model_executor/models/starcoder2.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/vllm/model_executor/models/starcoder2.py b/vllm/model_executor/models/starcoder2.py index 15e8f2af52cda..22189a517d313 100644 --- a/vllm/model_executor/models/starcoder2.py +++ b/vllm/model_executor/models/starcoder2.py @@ -37,7 +37,8 @@ from vllm.model_executor.layers.sampler import SamplerOutput, get_sampler from vllm.model_executor.layers.vocab_parallel_embedding import ( DEFAULT_VOCAB_PADDING_SIZE, ParallelLMHead, VocabParallelEmbedding) -from vllm.model_executor.model_loader.weight_utils import default_weight_loader +from vllm.model_executor.model_loader.weight_utils import ( + default_weight_loader, maybe_remap_kv_scale_name) from vllm.model_executor.sampling_metadata import SamplingMetadata from vllm.sequence import IntermediateTensors @@ -345,6 +346,10 @@ def load_weights(self, weights: Iterable[Tuple[str, weight_loader(param, loaded_weight, shard_id) break else: + name = maybe_remap_kv_scale_name(name, params_dict) + if name is None: + continue + if self.config.tie_word_embeddings and "lm_head.weight" in name: continue if is_pp_missing_parameter(name, self):