[에러 해결] RuntimeError: Error(s) in loading state_dict

 

RuntimeError: Error(s) in loading state_dict for GPT2LMHeadModel:
size mismatch for transformer.wte.weight: copying a param with shape torch.Size([51200, 768]) from checkpoint, the shape in current model is torch.Size([51215, 768]).
size mismatch for lm_head.weight: copying a param with shape torch.Size([51200, 768]) from checkpoint, the shape in current model is torch.Size([51215, 768]).

 

나의 경우 사전학습된 koGPT2 모델을 불러와 사용하고, tokenizer에는 15개의 additional special tokens을 추가해서 사용하는 상황이었다.

tokenizer = AutoTokenizer.from_pretrained(
    model_name,
    additional_special_tokens=[
        # 15개의 special tokens 추가
    ],
)
model = GPT2LMHeadModel.from_pretrained(
    model_name,
    from_tf=False,
    config=config,
)

model.resize_token_embeddings(len(tokenizer))

위 코드와 같이 tokenizer와 model을 선언한 후 바로 resize_token_embeddings()을 수행한 후에 train() 함수를 실행하는 순서였다.

모델을 불러오는 과정에서 Huggingface로부터 from_pretrained('skt/kogpt-base-v2')로 불러올 때는 에러가 발생하지 않았는데, 모델을 다운로드 받아서 로컬에 저장해두고 from_pretrained('로컬 경로')로 불러오도록 하니 위와 같은 에러가 발생했다.

 

에러 발생 지점

정확히 어디서 에러가 나는지 보니 PTM의 state_dict를 불러오는 부분에서 발생하고 있었다.

 

size mismatch for transformer.wte.weight: copying a param with shape torch.Size([51200, 768]) from checkpoint, the shape in current model is torch.Size([51215, 768])

위 에러 메시지를 보면 [51200, 768]은 PTM에서의 shape이고, [51215, 768]은 내가 special token들을 추가하고 resize를 한 후의 shape이다.

이를 해석해보면 즉, PTM의 파라미터를 현재 선언한 모델에 copy를 해야 하는데, 이를 하기도 전에 선언한 모델의 shape을 resize해줬기 때문에 차원이 맞지 않아서 에러가 발생하는 것이다.

 

해결 방법은 load_state_dict()를 먼저 한 후에 resize_token_embeddings()를 수행하면 된다.

나의 경우 trainer.py 파일 내의 self.load_state_dict_in_model(state_dict) 직후에 resize_token_embeddings()을 수행하도록 코드를 수정함으로써 에러를 해결하였다.