ローカルLLMのLoRAモデルをHuggingFaceにあげるまでのまとめ

概要

ローカルLLMについて日本語データセットを用いてLoRAを行い、それをHuggingFaceに保存するまでの手順を備忘録としてまとめてみました。
ベースモデルはllm-jp-3-13bで、使用したデータセットはELYZA-tasks-100です。
また、今回はunslothを使用しています。

本節

Google Colab上で実装を行ったのでまず初めにunslothを

!pip uninstall unsloth -y
!pip install --upgrade --no-cache-dir "unsloth[colab-new] @ git+https://github.com/unslothai/unsloth.git"

で最新の状態にしました。

LoRAの流れとしては
1.ベースモデル(FTを行う対象のモデル)をロード
2.使用するデータセットをロード
3.訓練時の各パラメータの設定
4.学習実行
5.学習したモデルでタスクの実行(推論)
6.推論結果をjsonl形式で保存
7.モデルとトークナイザーをHugging Faceにアップロード
です。

順にまとめていきます。

モデルのロード

モデルによって細かい部部分が違うかもしれないのでHuggingFaceのページを確認するとよいです。
今回のモデルだと

load_model.py

from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from unsloth import FastLanguageModel
import torch

max_seq_length = 2048 # コンテキスト長（自由に設定可能）
dtype = None # Noneにしておけば自動で設定
load_in_4bit = True 

model_id = "llm-jp/llm-jp-3-13b"
new_model_id = "llm-jp-3-13b-LoRA" #Fine-Tuningしたモデルにつけたい名前
# FastLanguageModel インスタンスを作成
model, tokenizer = FastLanguageModel.from_pretrained(
    model_name=model_id,
    dtype=dtype,
    load_in_4bit=load_in_4bit,
    trust_remote_code=True,
)

ここでは4bit量子化によってメモリ効率を向上させています。

またここでまとめてLoRAの微調整のためのモデルを定義しておきます。

prepare_LoRA.py

# SFT用のモデルを用意
model = FastLanguageModel.get_peft_model(
    model,
    r = 32,
    target_modules = ["q_proj", "k_proj", "v_proj", "o_proj",
                      "gate_proj", "up_proj", "down_proj",],
    lora_alpha = 256,
    lora_dropout = 0.05,
    bias = "none",
    use_gradient_checkpointing = "unsloth",
    random_state = 3407,
    use_rslora = False,
    loftq_config = None,
    max_seq_length = max_seq_length,
)

各パラメータについて
model:上でロードしたモデル
r:LoRAでの使う行列のランクのこと。大きいほど表現力は増すがメモリ消費増
target_modules:微調整したい層
lora_alpha:学習率スケーリング係数
lora_dropout:ドロップアウト率
bias:"none"にするとバイアスのパラメータは更新しない
use_gradient_checkpointing: 勾配チェックポイントを有効にしてメモリ効率向上
random_state:シード値決めることで再現性を保証
use_rslora: FalseでRSLoRA使わないようにしている
loftq_config = None: LOFT-Qを使わないように
max_seq_length: シーケンスの最大長さを指定

データセットのロード

HuggingFaceのページにコードが書いているのでそれを使いましょう。
今回は以下で読み込みできます。

load_dataset.py

from datasets import load_dataset
dataset = load_dataset("elyza/ELYZA-tasks-100")

さらにトレーニングセットと検証セットに分けておきます。今回は8:2で分けます。

split_data.py

dataset = test_dataset.train_test_split(test_size=0.2, shuffle=True, seed=42)

訓練のための各種設定

はじめに学習のためにデータセットの形式を統一しておきます。

format.py

prompt = """### 指示
{}
### 回答
{}"""



"""
formatting_prompts_func: 各データをプロンプトに合わせた形式に合わせる
"""
EOS_TOKEN = tokenizer.eos_token # トークナイザーのEOSトークン（文末トークン）
def formatting_prompts_func(examples):
    input = examples["text"] # 入力データ
    output = examples["output"] # 出力データ
    text = prompt.format(input, output) + EOS_TOKEN # プロンプトの作成
    return { "formatted_text" : text, } # 新しいフィールド "formatted_text" を返す
