【備忘録】大規模言語モデルをLoRA/QLoRAでファインチューニング

はじめに

huggingfaceに公開されている大規模言語モデルをLoRA/QLoRAでファインチューニングするのに調べた情報をまとめた備忘録。
各LoRA実行用のコードとコード内容の理解にあたって調査した内容をまとめている。

想定読者

大規模言語モデルをLoRAでファインチューニングしたい人（画像は対象外）
LoRAや大規模言語モデルについては説明省略。以下の記事は参考になった。

実行環境

• GPU
  • RTX3090
  • V100×4
• pip
  • transformers: 4.38.2
  • torch: 2.2.0
  • trl: 0.7.10
  • peft: 0.10.0

コード集

後々Githubにまとめて公開予定

コードの説明

モデルの読み込み

from transformers import AutoModelForCausalLM, AutoTokenizer
# 利用するLLMを指定
model_id = "elyza/ELYZA-japanese-Llama-2-7b-instruct"
# モデルによっては「use_fast=False」を追加
tokenizer = AutoTokenizer.from_pretrained(model_id, add_eos_token=True)
tokenizer.padding_side = "right" # おまじない
model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype="auto", device_map="auto")

LoRAの設定適用

ひとまず、LoRAでファインチューニングする場合、HuggingFaceが提供している「peft」ライブラリを扱うのがよい。

from peft import get_peft_model, LoraConfig, TaskType
peft_config = LoraConfig(task_type=TaskType.CAUSAL_LM,
                         inference_mode=False, # 学習時はFalse
                         r=8, 
                         lora_alpha=16,
                         lora_dropout=0.05,
                         bias="none",
                         target_modules=["q_proj", "v_proj", "k_proj", "o_proj", "gate_proj", "up_proj", "down_proj", "embed_tokens", "lm_head"],
                         )

# モデルにLoRAアダプター適用、更新対象のパラメータ数の確認
model = get_peft_model(model, peft_config)
model.print_trainable_parameters()

target_modulesには元のモデルのモジュール名を入れる。
「elyza/ELYZA-japanese-Llama-2-7b-instruct」で扱っているモジュールをすべて対象にしている。

ドキュメントがある程度まとまっているが、PEFT手法の一部にバグがあるようなので、要注意。

学習用データセットの作成

from datasets import load_dataset
# 指示に対する回答がついているデータセットを語尾が「ござる」になるように変更したデータセット
dataset = load_dataset("bbz662bbz/databricks-dolly-15k-ja-gozaru", split="train")

def dolly_prompt_template(data):
    if data["input"]:
        result = f"[INST] {data['instruction']}\n\n{data['input']} [/INST] {data['output']}"
    else:
        result = f"[INST] {data['instruction']} [/INST] {data['output']}"
    return result

def add_text(example):
    example["text"] = dolly_prompt_template(example)
    for key in ["index", "category", "instruction", "input", "output"]:
        del example[key]
    return example

dataset = dataset.map(add_text)

LoRAのファインチューニング対象となるベースのモデルのモデルカードの情報やtokenizer_config.jsonを参考にデータを加工。
今回は

[INST]
指示
[/INST]
回答

もしくは

[INST]
指示

入力情報
[/INST]
回答

の形式で作成している。
学習結果としては指示された内容にしたがって回答を生成（語尾は～でござる）となることを期待。

学習

こちらもhuggingfaceがtrlというライブラリを出しているのでそれを利用。
Transformer Reinforcement Learningの略ということだが、強化学習特化という感じはなく、huggingfaceで扱うモデルを学習するためのライブラリとして使えるという認識。
ChatGPTを代表にしたLLMで取り扱われている学習方法は一通りサポートされていそう。（DPOやKTOもあった。）

from transformers import TrainingArguments
from trl import SFTTrainer

# モデル名と出力先の設定
peft_name = "lora-elyza-7b-instruct-gozaru"
output_dir = "lora-elyza-7b-instruct-gozaru-results"

save_steps = 100
logging_steps = 20

# 学習パラメータ
training_arguments = TrainingArguments(
    output_dir=output_dir,  # 出力ディレクトリ
    max_steps=300,  # 学習ステップ数
    per_device_train_batch_size=4,  # 学習用のGPUあたりのバッチサイズ
    gradient_accumulation_steps=1,  # 勾配を蓄積するための更新ステップの数
    optim="paged_adamw_32bit",
    learning_rate=1e-4,  # 初期学習率
    lr_scheduler_type="cosine",  # 学習率スケジュール
    max_grad_norm=0.3,  # 最大法線勾配 (勾配クリッピング)
    warmup_ratio=0.03,  # 線形ウォームアップのステップ比率 (0から学習率まで)
    weight_decay=0.001,  # bias/LayerNormウェイトを除く全レイヤーに適用するウェイト減衰
    save_steps=save_steps,  # 何ステップ毎にチェックポイントを保存するか
    logging_steps=logging_steps,  # 何ステップ毎にログを記録するか
    group_by_length=True,  # シーケンスを同じ長さのバッチにグループ化 (メモリ節約)
    report_to=None  # レポート→noneにするとwandbを使わなくなる。
)

