LLMをSFTTrainerでファインチューニング

LLMをファインチューニングする

LLMは汎用的なデータセットを学習しているため、汎用的な質問に対しては正しく答えてくれることが多い。しかし、専門的な質問の場合にはその専門分野のデータセットでLLMがトレーニングされていないため、正しく答えてくれない。
そこで、ファインチューニングによって特定分野のデータセットなどを使ってLLMを追加で学習させることができる。

LLMファインチューニング環境

• 実行環境: Colaboratory (T4: GPU 16GB memory環境)
• 言語: python
• 対象LLMモデル: Mistral (mistralai/Mistral-7B-Instruct-v0.3)
• トレーニングデータセット: 日本語のデータセット(kunishou/databricks-dolly-15k-ja)

実際のColaboratoryのノートは以下リンクから開くことができる。
SFTTrainerのColaboratory

必要なライブラリのインストール

Colaboratory環境で以下のライブラリをインストールする。

!pip3 install datasets
!pip3 install -U bitsandbytes
!pip3 install transformers
!pip3 install torch
!pip3 install accelerate
!pip3 install trl
!pip3 install peft

HuggingFaceからモデル、データセットをダウンロードできるようにHuggingFaceにログイン

あらかじめColaboratoryのSecretタブでHuggingFaceのアクセストークンを設定しておき、コードの実行時にアクセスできるようにしておく。(HF_TOKENという名前で保存)

from huggingface_hub import login

from google.colab import userdata
login(token=userdata.get('HF_TOKEN'))

LLMモデルを読み込み

今回実行環境のGPUメモリ容量に制限がありLLMのトレーニングを行うとOut of Memoryが発生してトレーニングに失敗するため、4bit量子化を適用したモデルとする。
GPU環境の性能（メモリ）が十分にある場合には量子化は適用しない。

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
device = "cuda"

model_name = "mistralai/Mistral-7B-Instruct-v0.3"

quantization_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_compute_dtype=torch.float16
)

model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.float16,
    quantization_config=quantization_config,
)
tokenizer = AutoTokenizer.from_pretrained(model_name)
if tokenizer.pad_token is None:
  tokenizer.add_special_tokens({'pad_token': tokenizer.eos_token})
  model.resize_token_embeddings(len(tokenizer))

model.to(device)

読み込んだLLMの応答を確認する

読み込んだLLMが適切に応答するかどうかを確認する。

def talk(text):
    conversation = [{"role": "user", "content": text}]

    model_inputs = tokenizer.apply_chat_template(
        conversation,
        return_dict=True,
        add_generation_prompt=True,
        return_tensors="pt"
    ).to(device)

    generated_ids = model.generate(
        **model_inputs,
        pad_token_id=tokenizer.eos_token_id,
        temperature=0.1,
        do_sample=True,
        top_p=0.95,
        top_k=40,
        max_new_tokens=256,
    )

    answer = tokenizer.batch_decode(
        generated_ids,
        #skip_special_tokens = True,
    )[0]

    answer = answer.split("[/INST]")[1][:-len(tokenizer.eos_token)]

    return answer

talk("What is relation between Alice and Bob?")

出力は以下となり、適切に応答できていることが確認できた。

The relation between Alice and Bob is not specified in the information provided. They could be friends, family members, colleagues, or any other relationship. In computer networking, Alice and Bob are often used as placeholder names for two entities communicating with each other.

データセットを読み込む

以下を実行してデータセットを読み込む。今回学習に使用するデータは、海外で作成された学習データが日本語に翻訳されたものである。以下データをHuggingFace経由で使用させていただいた。

databricks-dolly-15k-ja

from datasets import load_dataset

dolly_dataset = load_dataset("kunishou/databricks-dolly-15k-ja")
dolly_train_dataset = dolly_dataset['train'].filter(lambda example: example['input'] == '')
print(dolly_train_dataset)
print(dolly_train_dataset[0])

