[HuggingFace][Gemma] RuntimeError: shape '[1, 20, 3072]' is invalid for input of size 81920

문제 상황

Hugging Face Gemma 모델 실행 중 다음 에러 발생

text = "Quote: Imagination is more"
device = "cuda:0"
inputs = tokenizer(text, return_tensors="pt").to(device)

outputs = model.generate(**inputs, max_new_tokens=20)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))

 

RuntimeError: shape '[1, 20, 3072]' is invalid for input of size 81920

 

해결 방법

transformer 4.38.1 버전으로 설치

pip3 install -q -U transformers==4.38.1

 

문제 원인

transformer 4.38.0 이하 버전에서의 버그

https://github.com/huggingface/transformers/issues/29181

Unable to reproduce the simple code snippets in official doc for google/gemma-7b · Issue #29181 · huggingface/transformers