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

AI・機械学習

はじめに

この記事では、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を使ったモデル学習の基本
機械学習のモデル開発において、実験管理は非常に重要です。この記事では、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()

file

まとめ

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

Mistral-7Bは、そのサイズにおいて最も強力な言語モデルの1つであり、様々なベンチマークでLlama 2 13BやLlama 1 34Bを上回る性能を示しています。また、GQAとSWAを使用することで、高速な推論と長いシーケンスの効率的な処理を可能にしています。

unslothとMistral-7Bを組み合わせることで、高品質のファインチューニングモデルを効率的に作成することができます。ぜひ自分のプロジェクトでも試してみてください!

ノートブック

Google Colaboratory

参考サイト

Hugging Faceトランスフォーマー | Weights & Biases 선적 서류 비치
Hugging Faceトランスフォーマー ライブラリは、BERTのような最先端のNLPモデルや、混合精度演算や勾配チェックポイントのようなトレーニング手法を簡単に利用できるようにしています。W&Bインテグレーションは、使いやすさを損なうことなく、インタラクティブな中央集約ダッシュボードに豊富で柔軟な実験トラッキングと...
GitHub - unslothai/unsloth: 2-5X faster 80% less memory LLM finetuning
2-5X faster 80% less memory LLM finetuning. Contribute to unslothai/unsloth development by creating an account on GitHub.

コメント

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