pass

# # 各データにフォーマットを適用
dataset = dataset.map(
    formatting_prompts_func,
    num_proc= 4, # 並列処理数を指定
)

ここでformatting_prompts_func中でのinput、outputについては実際にデータの列名を確認して設定する必要があります。

続いて訓練時に使うパラメータを設定します

train_params.py

from trl import SFTTrainer
from transformers import TrainingArguments
from unsloth import is_bfloat16_supported

trainer = SFTTrainer(
    model = model,
    tokenizer = tokenizer,
    train_dataset=dataset["train"],
    max_seq_length = max_seq_length,
    dataset_text_field="formatted_text",
    packing = False,
    args = TrainingArguments(
        per_device_train_batch_size = 4,
        gradient_accumulation_steps = 2,
        num_train_epochs = 1,
        logging_steps = 10,
        warmup_steps = 10,
        save_steps=100,
        save_total_limit=2,
        max_steps=-1,
        learning_rate = 2e-4,
        fp16 = not is_bfloat16_supported(),
        bf16 = is_bfloat16_supported(),
        group_by_length=True,
        seed = 3407,
        output_dir = "outputs",
        report_to = "none",
    ),
)

訓練の実行

では実際に訓練を行ってみましょう。以下のコードでできます。

train.py

trainer_stats = trainer.train()

今回はGPUとしてA100を使用して、訓練にかかる時間は6分程度でした。

次に推論用のデータセットをロードしますが、データのロードは同様なので割愛します。

推論の実行

inference.py

from tqdm import tqdm

# 推論するためにモデルのモードを変更
FastLanguageModel.for_inference(model)

results = []
for dt in tqdm(datasets):
  input = dt["input"]

  prompt = f"""### 指示\n{input}\n### 回答\n"""

  inputs = tokenizer([prompt], return_tensors = "pt").to(model.device)

  outputs = model.generate(**inputs, max_new_tokens = 512, use_cache = True, do_sample=False, repetition_penalty=1.2)
  prediction = tokenizer.decode(outputs[0], skip_special_tokens=True).split('\n### 回答')[-1]

  results.append({"task_id": dt["task_id"], "input": input, "output": prediction})

for文の中について、
まずプロンプト(モデルへの入力)の生成をdatasetから行っています。
ここでプロンプトの形式は

### 指示
<ユーザーからの入力>
### 回答

という形式にそろえています。
さらにトークナイザーを使用してプロンプトをトークン化しています。
その後、実際にモデルにトークン化したプロンプトを入れて出力を得ます。（推論の実行）
出力もトークン化されているのでそれを自然言語に変換しています。
.split('\n### 回答')[-1]で回答部分のみを抽出しています。
この一連の操作をdataset中のdataについて繰り返し行います。
得られた結果についてそれを(タスクID、入力、出力)の辞書形式にしてresultsリスとに格納しています。

モデルの保存

推論した結果得られたモデルをjsonl形式で保存しておきます

preserve.py

with open(f"llm-jp-3-13b-LoRA_output.jsonl", 'w', encoding='utf-8') as f:
    for result in results:
        json.dump(result, f, ensure_ascii=False)
        f.write('\n')

モデルのHugging Faceへのアップロード

model.push_to_hub_merged(
    "llm-jp-3-13b-LoRA", #モデルid名
    tokenizer=tokenizer,
    save_method="lora",
    token=HF_TOKEN, #HuggingFaceのトークン
    private=True #公開設定この場合はprivateになる
)

このコードによってモデルをあげることができます。
HuggingFaceのトークンについては

huggingface-cli login

で表示されるURLに飛ぶと作ることができます。

まとめ

今回はLoRAの手法に絞ってLLMモデルのファインチューニングを行ってみました。
ファインチューニングの流れを一つずつ確認することでよりFTの手法について理解が深まりました。
次はLoRA以外の手法についても試してみようと思います。