[PyTorch 에러 해결] Bert 모델 token embedding layer 차원 오류

실행 환경

• OS : Linux
• PyTorch : 2.0.0 version

 

첫번째 에러

RuntimeError: CUDA error: device-side assert triggered
CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect. For debugging consider passing CUDA_LAUNCH_BLOCKING=1.
Compile with TORCH_USE_CUDA_DSA to enable device-side assertions.

에러 메시지에서 확인할 수 있듯이, stacktrace의 내용만으로는 어디에서 에러가 발생했는지 파악하기 어렵다.

이를 해결하기 위해서는 아래 코드를 기존 코드에 추가하여 더 구체적인 에러 메시지를 확인해야 한다.

import os

os.environ['CUDA_LAUNCH_BLOCKING'] = "1"

 

두번째 에러

RuntimeError: CUDA error: device-side assert triggered
Compile with TORCH_USE_CUDA_DSA to enable device-side assertions.

안타깝게도 위 코드 적용 이후에도 추가적인 정보를 얻을 수 없었다.

계속 조사해본 결과, 입출력 차원이 제대로 맞지 않는 경우에 자주 발생하는 에러라는 것을 알아냈다.

코드를 살펴보니, tokenizer에 special token을 추가해줬는데 모델의 token embedding layer의 차원을 수정하지 않았다는 것을 확인했다.

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, additional_special_tokens = ["[ENT]", "[/ENT]"])

model = AutoModelForSequenceClassification.from_pretrained(klue/bert-base)

즉, tokenizer의 vocab size는 32000개에서 32002개로 증가했지만, token embedding layer의 차원은 [32000, 768]이므로 32000번 이후의 vocab에 대한 처리를 할 수 없었던 것이었다.

따라서 모델의 token embedding layer의 입력 차원을 수정하여 에러를 해결했다.

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, additional_special_tokens = ["[ENT]", "[/ENT]"])

model = AutoModelForSequenceClassification.from_pretrained(klue/bert-base)
model.resize_token_embeddings(len(tokenizer)) # 모델의 token embedding layer 입력 차원 수정

* 위 코드에서 tokenizer.vocab_size 대신 len(tokenizer)를 사용한 이유는 tokenizer.vocab_size는 추가된 토큰을 제외한 vocab size를 보여준다. 따라서 추가된 토큰을 포함한 vocab size를 확인하기 위해서는 len(tokenizer)를 사용해야한다.

• tokenizer.vocab_size : 추가된 토큰을 제외한 vocab size
• len(tokenizer) : 추가된 토큰을 포함한 vocab size