LLama 3のSFTTrainer+Weights & Biasesでファインチューニング

AI・機械学習

はじめに

LLama 3は、Meta社が開発した大規模言語モデルです。高性能でありながら、一般的なGPUでも扱えるサイズのモデルが提供されています。このモデルをファインチューニングすることで、様々なタスクに適用できます。

本記事では、Hugging Faceのtransformersライブラリとtrlライブラリを使って、LLama 3をファインチューニングする方法を解説します。


こちらの記事もおすすめ

Meta社の新しいオープンソースLLMモデル「Llama 3」がリリースされました
Meta社が、オープンソースの大規模言語モデル(LLM)ファミリー「Llama」の最新版である「Llama 3」を2024年4月18日にリリースしました。Llama 3モデルはHugging Faceからアクセス可能で、Hugging Fa...

環境のセットアップ

最初に必要なライブラリをインストールします。

!pip install -U transformers trl accelerate bitsandbytes
!pip install -U peft
!pip install --upgrade huggingface_hub

次に、Hugging FaceとWandbにログインします。

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

import wandb
from google.colab import userdata
wandb_api_key = userdata.get('MY_WANDB_KEY') 
!wandb login $wandb_api_key

データセットの準備

ファインチューニングに使用するデータセットを用意します。ここでは例として、Hugging Faceのdatasetsライブラリから"no_robots"データセットを読み込みます。

from datasets import load_dataset

dataset = load_dataset("HuggingFaceH4/no_robots")

# データセットのサイズを10分の1に縮小
train_dataset = dataset["train"].train_test_split(test_size=0.1, seed=42)["test"]
test_dataset = dataset["test"].train_test_split(test_size=0.1, seed=42)["test"]

# 新しいデータセットを作成
dataset_mini = {
    "train": train_dataset, 
    "test": test_dataset
}

モデルの読み込み

LLama 3のモデルをHugging Faceのモデルハブから読み込みます。ここでは8Bのベースモデルを使用します。

from transformers import AutoModelForCausalLM, AutoTokenizer

model_name = "meta-llama/Meta-Llama-3-8B"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name, low_cpu_mem_usage=True, load_in_4bit=True)

low_cpu_mem_usageload_in_4bitオプションを使って、メモリ使用量を減らした4bitの量子化モデルを読み込んでいます。

ファインチューニングの実行

trlライブラリのSFTTrainerを使ってファインチューニングを実行します。あわせて、TrainingArgumentsを使ってハイパーパラメータを設定します。

from trl import SFTTrainer
from transformers import TrainingArguments

# LoRAアダプターの設定
peft_config = LoraConfig(
    task_type=TaskType.CAUSAL_LM,
    inference_mode=False, 
    r=8,
    lora_alpha=32,
    lora_dropout=0.1,
)

model = get_peft_model(model, peft_config)

# トレーニング設定
training_args = TrainingArguments(
    output_dir="./llama3-sft",
    per_device_train_batch_size=1,
    num_train_epochs=1,
    learning_rate=0.0001,
    logging_steps=10,
    report_to="wandb"  
)

trainer = SFTTrainer(
    model=model,
    tokenizer=tokenizer,
    train_dataset=dataset_mini["train"],
    args=training_args,
    max_seq_length=1024,
)

trainer.train()
  • per_device_train_batch_size: 1つのGPUあたりのバッチサイズ。メモリに収まる範囲で大きくすると学習が速くなります。
  • num_train_epochs: 学習するエポック数。
  • learning_rate: 学習率
  • logging_steps: ログを出力する頻度。10ステップごとにログを出力します。
  • max_seq_length: 入力系列の最大長。LLama 3は最大8192トークンまで対応しています。

また、LoRAアダプターを使ってパラメータ数を削減しています。これにより、ファインチューニングの計算コストを大幅に下げることができます。

wandbでWeights & Biasesにログを保存することもできます。

ファインチューニング済みモデルの保存とロード

ファインチューニングが完了したら、モデルを保存します。

trainer.save_model("/content/drive/MyDrive/PROJECT/llama3-sft-model")

保存したモデルは、通常のTransformersのモデルと同様に読み込むことができます。

finetuned_model = AutoModelForCausalLM.from_pretrained("/content/drive/MyDrive/PROJECT/llama3-sft-model")

まとめ

本記事では、LLama 3のファインチューニング方法について解説しました。
手順をまとめると以下の通りです。

  1. 必要なライブラリをインストールし、Hugging FaceとWandbにログイン
  2. ファインチューニングに使用するデータセットを準備
  3. LLama 3のモデルを読み込む
  4. SFTTrainerを使ってファインチューニングを実行
  5. ファインチューニング済みモデルを保存し、必要に応じて再度ロード

ファインチューニング済みのモデルを使って、様々なタスクに適用してみてください。
質の高い文章生成が可能になるはずです。

ぜひ、プロンプトを工夫したり、オプションを調整したりして、LLama 3の性能を体験してみてください。

ノートブック

Google Colaboratory

参考サイト

Welcome Llama 3 - Meta's new open LLM
We’re on a journey to advance and democratize artificial intelligence through open source and open science.
meta-llama/Meta-Llama-3-8B at main
We’re on a journey to advance and democratize artificial intelligence through open source and open science.
meta-llama/Meta-Llama-3-8B-Instruct · Hugging Face
We’re on a journey to advance and democratize artificial intelligence through open source and open science.

コメント

タイトルとURLをコピーしました