T5ForConditionalGeneration checkpoint size mismatch
Author: msamogh
Pull request: no

### System Info

## Error Description
I trained a `T5ForConditionalGeneration` model and saved the checkpoint using PyTorch Lightning's Trainer to a `.ckpt` file. But when I try to load back the state_dict using `model.from_state_dict()`, I get this error:
```python
RuntimeError: Error(s) in loading state_dict for T5ForConditionalGeneration:
        Unexpected key(s) in state_dict: "decoder.block.0.layer.1.EncDecAttention.relative_attention_bias.weight". 
        size mismatch for shared.weight: copying a param with shape torch.Size([32103, 512]) from checkpoint, the shape in current model is torch.Size([32128, 512]).
        size mismatch for encoder.embed_tokens.weight: copying a param with shape torch.Size([32103, 512]) from checkpoint, the shape in current model is torch.Size([32128, 512]).
        size mismatch for decoder.embed_tokens.weight: copying a param with shape torch.Size([32103, 512]) from checkpoint, the shape in current model is torch.Size([32128, 512]).
        size mismatch for lm_head.weight: copying a param with shape torch.Size([32103, 512]) from checkpoint, the shape in current model is torch.Size([32128, 512]).
```

I have not changed the model definition in any way. The keys also match. So, I'm really not sure how the sizes could mismatch magically when loading?

## Loading the model
This is how I'm loading the model:
```python
tokenizer = T5Tokenizer.from_pretrained(args["model_checkpoint"], bos_token="[bos]", eos_token="[eos]", sep_token="[sep]")
    model = T5ForConditionalGeneration.from_pretrained(args["model_checkpoint"], ignore_mismatched_sizes=True)
    model.load_state_dict({k[6:]: v for k, v in ckpt["state_dict"].items()})
```

I even tried to pass `ignore_mismatched_sizes=True` to the `from_pretrained` call, and that didn't help either.

### Who can help?

_No response_

### Information

- [ ] The official example scripts
- [X] 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

As described above.

### Expected behavior

No error

Comments:
Hey @msamogh 👋 

