gpt-j input shape after finetuning
Author: calderma
Pull request: no

Hello after finetuning the EleutherAI/gpt-j-6B model and loading one of the checkpoints I'm getting the following error:

size mismatch for transformer.wte.weight: copying a param with shape torch.Size([50257, 4096]) from checkpoint, the shape in current model is torch.Size([50400, 4096])

Based on the GPT-J description there are some extra tokens beyond what the base GPT2 tokenizer has. Running the normal EleutherAI/gpt-j-6B model and tokenizer works fine, so perhaps the config that's being saved with the checkpoint is not accounting for those extra tokens? I pulled master this morning so I believe this is still the case with the current codebase. Training script is:
deepspeed run_clm.py \
    --model_name_or_path EleutherAI/gpt-j-6B \
    --train_file $TRAIN_FILE \
    --validation_file $VAL_FILE \
    --gradient_accumulation_steps 8 \
    --num_train_epochs 2.0 \
    --evaluation_strategy epoch \
    --save_steps XXX \
    --logging_steps XXX \
    --do_train \
    --do_eval \
    --save_total_limit 3 \
    --per_device_train_batch_size 1 \
    --per_device_eval_batch_size 1 \
    --deepspeed "deepspeed_zero3_config.json" \
    --overwrite_output_dir \
    --output_dir gpt-j_model

test script breaks at model loading:

from transformers import GPTJForCausalLM, AutoTokenizer
import torch
from pytictoc import TicToc

t = TicToc()

device = "cuda:0" if torch.cuda.is_available() else "cpu"

model = GPTJForCausalLM.from_pretrained("gpt-j_end2end/checkpoint-XXX", torch_dtype=torch.float16,max_length=2048).to(device)
print("model loaded")
...

Any help would be appreciated. Thanks!



Comments:
Hi @calderma 

Thank you for opening the issue. The actual vocab size is 50257  the remaining are extra tokens. 
The issue is that the vocab size is `50400` in the `config` file, but the tokenizer size is 50257, and the `run_clm.py` script, resizes  token embeddings here 
https://github.com/huggingface/transformers/blob/421929b556aedf022a1c4a1f3b2f116b14a7b88a/examples/pytorch/language-modeling/run_clm.py#L358

which reduces the model's embedding size to 50257, hence the shape mismatch.

@LysandreJik  should we remove the extra tokens in the embedding and update the model on the hub?

Ok is there a way around this issue by manually changing the config.json or something or should I just wait until this gets updated on the hub?

Yes, if you have already fine-tuned a model, then you should change the `vocab_size` in `config.json` to 50257.

The fix is on the way #13617 


@patil-suraj but in fine-tuned model config.json file vocab_size is already 50257 :) 

yes I noticed that as well. I got it to run by doing:

from transformers import GPTNeoForCausalLM, AutoTokenizer, GPT2Tokenizer, GPTJForCausalLM
import torch
from pytictoc import TicToc
t = TicToc()

device = "cuda:0" if torch.cuda.is_available() else "cpu"

model = GPTJForCausalLM.from_pretrained("gpt-j_model",ignore_mismatched_sizes=True,torch_dtype=torch.float16).to(device)

print("model loaded")

tokenizer = AutoTokenizer.from_pretrained("gpt-j_model",truncation=True)
model.resize_token_embeddings(len(tokenizer))

but then it output nonesense so I don't think the tokenization is being handled correctly when I do that. I also tried passing the base model as the tokenizer and the gpt2 tokenizer but they all output nosense.

> @patil-suraj but in fine-tuned model config.json file vocab_size is already 50257 :)

If that is the case, then the model should load without any issues, as the shapes would match.

> yes I noticed that as well. I got it to run by doing:
> 
> from transformers import GPTNeoForCausalLM, AutoTokenizer, GPT2Tokenizer, GPTJForCausalLM
> import torch
> from pytictoc import TicToc
> t = TicToc()
> 
> device = "cuda:0" if torch.cuda.is_available() else "cpu"
> 
> model = GPTJForCausalLM.from_pretrained("gpt-j_model",ignore_mismatched_sizes=True,torch_dtype=torch.float16).to(device)
> 
> print("model loaded")
> 
> tokenizer = AutoTokenizer.from_pretrained("gpt-j_model",truncation=True)
> model.resize_token_embeddings(len(tokenizer))
> 
> but then it output nonesense so I don't think the tokenization is being handled correctly when I do that. I also tried passing the base model as the tokenizer and the gpt2 tokenizer but they all output nosense.

This is because `ignore_mismatched_sizes` argument, when this argument is `True` it will ignore the weights from the checkpoint that do not have the same shape as the ones inside the model and leave the randomly initialized weights. If you have already fine-tuned a model, then the simplest solution, for now, would be to change the `vocab_size` in `config.json` manually.

@calderma How exactly are you finetuning? Can you share a complete example somehow?


@calderma Can you share what machine spec are you using for finetuning? It would be really helpful.

@working12 the finetuning I'm doing is just what I pasted in the issue. run_clm.py with the arguments I listed. For machine specs if you are asking about what kind of GPUs I'm using etc. I'm using a machine with 8x 32g GPUs. In order to fit it and train I had to use deepspeed either zero2 or zero3 and then could only do batch size 1.

