issues when i change the lm_head to a 32 node layer
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?: <fill in>
- Using GPU in script?: <fill in>
- GPU type: NVIDIA RTX A4000


### 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, ...)
- [X] My own task or dataset (give details below)

### Reproduction

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
import torch
from transformers import AutoModelForSeq2SeqLM, Trainer, TrainingArguments, T5Config
from datasets import Dataset
import pandas as pd
from sklearn.model_selection import train_test_split
from tokenizer import SourceFuzzyTokenizer as SourceTokenizer, TargetSolidTokenizer as TargetTokenizer

# 路径和超参数
source_vocab_file = "./datasets/tokenizer/source_vocab.json"  # 输入词表文件路径
target_vocab_file = "./datasets/tokenizer/target_vocab.json"  # 输出词表文件路径
data_file_fuzzy = "/home/zhi/Desktop/fuzzy_complement/datasets/raw/9specie_fuzzy_seqs.txt"  # 输入数据文件
data_file_gt = "/home/zhi/Desktop/fuzzy_complement/datasets/raw/9specie_gt_seqs.txt"       # 目标数据文件
model_name = "./checkpoint-1817245"                                             # 模型路径
max_length = 50
batch_size = 80
num_epochs = 100
learning_rate = 8e-5
source_vocab_size = 59746
target_vocab_size = 32
output_dir = "./results"


# train from stratch
config = T5Config(
    vocab_size=source_vocab_size,  # 输入词汇表大小
    decoder_start_token_id=0,
    eos_token_id=1,
    pad_token_id=0,
    d_model=512,  # 模型维度
    num_decoder_layers=6,
    num_encoder_layers=6,
)

# 加载数据
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)

# splite Tokenizer
source_tokenizer = SourceTokenizer(source_vocab_file)
target_tokenizer = TargetTokenizer(target_vocab_file)

# use pretrained model
# model = AutoModelForSeq2SeqLM.from_pretrained(model_name, ignore_mismatched_sizes=True)
# use from_stratch model
model = AutoModelForSeq2SeqLM.from_config(config)
# model = AutoModelForSeq2SeqLM.from_pretrained("t5-base")

# change model's embedding layer
model.encoder.embed_tokens = torch.nn.Embedding(source_vocab_size, model.config.d_model)
model.decoder.embed_tokens = torch.nn.Embedding(target_vocab_size, model.config.d_model)

# init the weight of model
model.encoder.embed_tokens.weight.data.normal_(mean=0.0, std=model.config.initializer_factor)
model.decoder.embed_tokens.weight.data.normal_(mean=0.0, std=model.config.initializer_factor)

model.lm_head = torch.nn.Linear(model.config.d_model, target_vocab_size, bias=False)
model.lm_head.weight = model.decoder.embed_tokens.weight

model.config.vocab_size = source_vocab_size
model.config.decoder_vocab_size = target_vocab_size

def preprocess_data(examples):
    inputs = [source_tokenizer.tokenize(seq) for seq in examples["input"]]
    targets = [target_tokenizer.tokenize(seq) for seq in examples["target"]]

    input_ids = [source_tokenizer.convert_tokens_to_ids(tokens)[:max_length] for tokens in inputs]
    target_ids = [target_tokenizer.convert_tokens_to_ids(tokens)[:max_length] for tokens in targets]

    pad_id_source = source_tokenizer.vocab.get("[PAD]", 0)
    pad_id_target = target_tokenizer.vocab.get("[PAD]", 0)
    input_ids = [seq + [pad_id_source] * (max_length - len(seq)) for seq in input_ids]
    target_ids = [seq + [pad_id_target] * (max_length - len(seq)) for seq in target_ids]

 
    attention_mask = [[1 if token != pad_id_source 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="epoch",
    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="epoch",
    save_total_limit=2,
    logging_dir="./logs",
    logging_steps=10,
    evaluation_strategy="epoch",
    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

My task is to translate a fuzzy sequence into a solid sequence, which involves peptide sequences. When I used T5-base without distinguishing between the source and target dictionaries, everything worked fine. However, I noticed that the solid sequence contains only 32 unique characters. To optimize, I separated the source and target character sets and adjusted the lm_head to match the 32 characters. Training worked fine, and the model converged. But during inference, I encountered an issue where the weight sizes do not match, even though I had adjusted the size during training. What could be causing this, and how can I resolve it?

here is the error message:
python bs_based_on_massdic.py 
Traceback (most recent call last):
  File "bs_based_on_massdic.py", line 100, in <module>
    model = AutoModelForSeq2SeqLM.from_pretrained(checkpoint_path, ignore_mismatched_sizes=True)
  File "/home/zhi/anaconda3/envs/peptide_completion/lib/python3.8/site-packages/transformers/models/auto/auto_factory.py", line 564, in from_pretrained
    return model_class.from_pretrained(
  File "/home/zhi/anaconda3/envs/peptide_completion/lib/python3.8/site-packages/transformers/modeling_utils.py", line 4225, in from_pretrained
    ) = cls._load_pretrained_model(
  File "/home/zhi/anaconda3/envs/peptide_completion/lib/python3.8/site-packages/transformers/modeling_utils.py", line 4785, in _load_pretrained_model
    raise RuntimeError(f"Error(s) in loading state_dict for {model.__class__.__name__}:\n\t{error_msg}")
RuntimeError: Error(s) in loading state_dict for T5ForConditionalGeneration:
	size mismatch for shared.weight: copying a param with shape torch.Size([59750, 512]) from checkpoint, the shape in current model is torch.Size([59746, 512]).
	size mismatch for decoder.embed_tokens.weight: copying a param with shape torch.Size([32, 512]) from checkpoint, the shape in current model is torch.Size([59746, 512]).
	You may consider adding `ignore_mismatched_sizes=True` in the model `from_pretrained` method.

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.