Can not load model after finetuning PHI2 model
Author: zhangmiaosen2000
Pull request: no

### System Info

Here is my code in finetuning PHI-2 model
```
model = transformers.AutoModelForCausalLM.from_pretrained("microsoft/phi-2", torch_dtype="auto", flash_attn=True, flash_rotary=True, fused_dense=True, trust_remote_code=True)
tokenizer = transformers.AutoTokenizer.from_pretrained(
        "microsoft/phi-2",
        cache_dir=training_args.cache_dir,
        model_max_length=training_args.model_max_length,
        padding_side="right",
        use_fast=True,
        trust_remote_code=True
    )
if tokenizer.pad_token is None:
    smart_tokenizer_and_embedding_resize(
        special_tokens_dict=dict(pad_token=DEFAULT_PAD_TOKEN),
        tokenizer=tokenizer,
        model=model,
    )
```
where:
```
def smart_tokenizer_and_embedding_resize(
    special_tokens_dict: Dict,
    tokenizer: transformers.PreTrainedTokenizer,
    model: transformers.PreTrainedModel,
):
    """Resize tokenizer and embedding.

    Note: This is the unoptimized version that may make your embedding size not be divisible by 64.
    """
    num_new_tokens = tokenizer.add_special_tokens(special_tokens_dict)
    model.resize_token_embeddings(len(tokenizer))

    if num_new_tokens > 0:
        input_embeddings = model.get_input_embeddings().weight.data
        output_embeddings = model.get_output_embeddings().weight.data

        input_embeddings_avg = input_embeddings[:-num_new_tokens].mean(dim=0, keepdim=True)
        output_embeddings_avg = output_embeddings[:-num_new_tokens].mean(dim=0, keepdim=True)

        input_embeddings[-num_new_tokens:] = input_embeddings_avg
        output_embeddings[-num_new_tokens:] = output_embeddings_avg
```
Then I successfully finetuned model and save to xxx/checkpoint-500.
But when I try to use this code to load finetuned model:
```
model = AutoModelForCausalLM.from_pretrained(
                base_model,
                torch_dtype=torch.float16,
                device_map="auto",
                trust_remote_code = True
            )
```
I alway got this:
```
Traceback (most recent call last):                                                                                                                               
  File "xxx.py", line 272, in <module>                                                                                                                 
    main()                                                                                                                                                       
  File "xxx.py", line 151, in main                                                                                                                     
    tokenizer, model = get_model(base_model=args.model, page_attention=args.vllm)                                                                                
  File "xxx.py", line 80, in get_model                                                                                                                 
    model = AutoModelForCausalLM.from_pretrained(                                                                                                                
  File "/usr/local/lib/python3.8/dist-packages/transformers/models/auto/auto_factory.py", line 561, in from_pretrained                                           
    return model_class.from_pretrained(                                                                                                                          
  File "/usr/local/lib/python3.8/dist-packages/transformers/modeling_utils.py", line 3480, in from_pretrained
    ) = cls._load_pretrained_model(
  File "/usr/local/lib/python3.8/dist-packages/transformers/modeling_utils.py", line 3870, in _load_pretrained_model
    new_error_msgs, offload_index, state_dict_index = _load_state_dict_into_meta_model(
  File "/usr/local/lib/python3.8/dist-packages/transformers/modeling_utils.py", line 743, in _load_state_dict_into_meta_model
    set_module_tensor_to_device(model, param_name, param_device, **set_module_kwargs)
  File "/usr/local/lib/python3.8/dist-packages/accelerate/utils/modeling.py", line 285, in set_module_tensor_to_device
    raise ValueError(
ValueError: Trying to set a tensor of shape torch.Size([50296, 2560]) in "weight" (which has shape torch.Size([50304, 2560])), this look incorrect.
```

**Can you help me with that?**

Note that, the saved config.json is:
```
{
  "_name_or_path": "microsoft/phi-2",
  "activation_function": "gelu_new",
  "architectures": [
    "PhiForCausalLM"
  ],
  "attn_pdrop": 0.0,
  "auto_map": {
    "AutoConfig": "microsoft/phi-2--configuration_phi.PhiConfig",
    "AutoModelForCausalLM": "microsoft/phi-2--modeling_phi.PhiForCausalLM"
  },
  "embd_pdrop": 0.0,
  "flash_attn": true,
  "flash_rotary": true,
  "fused_dense": true,
  "img_processor": null,
  "initializer_range": 0.02,
  "layer_norm_epsilon": 1e-05,
  "model_type": "phi-msft",
  "n_embd": 2560,
  "n_head": 32,
  "n_head_kv": null,
  "n_inner": null,
  "n_layer": 32,
  "n_positions": 2048,
  "resid_pdrop": 0.1,
  "rotary_dim": 32,
  "tie_word_embeddings": false,
  "torch_dtype": "bfloat16",
  "transformers_version": "4.36.2",
  "use_cache": false,
  "vocab_size": 50296
}

```

### Who can help?

_No response_

### Information

- [ ] The official example scripts
- [ ] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

,

### Expected behavior

,

Comments:
Running into the exact same issue after fine-tuning with embedding reshaped 

> But when I try to use this code to load finetuned model:

I assume that the error appears when you try to merge the adapter and the base model. The problem is that phi-2 expects the vocab_size to be a multiple of 64 (see [configuration_phi.py](https://huggingface.co/microsoft/phi-2/blob/834565c23f9b28b96ccbeabe614dd906b6db551a/configuration_phi.py#L44)). 

Whenever you add new tokens, you have to pad the embedding dimension to a multiple of 64. This has to be done after adding tokens when training **and** before merging the trained adapter. Example:

```python
tokenizer.add_tokens(["
", "<PAD>"])
tokenizer.pad_token = "<PAD>"
tokenizer.add_special_tokens(dict(eos_token="
"))
model.resize_token_embeddings(
    new_num_tokens=len(tokenizer),
    pad_to_multiple_of=64)   # phi2 default is 64, see configuration_phi.py
```

[Notebook](https://github.com/geronimi73/phi2-finetune/blob/main/nb_qlora.ipynb) with a complete example for a QLoRA finetune of phi2.

Thanks @geronimi73 for the answer 🤗 

> > But when I try to use this code to load finetuned model:
> Whenever you add new tokens, you have to pad the embedding dimension to a multiple of 64. This has to be done after adding tokens when training **and** before merging the trained adapter. Example:

sorry, forget everything I wrote please. It works but it's not necessary. The model's embedding layer is big enough to accomodate additional tokens, see this: https://huggingface.co/microsoft/phi-2/discussions/22#659d8ba950c1bbee5be6f179

in other words, when you add tokens to the tokenizer, no need to resize the embeddings. in code:

```python
if tokenizer.pad_token is None:
  tokenizer.add_special_tokens(dict(pad_token=DEFAULT_PAD_TOKEN))
   # not needed:
   # smart_tokenizer_and_embedding_resize(
   #     special_tokens_dict=dict(pad_token=DEFAULT_PAD_TOKEN),
   #     tokenizer=tokenizer,
   #    model=model,
   #)

```

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.