To explain why there is a mismatch, we would need to know exactly how the model was trained :) However, the most important part -- you may be able to load the checkpoint with these two strategies:
1. Load the model architecture from the same configuration as your trained model
2. After initializing the model architecture (and before loading the checkpoint), [resize the embeddings](https://huggingface.co/docs/transformers/v4.22.2/en/main_classes/model#transformers.PreTrainedModel.resize_token_embeddings)

Both strategies should change the shape of your architecture to match your checkpoint

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.

Hi @gante, does interpolate_pos_embedding type functions do what you mentioned in point 2?

Here is a snippet:

```
def interpolate_pos_embed_multimae(model, checkpoint_model):
    pattern = "input_adapters\.(.*)\.pos_emb"
    matched_keys = [k for k in checkpoint_model if bool(re.match(pattern, k))]

    for key in matched_keys:
        domain = re.match(pattern, key).group(1)  # group(0) is entire matched regex
        if getattr(model.input_adapters, domain, None) is not None:
            pos_embed_checkpoint = checkpoint_model[key]
            _, _, orig_H, orig_W = pos_embed_checkpoint.shape
            _, _, new_H, new_W = getattr(model.input_adapters, domain).pos_emb.shape
            if (orig_H != new_H) or (orig_W != new_W):
                print(f"Key {key}: Position interpolate from {orig_H}x{orig_W} to {new_H}x{new_W}")
                pos_embed_checkpoint = torch.nn.functional.interpolate(
                    pos_embed_checkpoint, size=(new_H, new_W), mode='bicubic', align_corners=False)
                checkpoint_model[key] = pos_embed_checkpoint

```



Hey @forkbabu 👋 I do not know the answer to your question. However, from your code snippet, it seems like you are working with a vision model -- my recommendation would be to open a new issue and tag one of our vision experts 

Usually you don't encounter any problems when loading the model for which you've added some extra tokens during the training. In my case, it was the `pad_to_multiple_of` parameter that caused the trouble. It is claimed to do some Nvidia magic for a more efficient utilization of modern GPUs, so I used it when I created the model for training and then happily forgot about it: 
```
model.resize_token_embeddings(len(tokenizer), pad_to_multiple_of=16)
```

But as it seems, the current API (4.33.0.dev0) struggles to load such models. The workaround would be:
```
MODEL_CHECKPOINT = ''  # your directory here
config_path = path.join(MODEL_CHECKPOINT, 'config.json')
weights_path = path.join(MODEL_CHECKPOINT, 'pytorch_model.bin')

tokenizer = GPT2Tokenizer.from_pretrained(MODEL_CHECKPOINT)

config = AutoConfig.from_pretrained(config_path)
model = T5ForConditionalGeneration(config)  
model.resize_token_embeddings(len(tokenizer), pad_to_multiple_of=16)
model.load_state_dict(torch.load(weights_path, map_location=torch.device('cpu')))
```
Which outputs: \<All keys matched successfully>

Yep, the vocab size is missing some changes, fix is here: #25732

@gante , the fix that @ArthurZucker started in #25732  looks like a good one, and it addresses this error that I bet more people are going to run into, now that there is a warning when training that encourages us to pad to a multiple of 8.

Here are some helper functions I wrote to correctly resize embeddings for a model (before this bug is fixed), and repair broken models that were saved with this incorrect vocab size. (I might have typed things too narrowly without realizing it.)

This is tiding me over until #25732 or an equivalent is merged.

```python
def resize_model_token_embeddings_correctly(model: PreTrainedModel, 
                                            new_num_tokens: Optional[int] = None, 
                                            pad_to_multiple_of: Optional[int] = None)-> nn.Embedding:
    """This is a workaround for the bug: https://github.com/huggingface/transformers/issues/25729 that doesn't save embedding vocab sizes correctly"""
    vocab_size_before:int = model.get_input_embeddings().weight.shape[0]
    model_embeds: nn.Embedding = model.resize_token_embeddings(new_num_tokens, pad_to_multiple_of=pad_to_multiple_of)
    vocab_size_after:int = model.get_input_embeddings().weight.shape[0]
    # This appears to be the correct way to resize the embeddings since the pad_to_multiple_of is broken in the current version of Transformers:
    # This has been fixed in this PR: https://github.com/huggingface/transformers/pull/25732/files#diff-6b72b98c4c2dcfc6cc606843917733f5d858374fbc22a735ff483bbc0c1e63ea
    # But in the meantime, we do this part manually: (Most of this is copied from the resize_token_embeddings method or that PR)

    # You might think we need to early-out in case they didn't pass new_num_tokens, but the number of tokens might _still_ change
    # if they passed a pad_to_multiple_of in, so we do this to be safe for now.
    if vocab_size_before == vocab_size_after:
        return model_embeds
    
    model.config.vocab_size = vocab_size_after
    model.vocab_size = vocab_size_after
    model.tie_weights()
    return model_embeds

def fix_model_folder_with_incorrect_vocab_size(model_folder: Path, pad_vocab_size_to_multiple_of: int=64):
    """If you got bitten by the bug https://github.com/huggingface/transformers/issues/25729, and your model got saved with an incorrect vocab size, this will fix it."""
    config_file = model_folder / "config.json"
    config = AutoConfig.from_pretrained(config_file)
    vocab_size:int = config.vocab_size
    # is vocab_size an int, and is it bigger than 0?
    assert isinstance(vocab_size, int) and vocab_size > 0
        
    if vocab_size % pad_vocab_size_to_multiple_of != 0:
        print(f"vocab_size ({vocab_size}) is not a multiple of {pad_vocab_size_to_multiple_of}. Fixing it...")
        # fix the vocab_size:
        config.vocab_size = vocab_size + (pad_vocab_size_to_multiple_of - vocab_size % pad_vocab_size_to_multiple_of)
        # save the config file back to the checkpoint folder:
        config.save_pretrained(model_folder)
        print(f"Fixed the vocab_size to {config.vocab_size}")

def load_model_with_missized_vocab_size(model_folder: Path, pad_vocab_size_to_multiple_of: int=64) -> PreTrainedModel:
    """There is a bug in how a model was getting saved, where the vocab size was not set to a multiple of a given padding, but it should have been."""
    fix_model_folder_with_incorrect_vocab_size(model_folder, pad_vocab_size_to_multiple_of)
    model:PreTrainedModel = AutoModelForSeq2SeqLM.from_pretrained(model_folder)
    return model
    ```