Zenn
🦙

UnslothでLlama3をファインチューニングする

2024/06/15に公開
LLM
Fine-tuning
Llama 3
unsloth
tech

はじめに

Unslothはより高速なLLMファインチューニングのための軽量ライブラリです。
Unslothを使う理由としては単に高速なだけでなく、Llama3公式推奨のFine-tuningフレームワークより低いスペックのマシンでも動作するところが良いからです。
この記事では、Unslothを使ってLlama3をFine-tuningする方法について説明します。

開発環境

  • WSL2 (Ubuntu 24.04)
  • Anaconda
  • CUDA 12.5

本記事で使うソースコード及びデータセットは、GitHub上で公開しています。

前準備

まずは環境構築です。正直ここが一番苦労でした。
自分でなんとかライブラリーをインストールしたりしないで、Anacondaを使って環境構築することをおすすめします。

CUDAのインストール

sudo apt-key del 7fa2af80
wget https://developer.download.nvidia.com/compute/cuda/repos/wsl-ubuntu/x86_64/cuda-wsl-ubuntu.pin
sudo mv cuda-wsl-ubuntu.pin /etc/apt/preferences.d/cuda-repository-pin-600
wget https://developer.download.nvidia.com/compute/cuda/12.5.0/local_installers/cuda-repo-wsl-ubuntu-12-5-local_12.5.0-1_amd64.deb
sudo dpkg -i cuda-repo-wsl-ubuntu-12-5-local_12.5.0-1_amd64.deb
sudo cp /var/cuda-repo-wsl-ubuntu-12-5-local/cuda-*-keyring.gpg /usr/share/keyrings/
sudo apt-get update
sudo apt-get -y install cuda-toolkit-12-5

# CUDAが正しくインストールされたか確認
nvidia-smi

Sat Jun 15 19:29:30 2024       
+-----------------------------------------------------------------------------------------+
| NVIDIA-SMI 555.52.01              Driver Version: 555.99         CUDA Version: 12.5     |
|-----------------------------------------+------------------------+----------------------+
| GPU  Name                 Persistence-M | Bus-Id          Disp.A | Volatile Uncorr. ECC |
| Fan  Temp   Perf          Pwr:Usage/Cap |           Memory-Usage | GPU-Util  Compute M. |
|                                         |                        |               MIG M. |
|=========================================+========================+======================|
|   0  NVIDIA GeForce RTX 4090        On  |   00000000:01:00.0  On |                  Off |
|  0%   43C    P8             17W /  450W |    8480MiB /  24564MiB |      4%      Default |
|                                         |                        |                  N/A |
+-----------------------------------------+------------------------+----------------------+

+-----------------------------------------------------------------------------------------+
| Processes:                                                                              |
|  GPU   GI   CI        PID   Type   Process name                              GPU Memory |
|        ID   ID                                                               Usage      |
|=========================================================================================|
|    0   N/A  N/A        26      G   /Xwayland                                   N/A      |
|    0   N/A  N/A    153395      C   /python3.10                                 N/A      |
+-----------------------------------------------------------------------------------------+

Anacondaのインストール、仮想環境の作成

wget https://repo.anaconda.com/archive/Anaconda3-2024.02-1-Linux-x86_64.sh
chmod +x Anaconda3-2024.02-1-Linux-x86_64.sh
./Anaconda3-2024.02-1-Linux-x86_64.sh

# Anacondaで仮想環境を作成する
conda create --name unsloth_env \
    python=3.10 \
    pytorch-cuda=12.1 \
    pytorch cudatoolkit xformers -c pytorch -c nvidia -c xformers \
    -y

conda activate unsloth_env

# 必要なライブラリをインストールする
pip install "unsloth[colab-new] @ git+https://github.com/unslothai/unsloth.git"

pip install --no-deps "trl<0.9.0" peft accelerate bitsandbytes

pip install jupyter gguf protobuf

ここからは、Jupyter Notebookを使って進めます。

jupyter notebook

データセットの準備

今回はkigner/ruozhiba-llama3-ttを基に日本語のデータセットを作成しました。
日本語のデータセットは、こちらからダウンロードできます。

Fine-tuning

ここからは、Jupyter Notebookを使ってFine-tuningを行います。

1. LLMの読み込み