# トレーナーの準備
trainer = SFTTrainer(
    model=model,
    tokenizer=tokenizer,
    train_dataset=dataset,
    dataset_text_field="text",  # データセットのtext列
    peft_config=peft_config,  # PEFTパラメータ
    args=training_arguments,  # 学習パラメータ
    max_seq_length=None,  # 使用する最大シーケンス長
    packing=False,  # 同じ入力シーケンスに複数サンプルをパッキング(効率を高める)
)

トレーナーとして設定するSFTTrainerのパラメータにDataCollatorForCompletionOnlyLMを設定するとアシスタントの生成部分を学習対象にする。（lossの計算対象にする）

上記リンクのUsing token_ids directly for response_templateの欄にも記載されているが、collatorをうまく作成しないとすべて学習対象外になってしまい、lossが0になる。
DataCollatorForCompletionOnlyLMの設定と-100（学習対象外）ラベルを確認するコードを以下に記載。

def formatting_prompts_func(example):
    output_texts = []
    for i in range(len(example)):
        if example["input"][i]:
            result = f"[INST] {example['instruction'][i]}\n\n{example['input'][i]} [/INST] ### 応答: {example['output'][i]}"
        else:
            result = f"[INST] {example['instruction'][i]} [/INST] ### 応答: {example['output'][i]}"
        output_texts.append(result)
    return output_texts


response_template = "### 応答: "
response_template_ids = tokenizer.encode(response_template, add_special_tokens=False)[2:]
collator = DataCollatorForCompletionOnlyLM(response_template_ids, tokenizer=tokenizer)

trainer = SFTTrainer(
    model=model,
    train_dataset=dataset,
    formatting_func=formatting_prompts_func, # 追加
    data_collator=collator, # 追加
    peft_config=peft_config,
    args=training_arguments,
    max_seq_length=None,
    packing=False,
)

# collatorの確認
from torch.utils.data import DataLoader
loader = DataLoader(trainer.train_dataset, collate_fn=collator, batch_size=8)

batch = next(iter(loader))
print(batch['labels'][0])
"""
tensor([ -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,
         -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,
         -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,
         -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,
         -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,
         -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,
         -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,
         -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,
         -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,
         -100,  -100,  -100, 31812, 30767, 29909, 29901, 29871,   232,   178,
          160, 30466, 30298, 30332, 30364, 30538, 30353, 31688, 30458, 30723,
          230,   132,   161, 30723,   230,   132,   161, 30326, 30366, 30199,
        30499, 30330, 31558, 30538, 30466, 31192, 30466, 30735, 30366, 30513,
        30579, 30439, 30597, 30458, 30298, 30366, 30199, 30584, 29871,  -100,
         -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,
         -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,
         -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,
         -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,
         -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,
         -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,
         -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,
         -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,
         -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,
         -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,
         -100,  -100,  -100,  -100,  -100,  -100,  -100])
"""

学習の実行とモデルの保存は以下のコードを利用。LoRAモデルはadapter付きのモデルになる。

# 学習の実行
model.config.use_cache = False
trainer.train()

# LoRAモデルの保存
trainer.model.save_pretrained(peft_name)

LoRAモデルで推論

from peft import AutoPeftModelForCausalLM
from transformers import AutoTokenizer

model_id = "elyza/ELYZA-japanese-Llama-2-7b-instruct"

# モデルの読み込み
model = AutoPeftModelForCausalLM.from_pretrained("./lora-elyza-7b-instruct-gozaru-results/checkpoint-300", torch_dtype="auto", device_map="auto")

# トークナイザーの準備
tokenizer = AutoTokenizer.from_pretrained(model_id)

「ござる」がうまくいった例

ハルシネーションを起こしているが、「ござる」のファインチューニングはうまくいっている。

prompt = "[INST] 富士山の高さは？ [/INST] "

# 推論の実行
input_ids = tokenizer(prompt, add_special_tokens=False, return_tensors='pt').to(model.device)
output_ids = model.generate(
    **input_ids,
    max_new_tokens=100
)
output = tokenizer.decode(output_ids.tolist()[0])
print(output)
"""
[INST] 富士山の高さは？ [/INST] 2542mでござる。
富士山には、多くの伝説があり、日本山岳賛歌では日本人から愛されていますでござる。富士山は国土地理院から「日本百名山」に指定されていますでござる。富士山は、
"""

「ござる」がうまくいかない例

時代劇でベルギー、スイス...？というのもあるが、最初以外「ござる」がついていない

prompt = "[INST] 熊と少女の時代劇を書いて [/INST] "

# 推論の実行
input_ids = tokenizer(prompt, add_special_tokens=False, return_tensors='pt').to(model.device)
output_ids = model.generate(
    **input_ids,
    max_new_tokens=100
)
output = tokenizer.decode(output_ids.tolist()[0])
print(output)
"""
[INST] 熊と少女の時代劇を書いて [/INST] 描けそうなテーマと時代に関しては、「ベン・ハー」でござる。
ベン・ハーは、狼や虎のような犬で、ベルギーでは家出の少女の愛玩動物として親しまれています。ベン・ハーと少女の物語を書くには、十九世紀のベルギーが目標になります。十九世紀になると、ベン・ハーはドス・コスと呼ばれるようになり、現代では小さなスイスで愛されている。
"""

最後に

LoRA以外のPEFT手法についても備忘録としてまとめる予定。
複数のLoRAモデルをマージした生成も試す