🤗 Weights & Biases とUnslothを使用したMistral-7Bモデルのファインチューニング

はじめに
環境設定
データセットの準備
モデルの準備
ファインチューニングの実行
まとめ
ノートブック
参考サイト
1. 関連

はじめに

この記事では、unslothライブラリを使用して、Mistral-7Bモデルをファインチューニングする方法について解説します。unslothは、メモリ使用量を大幅に削減しながら、高速にファインチューニングを行うことができるライブラリです。

Mistral-7Bは、そのサイズにおいて最も強力な言語モデルの1つであり、以下の特徴を持っています:

Llama 2 13Bのすべてのベンチマークで優れた性能を発揮
多くのベンチマークでLlama 1 34Bを上回る性能
英語のタスクでの良好な性能を維持しながら、コードタスクでCodeLlama 7Bに匹敵する性能
Grouped-query attention (GQA)を使用した高速な推論
Sliding Window Attention (SWA)を使用した低コストでの長いシーケンスの処理

こちらの記事もおすすめ

機械学習のモデル開発において、実験管理は非常に重要です。この記事では、Pytorchを用いたモデル学習の基本的な流れと、実験管理ツールであるWandbの使い方について解説します。必要なライブラリのインストールまず、必要なライブラリをインスト...

環境設定

まず、必要なライブラリをインストールし、環境を設定します。

!pip install --upgrade --force-reinstall --no-cache-dir torch==2.1.0 triton \
  --index-url https://download.pytorch.org/whl/cu121

import torch; torch.version.cuda

!pip install "unsloth[cu121] @ git+https://github.com/unslothai/unsloth.git"
!pip install wandb

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

データセットの準備

IMDbデータセットをロードし、テキストと対応するラベルを確認します。

import torch
from datasets import load_dataset

# Get dataset
dataset = load_dataset("imdb", split="train")

print(dataset["text"][0])  # テキストの例を表示
print(dataset["label"][0]) # ラベルの例を表示

モデルの準備

Mistral-7Bモデルをロードし、unslothを使用してファインチューニングのために準備します。

from unsloth import FastLanguageModel

max_seq_length = 2048  # RoPEスケーリングを内部でサポートしているため、任意の値を選択可能

# Mistralモデルをロード
model, tokenizer = FastLanguageModel.from_pretrained(
    model_name = "unsloth/mistral-7b-bnb-4bit", 
    max_seq_length = max_seq_length,
    dtype = None,
    load_in_4bit = True,
)

# モデルのパッチを適用し、高速なLoRA重みを追加
model = FastLanguageModel.get_peft_model(
    model,
    r = 16,
    target_modules = ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj",], 
    lora_alpha = 16,
    lora_dropout = 0,  # 任意の値をサポートしますが、0が最適化されています
    bias = "none",     # 任意の値をサポートしますが、"none"が最適化されています
    use_gradient_checkpointing = "unsloth",  # 4倍長いコンテキストウィンドウをサポート
    random_state = 3407,
    max_seq_length = max_seq_length,
)

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

SFTTrainerを使用して、準備したデータセットとモデルでファインチューニングを実行します。

from trl import SFTTrainer
from transformers import TrainingArguments

trainer = SFTTrainer(
    model = model,
    train_dataset = dataset,
    dataset_text_field = "text",
    max_seq_length = max_seq_length,
    tokenizer = tokenizer, 
    args = TrainingArguments(
      per_device_train_batch_size = 2,
      gradient_accumulation_steps = 4,
      warmup_steps = 10,
      max_steps = 60, 
      fp16 = not torch.cuda.is_bf16_supported(),
      bf16 = torch.cuda.is_bf16_supported(),   
      logging_steps = 1,
      output_dir = "outputs",
      optim = "adamw_8bit",
      seed = 3407,
      report_to="wandb",
  ),
)

trainer.train()