From cd41deecbc902c9157321d3e4ef5a662afff4095 Mon Sep 17 00:00:00 2001 From: Isotr0py <2037008807@qq.com> Date: Mon, 27 Jan 2025 15:20:05 +0800 Subject: [PATCH] remove lm_head ignorance Signed-off-by: Isotr0py <2037008807@qq.com> --- vllm/model_executor/models/gpt2.py | 13 ++++--------- 1 file changed, 4 insertions(+), 9 deletions(-) diff --git a/vllm/model_executor/models/gpt2.py b/vllm/model_executor/models/gpt2.py index c3349f43aae97..2f1aa2d68653c 100644 --- a/vllm/model_executor/models/gpt2.py +++ b/vllm/model_executor/models/gpt2.py @@ -258,14 +258,12 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.transformer = GPT2Model(vllm_config=vllm_config, prefix=maybe_prefix( prefix, "transformer")) - - # GPT-2 ties the weights of the embedding layer and the final - # linear layer. self.lm_head = ParallelLMHead(self.config.vocab_size, self.config.hidden_size, quant_config=quant_config, prefix=f"{prefix}.lm_head") - self.lm_head = self.lm_head.tie_weights(self.transformer.wte) + if self.config.tie_word_embeddings: + self.lm_head = self.lm_head.tie_weights(self.transformer.wte) self.logits_processor = LogitsProcessor(config.vocab_size) self.sampler = get_sampler() @@ -311,15 +309,12 @@ def load_weights(self, weights: Iterable[Tuple[str, params_dict = dict(self.named_parameters(remove_duplicate=False)) loaded_params: Set[str] = set() for name, loaded_weight in weights: - if name.startswith("lm_head"): - # GPT-2 ties the weights of the embedding layer and the final - # linear layer. - continue if ".attn.bias" in name or ".attn.masked_bias" in name: # Skip attention mask. # NOTE: "c_attn.bias" should not be skipped. continue - if not name.startswith("transformer."): + if not name.startswith("transformer.") and not name.startswith( + "lm_head"): name = "transformer." + name if is_pp_missing_parameter(name, self):