✔️

LLM model を, PEFT checkpoint から from_pretrained するメモ

2023/11/01に公開

LLM

peft

tech

背景

huggingface PEFT で LoRA 学習したり incremental pretraining したあとで, checkpoint からモデルをロードしたい.

checkpoint に対して from_pretrained で一発ではいけませんでした.

base model をロード
PEFT(LoRA)の weight で state_dict を上書き

の手順を踏む必要があります.

checkpoint からのロード

最新 checkpoint の取得

まず, PEFT Trainer の pretrained 保存 callback で huggingface transformers で save_pretrained (or を使った)場合, 以下で最新の checkpoint を取得できます.

from transformers.trainer_utils import get_last_checkpoint

checkpoint_dir="output_dir/"
last_checkpoint = get_last_checkpoint(checkpoint_dir)

そしてこのフォルダに対して from_pretrained すれば OK.

from transformers import LlamaForCausalLM, LlamaTokenizer
from transformers.trainer_utils import get_last_checkpoin
from peft import PeftModel, PeftConfig

checkpoint_dir="output_dir/"
last_checkpoint = get_last_checkpoint(checkpoint_dir)

config = PeftConfig.from_pretrained(last_checkpoint)
print(config.base_model_name_or_path)
model = LlamaForCausalLM.from_pretrained(config.base_model_name_or_path, device_map="cpu")

model = PeftModel.from_pretrained(model, last_checkpoint)

from_pretrained には local_files_only=True で, モデルデータをファイルから読むように指定するとよいかもです.

Module(Layer)情報の取得

確認のため, model のネットワーク情報を表示するとよいでしょう.

for name, module in model.named_modules():
    print(name)
    ...

と named_modules でいけます.

model                                                                         model.layers
model.layers.0
model.layers.self_attn
...

全 layer の parameter を取得したい場合は, named_parameters を使います.

PEFT(LoRA) 学習時の checkpoint 保存

学習時の checkpoint 保存では, Trainer の callback で PEFT model をセーブするのが推奨されています.

issue

いくつかの LoRA training コードでは, 以下のようなのがありますが,

      model = get_peft_model(model, peft_config)
      model.print_trainable_parameters()
  
      old_state_dict = model.state_dict
  
      # ?
      model.state_dict = (
          lambda self, *_, **__: get_peft_model_state_dict(self, old_state_dict())
      ).__get__(model, type(model))

これだと二回目以降の save_pretrained で正しく lora weight が保存されません.
とりあえず issue にあるように, コード削除でよいでしょう.

また, Trainer の callback で PEFT model をセーブするのが推奨されています.

embedding size は元モデルのまま?

Japanese LLaMa など, embeddings の大きさを　model.resize_token_embeddings() で変えた場合,

save_pretrained では, しかし embedding の次元は元モデルのまま(base model のまま)になります.

weight 自体は PEFT のほうの checkpoint にうまく保存されていますので,

base model ロード
resize_token_embeddings()
PEFT from_pretrained で weight(state_dict)の読み直し

で対応する必要があります.