from unsloth import FastLanguageModel
import torch
max_seq_length = 2048
dtype = None
load_in_4bit = True
model, tokenizer = FastLanguageModel.from_pretrained(
    model_name = "unsloth/llama-3-8b-bnb-4bit",
    max_seq_length = max_seq_length,
    dtype = dtype,
    load_in_4bit = load_in_4bit,
)

2. Fine-tuningする前に質問してみる

alpaca_prompt = """Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.

### Instruction:
{}
### Input:
{}
### Response:
{}"""

FastLanguageModel.for_inference(model)
inputs = tokenizer(
[
    alpaca_prompt.format(
        "日本語で回答してください", # instruction
        "赤壁の戦いで、曹操はなぜ119に電話して救助を求めなかったのですか？", # input
        "", # output
    )
], return_tensors = "pt").to("cuda")

from transformers import TextStreamer
text_streamer = TextStreamer(tokenizer)
_ = model.generate(**inputs, streamer = text_streamer, max_new_tokens = 128)

output:

<|begin_of_text|>Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.

### Instruction:
日本語で回答してください
### Input:
赤壁の戦いで、曹操はなぜ119に電話して救助を求めなかったのですか？
### Response:
<|end_of_text|>

まだFine-tuningしていないので、回答できないことがわかります。

3. データセットの読み込み

EOS_TOKEN = tokenizer.eos_token
def formatting_prompts_func(examples):
    instructions = examples["instruction"]
    inputs       = examples["input"]
    outputs      = examples["output"]
    texts = []
    for instruction, input, output in zip(instructions, inputs, outputs):
        text = alpaca_prompt.format(instruction, input, output) + EOS_TOKEN
        texts.append(text)
    return { "text" : texts }
pass

from datasets import load_dataset
dataset = load_dataset("json", data_files = "./dataset.json", split = "train")
dataset = dataset.map(formatting_prompts_func, batched = True)

4. パラメータの設定

from trl import SFTTrainer
from transformers import TrainingArguments

model = FastLanguageModel.get_peft_model(
    model,
    r = 16, # 8, 16, 32, 64, 128
    target_modules = ["q_proj", "k_proj", "v_proj", "o_proj",
                      "gate_proj", "up_proj", "down_proj",],
    lora_alpha = 16,
    lora_dropout = 0,
    bias = "none",
    use_gradient_checkpointing = "unsloth",
    random_state = 3407,
    use_rslora = False,
    loftq_config = None,
)

trainer = SFTTrainer(
    model = model,
    tokenizer = tokenizer,
    train_dataset = dataset,
    dataset_text_field = "text",
    max_seq_length = max_seq_length,
    dataset_num_proc = 2,
    packing = False,
    args = TrainingArguments(
        per_device_train_batch_size = 2,
        gradient_accumulation_steps = 4,
        warmup_steps = 5,
        max_steps = 60,
        learning_rate = 2e-4,
        fp16 = not torch.cuda.is_bf16_supported(),
        bf16 = torch.cuda.is_bf16_supported(),
        logging_steps = 1,
        optim = "adamw_8bit",
        weight_decay = 0.01,
        lr_scheduler_type = "linear",
        seed = 3407,
        output_dir = "outputs",
    ),
)

5. Fine-tuning

trainer_stats = trainer.train()

6. Fine-tuningしたモデルをテストする

FastLanguageModel.for_inference(model)
inputs = tokenizer(
    [
        alpaca_prompt.format(
            "日本語で回答してください", # instruction
            "赤壁の戦いで、曹操はなぜ119に電話して救助を求めなかったのですか？", # input
            "", # output
        )
    ], return_tensors = "pt").to("cuda")

from transformers import TextStreamer
text_streamer = TextStreamer(tokenizer)
_ = model.generate(**inputs, streamer = text_streamer, max_new_tokens = 128)

output:

<|begin_of_text|>Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.

### Instruction:
日本語で回答してください
### Input:
赤壁の戦いで、曹操はなぜ119に電話して救助を求めなかったのですか？
### Response:
三国時代には電話や現代の消防システムがなかったので、曹操は119に電話して救助を求めることはできませんでした。<|end_of_text|>

Fine-tuningしたモデルが正しく回答できることがわかります。

7. LoRAモデルを保存する

model.save_pretrained("lora_model")

8. モデルを結合し、4ビットのGGUFに量子化して保存する

model.save_pretrained_gguf("model", tokenizer, quantization_method = "q4_k_m")

まとめ

以上で、Unslothを使ってLlama3をFine-tuningする方法について説明しました。

Discussion