データセット中のデータの例は以下となる。このように、分類や質疑応答といったデータが含まれている。データレコード数は10417件となっている。

{'output': 'ジョン・デイヴィソン・ロックフェラー', 'input': '', 'index': '7', 'category': 'open_qa', 'instruction': '国連に本部を建てるためにニューヨークの土地を寄付したのは誰？'}

トレーニング用プロンプトの整形とトレーニング評価箇所（LLM応答）の設定

トレーニングのためにはデータセットから指示部分と応答部分を使ってLLM学習用のプロンプトを作成する必要がある。
また、LLMの応答を適切に評価するために、LLMに入力した指示部分を取り除いた応答箇所のみを対象とするためのcollatorを設定する必要がある。今回の場合は、「### Answer:」以降がLLMの応答部分となるためこちらを設定している。

def formatting_prompts_func(example):
    output_texts = []
    for i in range(len(example['instruction'])):
        text = f"以下は、タスクを説明する指示です。要求を適切に満たす応答を書きなさい\n\n### Instruction:\n{example['instruction'][i]}\n\n### Answer:\n{example['output'][i]}" + tokenizer.eos_token
        output_texts.append(text)
    return output_texts


from trl import DataCollatorForCompletionOnlyLM

response_template = "### Answer:\n"
collator = DataCollatorForCompletionOnlyLM(tokenizer.encode(response_template, add_special_tokens = False)[2:], tokenizer=tokenizer)

トレーニングを行う

SFTTrainerを使って実際にトレーニングを行う。効率的にファインチューニングを行うためにLoRAを使って学習するパラメータ数を削減している。
LLMのトレーニングは学習用データセットのサイズに応じてかなりの時間が必要となる。そのため、まずはターゲットとなるデータセット数を削減して正しく学習が進むかを確認したのち、全体のデータセットで学習を行う。以下ではデータセット数は30としてまずは学習が進むことを確認している。
なお、SFTTrainerでの学習の進捗を保存し可視化するサービスの Weights & Biases に会員登録し、API Keyを発行する必要がある。

# ターゲットのデータセットの削減
target_dataset = dolly_train_dataset.select(range(30))

from transformers import TrainingArguments
from trl import SFTTrainer
from peft import LoraConfig

peft_config = LoraConfig(
    r=8,
    lora_alpha=16,
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM",
)

args = TrainingArguments(
    output_dir='./output',
    num_train_epochs=2,
    gradient_accumulation_steps=8,
    per_device_train_batch_size=8,
    save_strategy="no",
    logging_steps=20,
    lr_scheduler_type="constant",
    save_total_limit=1,
    auto_find_batch_size=True,
    fp16=True,
)

trainer = SFTTrainer(
    model,
    args=args,
    train_dataset=target_dataset,
    formatting_func=formatting_prompts_func,
    #max_seq_length=1024,
    data_collator=collator,
    peft_config=peft_config,
)
trainer.train()

上記コードを実行したところ、30件のデータセットで1分弱で学習が完了した。

学習結果を確認

以下データセット内の質問の答えを確認して、学習が進んだことを確認。

talk("国連に本部を建てるためにニューヨークの土地を寄付したのは誰？")

以下が出力。

国連本部の土地をニューヨークに寄付したのは、ロードス・ジェイムズ・ストリートにある現在の国連本部の地帯は、1947年に、ロックフェラーの財団が買い取りした。

まとめ

ColaboratoryのT4環境でもLLMをファインチューニングすることができた。ファインチューニングによりLLMの追加学習を行い、専門的な知識を持ったLLMや自分専用のLLMを簡単に作成することができるとわかったので、LLM使用の際に応用範囲を広げることができそうだと感じた。
ただし、今回パラメータ数が7Bで4bit量子化したモデルのファインチューニングでも使用したGPUメモリは12.8GBと16GBに対してギリギリであったので、量子化せず大きなパラメータのLLMをファインチューニングするためにはかなり高性能なGPUが必要であることも認識できた。