@patil-suraj what's the status of this issue? Does your fix work for fine-tune command:

deepspeed --master_port 29603 --include localhost:3 run_clm.py \
--deepspeed zero3.json \
--model_revision float32 \
--model_name_or_path EleutherAI/gpt-j-6B \
--do_train --train_file train.txt \
--do_eval --validation_file test.txt \
--evaluation_strategy no \
--logging_strategy steps \
--logging_steps $logging_steps \
--save_strategy steps \
--save_steps $save_steps \
--overwrite_output_dir \
--output_dir output.$EXP \
--per_device_train_batch_size $BS \
--per_device_eval_batch_size $BS \
--gradient_accumulation_steps $AS \
--num_train_epochs $N > log.$EXP

I am blocked by the same issue when NOT use --fp16. I tried two fine-tune with the same parameters but

(1) one with "--model_revision float16 --fp16"
(2) the other with the command above, which uses "--model_revision float32" but no --fp16.

The output model from (1) can be loaded fine but (2) triggered size mismatch error. How do I fix the error? --fp16 fine tuning has so many OVERFLOW errors that I wanted to turn --fp16 off. 

A bit more details.

config.json-->vocab_size: 50400 in (1) but 50257 in (2). I tried to manually change 50257 to 50400 in (2) but triggered other errors.
tokenizer.json: many 
 in (1) but not in (2)

@patil-suraj @calderma @MantasLukauskas The issue is actually more fundamental. There seems to be a problem with the serialization of GPTJ. Even without fine-tuning, this issue is encountered if the embedding size of GPTJ is expanded. For example, if you run this code below you will also get a size mismatch error:

```
from transformers import GPTJForCausalLM
model = GPTJForCausalLM.from_pretrained("EleutherAI/gpt-j-6B")
model.resize_token_embeddings(54001)
model.save_pretrained("model_dir")
model = GPTJForCausalLM.from_pretrained("model_dir")
RuntimeError: Error(s) in loading state_dict for GPTJForCausalLM:
	size mismatch for lm_head.weight: copying a param with shape torch.Size([50400, 4096]) from checkpoint, the shape in current model is torch.Size([54001, 4096]).
	size mismatch for lm_head.bias: copying a param with shape torch.Size([50400]) from checkpoint, the shape in current model is torch.Size([54001]).
```

I fixed this one by using "--model_revision main" instead of "--model_revision float32". It looks the previous PR fixed only "main" and "float16" but not for "float32".

> @patil-suraj what's the status of this issue? Does your fix work for fine-tune command:
> 
> deepspeed --master_port 29603 --include localhost:3 run_clm.py --deepspeed zero3.json --model_revision float32 --model_name_or_path EleutherAI/gpt-j-6B --do_train --train_file train.txt --do_eval --validation_file test.txt --evaluation_strategy no --logging_strategy steps --logging_steps $logging_steps --save_strategy steps --save_steps $save_steps --overwrite_output_dir --output_dir output.$EXP --per_device_train_batch_size $BS --per_device_eval_batch_size $BS --gradient_accumulation_steps $AS --num_train_epochs $N > log.$EXP
> 
> I am blocked by the same issue when NOT use --fp16. I tried two fine-tune with the same parameters but
> 
> (1) one with "--model_revision float16 --fp16" (2) the other with the command above, which uses "--model_revision float32" but no --fp16.
> 
> The output model from (1) can be loaded fine but (2) triggered size mismatch error. How do I fix the error? --fp16 fine tuning has so many OVERFLOW errors that I wanted to turn --fp16 off.
> 
> A bit more details.
> 
> config.json-->vocab_size: 50400 in (1) but 50257 in (2). I tried to manually change 50257 to 50400 in (2) but triggered other errors. tokenizer.json: many 
 in (1) but not in (2)



@dunalduck0 Oh great I'll try that.

@dunalduck0  the `main` branch has the fp32 weights which should be used instead of the `float32` branch.
Also, there was another issue with resizing embeddings, #14190 should fix that.

Thank you @patil-suraj 

This is fixed by #14190

@alexorona I've got the same problem as you. After resizing, I cannot load the original model anymore due to size mismatch. I found this is due to resize_token_embeddings.() changing config.vocal_size. If I reset config.vocab_size to 54000 before loading, it works fine.

> @patil-suraj @calderma @MantasLukauskas The issue is actually more fundamental. There seems to be a problem with the serialization of GPTJ. Even without fine-tuning, this issue is encountered if the embedding size of GPTJ is expanded. For example, if you run this code below you will also get a size mismatch error:
> 
> ```
> from transformers import GPTJForCausalLM
> model = GPTJForCausalLM.from_pretrained("EleutherAI/gpt-j-6B")
> model.resize_token_embeddings(54001)
> model.save_pretrained("model_dir")
> model = GPTJForCausalLM.from_pretrained("model_dir")
> RuntimeError: Error(s) in loading state_dict for GPTJForCausalLM:
> 	size mismatch for lm_head.weight: copying a param with shape torch.Size([50400, 4096]) from checkpoint, the shape in current model is torch.Size([54001, 4096]).
> 	size mismatch for lm_head.bias: copying a param with shape torch.Size([50400]) from checkpoint, the shape in current model is torch.Size([54001]).
> ```