Llama2のLoRAをしていた時、うまくEOSが出力されず、文章生成がmax_new_tokensの数まで生成されてしまっていた。
解決するのに手間取ったので備忘録。
結論を先に示すとPAD = EOSにしていたのが原因で、EOSもattention_maskの対象となり学習されなかったのではないかと考えられる。
まずはtokenizerを見てみる。
meta公式のllama2はアクセス権の申請が必要で見れない人がいるかもしれないので、かわりにelyza/ELYZA-japanese-Llama-2-7bを例にする。
tokenizer_config.jsonを見ると、PADが設定されていない。
"pad_token": null,
そのせいか、ネット上の記事ではPAD = EOSにしている例が多い。
tokenizer.pad_token = tokenizer.eos_token
しかし、冒頭に述べた通りEOSがattention_maskの対象になり、文章生成が止まらなくなってしまう。
なので、明示的にPADを追加する必要がある。
tokenizer.add_special_tokens({'pad_token': '[PAD]'})
qlora.pyでも、PADを追加するような実装になっていた。
tokenizerの語彙数が増えたので、モデルのembeddingも増やす。
model.resize_token_embeddings(len(tokenizer))
tokenizer.vocab_sizeはspecial_tokens以外の語彙数なので、len(tokenizer)を指定する必要がある。
ここまででtokenizerとmodelの準備は終わり。
後は通法に従ってLoRAの学習を行う。
学習が終わったらアダプターをマージして保存する。
なぜマージするのかと言うと、resize_token_embeddingsしてPADという新たなトークンを学習させているため、推論時に元のモデルをロード→アダプター適応を行ってしまうと、元のモデルにとってPADは未知のままになってしまうから。
なので、マージして新たなモデルとして保存する方が望ましいと思う。
add_special_tokensしたtokenizerも忘れず保存する。
merged = model.merge_and_unload()
merged.save_pretrained(save_dir, safe_serialization=True)
tokenizer.save_pretrained(save_dir)
推論時は、tokenizerもmodelもfrom_pretrainedで保存先を指定すれば読み込める。
スタンドアロンで動くのでtext generation web uiでも使える。
これでEOSが出力されるようになり、想定通りにテキスト生成が止まるようになった。
コメント