💾

LLMのプロンプトで口調を指示すると言わせてる感があるのでファインチューニングする

2024/12/02に公開

機械学習

LLM

Fine-tuning

tech

この記事は「LLM・LLM活用 Advent Calendar 2024」の2日目の記事になります！
(間に合ってよかった。。)

記事を書くに至った経緯

性能がいいモデルがどんどん公開されて嬉しい、けど皆「ですます口調」でちょっと味気ない
指示性能もいいのでプロンプトに口調の指示を書けばそれっぽく話してくれる
でも、なんか言わせてる感すごくてちょっと嫌だ。。
じゃあファインチューニングして重みから改造しちゃお！

すること

サンプルの質問(5個)に対して、求める口調での回答を手作業で作成
ファインチューニングしたいモデル(以下、被FTモデル)を選ぶ
指示に従って口調を変換できるモデル(以下、言い換えモデル)を選ぶ
事前に用意した1000件の質問に対し被FTモデルで回答を生成
1を基にFew-Shotプロンプトを作成し、言い換えモデルで4の1000件の回答を言い換え
5を学習データとして被FTモデルをLoRAを用いてSFT
W&B Sweepsを使って学習を見守りつつハイパーパラメータ探索
できた一番良いAdapterを被FTモデルにマージしてGGUF変換してllama.cppで推論

各ステップの詳細

1. 口調サンプルの作成

これは単純にどんな口調にしたいのかという部分です。
今回はフランクで絵文字を使うような口調にしていきたいと思います。

{
  "input": "富士山の高さは？",
  "expected": "富士山の高さは3,776mだよ！すごいよね！🤗"
}

2. ファインチューニングするモデルの選択

こちらもそのままです。
今回の選択基準は次の2つです。

VRAM16GBで推論、学習ができること(ローカルPCがRTX4060Ti16GBのため)
ライセンスがApache-2.0、MIT、CC-○○(NCは問わない)であること

Nejumi LLMリーダーボード3などで探したところ、llm-jp/llm-jp-3-3.7b-instructが良さそうでしたので、こちらで進めます。

3. 口調を言い換えできるモデルの選択

なぜモデルを2つに分けるのかというと、今回のハード要件的にモデルを1つだけにしてデータ生成までさせると、出力フォーマットを上手く守れなかったり、口調を反映できていなかったりしたためです。

推論だけであればVRAM16GBならcontext_length=2048でgemma2-27BのIQ4_XSが動きますし、時間をかけられるのならQwen2.5-32BをCPUで動かすこともできます。(恐らく今回のケースだと丸二日くらいでいけそうだった)

ということで、被FTモデルはLoRAを使ってVRAMに載る最大サイズのモデルを、言い換えモデルはCPU推論も含めて推論できる一番指示性能の高いモデルを選ぶことにしました。
今回はgoogle/gemma-2-27b-itのIQ4_XSを選択しました。

4. 被FTモデルでの素データ生成

今回の学習用データの元となる1000件のデータはkunishou/databricks-dolly-15k-jaよりランダムに抽出しました。

それなら、データセットのoutputをそのまま言い換えればいいのでは？と思ったのですが、「データセットoutputを言い換えパターン」と「一度被FTモデルで回答を生成して言い換えするパターン」では、学習時のlossの始まりがそれぞれ2.1と1.75で差がありました。

今回は口調を変更するという比較的変更要素が少ない(と思われる)学習のため、いたずらにlossを高いところから始めるよりも低いほうがいいだろうと思い後者を取っています。

5. 言い換えモデルで口調を反映

1で作成した数件のサンプルをFew-Shotとして与えて、4で生成したデータを言い換えていきます。
プロンプトは次のようにしています。

以下に質問と回答の例が5つあります。それぞれの回答の口調に注意してください。
これらを参考にして、新しい回答を同じ口調に修正してください。
修正する際は元の意味や文脈を変えないようにしてください。
修正した回答は出力フォーマットに従って出力してください。

---

**例1**
**質問:** 富士山の高さは？
**回答:** 富士山の高さは3,776mだよ！すごいよね！🤗

---

**例2**
**質問:** ...
**回答:** ...

---

**例3**
**質問:** ...
**回答:** ...

---

**例4**
**質問:** ...
**回答:** ...

---

**例5**
**質問:** ...
**回答:** ...

---

**質問:** ポテトチップスの袋は、なぜ開封後に古くなるのでしょうか？
**新しい回答:** 

