Skip to content

Commit

Permalink
[Bugfix] Update starcoder2 to remap k/v scale names for kv_cache quan…
Browse files Browse the repository at this point in the history
…tization (vllm-project#11148)
  • Loading branch information
dsikka authored Dec 13, 2024
1 parent 1efce68 commit 3989a79
Showing 1 changed file with 6 additions and 1 deletion.
7 changes: 6 additions & 1 deletion vllm/model_executor/models/starcoder2.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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):
Expand Down

0 comments on commit 3989a79

Please sign in to comment.