Repository: vllm-project/vllm

Issue:
KeyError: lm_head.weight in GemmaForCausalLM.load_weights when loading finetuned Gemma 2B
Hello,

I finetuned Gemma 2B with [Unsloth](https://github.com/unslothai/unsloth). It uses LoRA and merges the weights back into the base model.

When I try to load this model, it gives me the following error:

```
...
File "/home/ubuntu/projects/cql-ml/.venv/lib/python3.10/site-packages/vlm/model_executor/model_loader.py", line 86, in get _model model. load weights(model_config.model, model_config.download_
config. model, model_ config. download dir,
File "/home/ubuntu/projects/cql-ml/.venv/lib/python3.10/site-packages/vlm/model_executor/models/gemma.py", line 339, in load weights
param = params_dict [name]
KeyError: 'lm_head.weight'
```

My `pytorch_model.bin.index.json` looks like this:

```
{
  "metadata": {
    "total_size": 6060920832
  },
  "weight_map": {
    "lm_head.weight": "pytorch_model-00002-of-00002.bin",
    "model.embed_tokens.weight": "pytorch_model-00001-of-00002.bin",
    "model.layers.0.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
    "model.layers.0.mlp.down_proj.weight": "pytorch_model-00001-of-00002.bin",
    "model.layers.0.mlp.gate_proj.weight": "pytorch_model-00001-of-00002.bin",
...
```

I saw in a few of the other classes a similar check for `lm_head.weight` so I replicated it in `load_weights` and the model loads correctly and works as intended. [1](https://github.com/vllm-project/vllm/blob/4c922709b65ff5c0652ae36b93047016bdeaace8/vllm/model_executor/models/bloom.py#L306) [2]([vllm/model_executor/models/gpt_bigcode.py](https://github.com/vllm-project/vllm/blob/4c922709b65ff5c0652ae36b93047016bdeaace8/vllm/model_executor/models/gpt_bigcode.py#L270)) [3](https://github.com/vllm-project/vllm/blob/4c922709b65ff5c0652ae36b93047016bdeaace8/vllm/model_executor/models/opt.py#L331)

The modified load_weights function:

https://github.com/vllm-project/vllm/commit/13333220e5e37c1a5e96e5ec879841d6f3774344

I'm not sure if this is an issue with vllm, or an issue with the output of Unsloth. The model works correctly when `load_weights` is modified. I don't know what the internals of the model should look like. Any help would be appreciated!


I'm unsure if this is related to https://github.com/vllm-project/vllm/issues/2816

My model is Private, so unfortunately I can't share it. However I found [this other model](https://huggingface.co/Aabbhishekk/gemma-2b-coder-unsloth-merged/tree/main) on huggingface that's trained with the same tool with the `lm_head.weight` in the index. 

If the modified `load_weights` function is the desired fix, I can submit a PR if that will help.

Thank you for the help!


Discussion:
Same error occurs when using quantized gemma-7b by autoawq.

After fine-tuning gemma-7b, the same error occurs.

It should be relevant to this PR: [3050](https://github.com/vllm-project/vllm/pull/3050).

I'm not sure if directly add the deleted two lines will break the loading of lora weights.

@WoosukKwon 
fine-tuning gemma-7b, the same error occurs. +1
Same here I used axolotl and LoRa to finetune gemma-7b. +1.
vllm 0.3.3

Resolving patch:
diff --git a/vllm/model_executor/models/gemma.py b/vllm/model_executor/models/gemma.py
--- a/vllm/model_executor/models/gemma.py
+++ b/vllm/model_executor/models/gemma.py
@@ -340,6 +340,10 @@ def load_weights(self,
                 weight_loader(param, loaded_weight, shard_id)
                 break
             else:
+                # lm_head is not used in vllm as it is tied with embed_token.
+                # To prevent errors, skip loading lm_head.weight.
+                if "lm_head.weight" in name:
+                    continue
                 # Skip loading extra bias for GPTQ models.
                 if name.endswith(".bias") and name not in params_dict:
                     continue


Test patch: