size mismatch for lm_head.weight: copying a param with shape torch.Size([32, 768]) from checkpoint, the shape in current model is torch.Size([59744, 768]).
Author: SoSongzhi
Pull request: no
Labels: name: bug

### System Info

transformers version: 4.46.3
Platform: Linux-6.8.0-49-generic-x86_64-with-glibc2.17
Python version: 3.8.20
Huggingface_hub version: 0.26.2
Safetensors version: 0.4.5
Accelerate version: 1.0.1
Accelerate config: not found
PyTorch version (GPU?): 2.4.1+cu121 (True)
Tensorflow version (GPU?): not installed (NA)
Flax version (CPU?/GPU?/TPU?): not installed (NA)
Jax version: not installed
JaxLib version: not installed
Using distributed or parallel set-up in script?:
Using GPU in script?:
GPU type: NVIDIA RTX A4000



### Who can help?

@ArthurZucker 

Hi, I am using T5 1.1 base to train on my seq2seq task, and the source vocab size is 59744 and the target vocab size is only 32. I know by default the headLM's size and decoder embedding size are equal to the vocab's size as 59k so I change decoder embedding also headLM as follows when trainning and I found it all appear no error through training:

model = T5ForConditionalGeneration.from_pretrained("google/t5-v1_1-base")
model.resize_token_embeddings(new_vocab_size)

target_vocab_size = 32
model.decoder.embed_tokens = torch.nn.Embedding(target_vocab_size, model.config.d_model)

model.lm_head = torch.nn.Linear(model.config.d_model, target_vocab_size, bias=False)
torch.nn.init.normal_(model.lm_head.weight, mean=0.0, std=model.config.initializer_factor)



but when i try to load the checkpoint to predict as follows:

config = T5Config.from_pretrained("./resultstest/checkpoint-100")

    # model = T5ForConditionalGeneration(config)
    model = AutoModelForSeq2SeqLM.from_pretrained(checkpoint_path, config=config)


it shows like that:

size mismatch for decoder.embed_tokens.weight: copying a param with shape torch.Size([32, 768]) from checkpoint, the shape in current model is torch.Size([59744, 768]).
	size mismatch for lm_head.weight: copying a param with shape torch.Size([32, 768]) from checkpoint, the shape in current model is torch.Size([59744, 768]).


So how can I fix it to achieve the decoder embedding and headLM's size is only 32 to fit my target. because I need to calculate the probability of each token using softmax on 32 candidates rather than 59k. 

It would be very appreciate if someone could help me to fix it. 


Best.


### 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

this is my trainning code:

'''python

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
import torch
from transformers import AutoModelForSeq2SeqLM, Trainer, TrainingArguments, AutoTokenizer, T5Config, T5ForConditionalGeneration
from datasets import Dataset
import pandas as pd
from sklearn.model_selection import train_test_split
from tokenizer import SourceFuzzyTokenizer as CustomTokenizer 
vocab_file = "./datasets/tokenizer/source_vocab.json"  

data_file_fuzzy = "./datasets/test_data/testfuzzy.txt"
data_file_gt = "./datasets/test_data/testgt.txt"

    

max_length = 50                   
batch_size = 50                
num_epochs = 100                     
learning_rate = 5e-4 
new_vocab_size = 59744
output_dir = "./resultstest"  


with open(data_file_fuzzy, "r") as fuzzy_file, open(data_file_gt, "r") as gt_file:
    fuzzy_seqs = fuzzy_file.read().splitlines()
    gt_seqs = gt_file.read().splitlines()


assert len(fuzzy_seqs) == len(gt_seqs), "fuzzy_seqs.txt and gt_seqs.txt do not MATCH!"

data = {"input": fuzzy_seqs, "target": gt_seqs}
df = pd.DataFrame(data)


train_df, eval_df = train_test_split(df, test_size=0.1, random_state=42)

train_dataset = Dataset.from_pandas(train_df)
eval_dataset = Dataset.from_pandas(eval_df)


tokenizer = CustomTokenizer(vocab_file)

model = T5ForConditionalGeneration.from_pretrained("google/t5-v1_1-base")
model.resize_token_embeddings(new_vocab_size)

target_vocab_size = 32
model.decoder.embed_tokens = torch.nn.Embedding(target_vocab_size, model.config.d_model)

model.lm_head = torch.nn.Linear(model.config.d_model, target_vocab_size, bias=False)
torch.nn.init.normal_(model.lm_head.weight, mean=0.0, std=model.config.initializer_factor)

model.config.vocab_size = 59744 
model.config.decoder_vocab_size = target_vocab_size 

model.lm_head.weight.data.normal_(mean=0.0, std=model.config.initializer_factor)

def preprocess_data(examples):
    inputs = [tokenizer.tokenize(seq) for seq in examples["input"]]
    targets = [tokenizer.tokenize(seq) for seq in examples["target"]]


    input_ids = [tokenizer.convert_tokens_to_ids(tokens)[:max_length] for tokens in inputs]
    target_ids = [tokenizer.convert_tokens_to_ids(tokens)[:max_length] for tokens in targets]


    pad_id = tokenizer.vocab.get("[PAD]", 0)
    input_ids = [seq + [pad_id] * (max_length - len(seq)) for seq in input_ids]
    target_ids = [seq + [pad_id] * (max_length - len(seq)) for seq in target_ids]

    attention_mask = [[1 if token != pad_id else 0 for token in seq] for seq in input_ids]

    return {
        "input_ids": input_ids,"attention_mask": attention_mask,
        "labels": target_ids,
    }

train_dataset = train_dataset.map(preprocess_data, batched=True)
eval_dataset = eval_dataset.map(preprocess_data, batched=True)

training_args = TrainingArguments(
    output_dir=output_dir,
    eval_strategy="steps",
    learning_rate=learning_rate,
    per_device_train_batch_size=batch_size,
    per_device_eval_batch_size=batch_size,
    num_train_epochs=num_epochs,
    weight_decay=0.01,
    save_strategy="steps",
    save_total_limit=2,
    logging_dir="./logs",
    logging_steps=10,
    evaluation_strategy="epoch",
    save_steps = 50,
    # load_best_model_at_end=True,
    dataloader_num_workers = 1,
    fp16=False,
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
)


trainer.train()
'''

### Expected behavior

I expect that the headLM and decoder embedding layer can fit my source dictionary's size 32, rather than calculate logits witch have 59744 size.

Comments:
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.

Hey! Not sure what happened, hope you found a fix!