---

出力フォーマット

{
  **新しい回答:** 
}

6&7. LoRAを用いた学習(W&B Sweepsでハイパーパラメータ探索)

モデルとデータができたので学習に入っていきます。

みなさんご存じの通り、7Bや8Bといった今日では小規模と言われるモデルでもフルパラメータチューニングをしようとするとA100やH100といった超ハイエンドなハードを必要とします。

そのため、ある程度の性能を持ったモデルをファインチューニングするとなるとLoRAなどの学習コストを下げる工夫をすることがスタンダードになってきています。
(とはいえLoRAでは新たな概念の獲得は難しいとも言われておりパフォーマンスは低下する可能性も)

今回はVRAMが16GBという制約なので、3B～4Bのモデルを4bitでロードして、かつLoRAで学習パラメータを減らして学習していくことにしました。
RTX3090やRTX4090などを持たれているGPUリッチな方はCohereForAI/aya-expanse-8bの4bitロード&QとVをターゲットで22.3GBでしたのでどうぞ。(何が)

スクリプトは次のようにしています。

run.py

pip install transformers trl peft bitsandbytes wandb
wandb login

from datasets import load_from_disk
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from trl import SFTConfig, SFTTrainer, DataCollatorForCompletionOnlyLM
from peft import LoraConfig, get_peft_model
import bitsandbytes as bnb


dataset = load_from_disk("dataset/test")

model_path = "llm-jp/llm-jp-3-3.7b-instruct"
tokenizer = AutoTokenizer.from_pretrained(model_path)

def find_all_linear_names(model):
    target_class = bnb.nn.Linear4bit
    linear_layer_names = set()
    for name_list, module in model.named_modules():
        if isinstance(module, target_class):
            names = name_list.split('.')
            layer_name = names[-1] if len(names) > 1 else names[0]
            linear_layer_names.add(layer_name)
    if 'lm_head' in linear_layer_names:
        linear_layer_names.remove('lm_head')
    return list(linear_layer_names)

def formatting_prompts_func(example):
    output_texts = []
    for i in range(len(example['input'])):
        # print(example['input'][i])
        messages = [
            {"role": "user", "content": example['input'][i]},
            {"role": "assistant", "content": example['output'][i]}
        ]
        try:
            text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=False)

            output_texts.append(text)
        except TypeError:
            print("type error")
    return output_texts

response_template = [27338, 279, 70147, 28752, 18]
collator = DataCollatorForCompletionOnlyLM(
    response_template=response_template,
    tokenizer=tokenizer,
    mlm=False
)

def model_init(trial=None):
    torch.cuda.empty_cache()

    quantization_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_use_double_quant=True,
        bnb_4bit_compute_dtype=torch.float16
    )

    model = AutoModelForCausalLM.from_pretrained(model_path, quantization_config=quantization_config)

    peft_config=LoraConfig(
        r=16,
        lora_alpha=32,
        lora_dropout=0.05,
        bias="none",
        task_type="CAUSAL_LM",
        target_modules=find_all_linear_names(model)
    )

    return get_peft_model(model, peft_config)

def wandb_hp_space(trial):
    return {
        "method": "random",
        "metric": {"name": "eval_loss", "goal": "minimize"},
        "parameters": {
            "gradient_accumulation_steps": {"values": [2, 4, 8, 16]},
            "learning_rate": {"values": [1e-6, 4e-6, 8e-6, 1e-5, 4e-5, 8e-5, 1e-4, 4e-4, 8e-4]},
            "weight_decay": {"values": [0.001, 0.01, 0.1]},
            "adam_beta2": {"values": [0.95, 0.98, 0.999]},
            "num_train_epochs": {"values": [1, 2, 3, 4, 5]},
            "lr_scheduler_type": {"values": ["constant", "linear", "cosine"]},
            "neftune_noise_alpha": {"values": [5, 10, 15]},
        },
    }

def get_config(gradient_accumulation_steps, learning_rate, weight_decay, adam_beta2, num_train_epochs, lr_scheduler_type, neftune_noise_alpha):
    return SFTConfig(
        output_dir="output",
        do_train=True,
        do_eval=True,
        evaluation_strategy="steps",
        prediction_loss_only=True,
        per_device_train_batch_size=1,
        per_device_eval_batch_size=1,
        gradient_accumulation_steps=gradient_accumulation_steps,
        learning_rate=learning_rate,
        weight_decay=weight_decay,
        adam_beta1=0.9,
        adam_beta2=adam_beta2,
        adam_epsilon=1.0e-4,
        num_train_epochs=num_train_epochs,
        lr_scheduler_type=lr_scheduler_type,
        logging_dir="output/logs",
        logging_strategy="steps",
        logging_steps=5,
        save_strategy="steps",
        save_steps=1000,
        save_total_limit=1,
        save_safetensors=True,
        seed=42,
        fp16=True,
        eval_steps=10,
        dataloader_num_workers=4,
        run_name="sweep_run",
        remove_unused_columns=True,
        optim="adamw_bnb_8bit",
        report_to="wandb",
        neftune_noise_alpha=neftune_noise_alpha,
        packing=False,
    )

lora_config = LoraConfig(
    r=16,
    lora_alpha=32,
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM",
)

trainer = SFTTrainer(
    model_init=model_init,
    tokenizer=tokenizer,
    train_dataset=dataset["train"],
    eval_dataset=dataset["test"],
    args=get_config(2, 1e-6, 0.1, 0.95, 5, "cosine", 10),
    formatting_func=formatting_prompts_func,
    data_collator=collator,
)

best_trial = trainer.hyperparameter_search(
    direction="minimize",
    backend="wandb",
    hp_space=wandb_hp_space,
    n_trials=100,
)

print(f"Best trial results: {best_trial}")

例外を握りつぶしていたり、ところどころ変なのは気にしないでください。。

線形層の名前の抽出は下記を参考にしています。
https://note.com/npaka/n/na506c63b8cc9

本当はLoraConfigもハイパーパラメータの探索をしたかったのですが、どうやらこの方法ではまだtransformersライブラリで対応されていない様子。

ちなみに今回のケースでのVRAM使用量は15.1GBでした。結構ぎりぎり。
1つのrunにかかった時間は数十分です。1日～2日くらい回せばよさげなパラメータが出てくるかも。

8. Adapterマージ&GGUF変換&llama.cpp推論

ここからはほとんどllama.cppの説明になりますので手短に行きます。

まず作成したLoRA Adapterを学習に使用したモデルにマージします。

from peft import PeftModel
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

name = "llm-jp/llm-jp-3-3.7b-instruct"
tokenizer = AutoTokenizer.from_pretrained(name)

model = AutoModelForCausalLM.from_pretrained(name, device_map="cpu", torch_dtype=torch.float16)

peft_name = "./output/run-74exhbi2/checkpoint-100"
model = PeftModel.from_pretrained(
    model,
    peft_name,
    device_map="cpu",
)
model.eval()

merged_model = model.merge_and_unload()
merged_model.save_pretrained("./merged_model")
tokenizer.save_pretrained("./merged_model")

重みをロードしてそれを保存するだけなのでRAMを使用しています。
以下の記事を参考にさせていただきました。
https://aipracticecafe.site/detail/7

次にllama.cppでGGUFに変換、量子化して推論していきます。

# llama.cppをクローンしてビルド(変換にGPUは不要ですが、推論で必要になるので)
git clone https://github.com/ggerganov/llama.cpp.git
cd llama.cpp
cmake -B build -DGGML_CUDA=ON
cmake --build build --config Release -j8

# pythonの仮想環境を作ってpip install
python -m venv --prompt . .venv
source ./.venv/bin/activate
pip install -r requirements/requirements-convert_hf_to_gguf.txt

# GGUF変換
python convert_hf_to_gguf.py ../merged_model --outfile ../model.gguf --outtype f16

# 量子化
./build/bin/llama-quantize ../model.gguf ../model-Q4_K_M.gguf Q4_K_M

# 推論(Server)
./build/bin/llama-server -m ../model-Q4_K_M.gguf -c 4096 -n 1024 --ngl 29

推論結果はこちら

### 指示:
もうすぐクリスマスですね。

### 応答:
クリスマス、楽しみだね！🎄✨ サンタさんに何お願いしようか考えてるのかな？😊

まとめ

何番煎じかわかりませんが、LoRAを用いたSFTで口調を調整してみました。
もう少し詰められた部分もあるので、今後も引き続きブラッシュアップしていきたいと思います。

気が向けばUIつけて使いやすいようにして公開するかも？
その際はまた記事を書きます！

読んでいただきありがとうございました！