【Llama3】Unslothで爆速ファインチューニング(QLoRA)

Dockerfileにインストールするパッケージを記述します。

Unsloth、CUDA、PyTorch等は、バージョン依存関係によるエラーが起きやすいので、動作検証済のバージョン指定してインストールしています。

特にUnslothには、細かいバージョンの組み合わせがありますので、詳細はGithubをご参照ください。

GitHub

GitHub - unslothai/unsloth: Finetune Llama 3.1, Mistral, Phi & Gemma LLMs 2-5x faster with 80% less ... Finetune Llama 3.1, Mistral, Phi & Gemma LLMs 2-5x faster with 80% less memory - unslothai/unsloth

Ubuntuのコマンドラインから、Dockerfileを作成します。

mkdir llama3_unsloth
cd llama3_unsloth
nano Dockerfile

次の記述をコピーしてDockerfileに貼り付けます。

# ベースイメージ(CUDA)の指定
FROM nvidia/cuda:12.1.0-cudnn8-devel-ubuntu22.04

# 必要なパッケージをインストール
RUN apt-get update && apt-get install -y python3-pip python3-venv git nano

# 作業ディレクトリを設定
WORKDIR /app

# アプリケーションコードをコピー
COPY . /app

# Python仮想環境の作成
RUN python3 -m venv /app/.venv

# 仮想環境をアクティベートするコマンドを.bashrcに追加
RUN echo "source /app/.venv/bin/activate" >> /root/.bashrc

# JupyterLab,HuggingFaceHub,WandBのインストール
RUN /app/.venv/bin/pip install Jupyter \
    jupyterlab \
    huggingface_hub[cli] \
    wandb

# PyTorch,Tritonのインストール
RUN /app/.venv/bin/pip install torch==2.2.0 torchvision==0.17.0 torchaudio==2.2.0 triton --index-url https://download.pytorch.org/whl/cu121

# Wheelのインストール
RUN /app/.venv/bin/pip install wheel

#Unslothをインストール 
RUN /app/.venv/bin/pip install "unsloth[cu121-ampere-torch220] @ git+https://github.com/unslothai/unsloth.git"

# コンテナの起動時にbashを実行
CMD ["/bin/bash"]

[Ctrl + S]キーで変更内容を保存し、[Ctrl + X]キーで編集モードから抜けます。

docker-compose.ymlファイルを使ってDockerコンテナの設定をします。

docker-compose.ymlファイルを作成します。

nano docker-compose.yml

次の記述をコピーしてdocker-compose.ymlに貼り付けます。

services:
  llama3_unsloth:
    build:
      context: .
      dockerfile: Dockerfile
    image: llama3_unsloth
    runtime: nvidia
    container_name: llama3_unsloth
    ports:
      - "8888:8888"
    volumes:
      - .:/app/unsloth
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: 1
              capabilities: [gpu]
    command: >
      bash -c '/app/.venv/bin/jupyter lab --ip="*" --port=8888 --NotebookApp.token="" --NotebookApp.password="" --no-browser --allow-root'

[Ctrl + S]キーで変更内容を保存し、[Ctrl + X]キーで編集モードから抜けます。

Dockerfileからビルドしてコンテナを起動します。   

docker compose up

Dockerの起動後にブラウザの検索窓に”localhost:8888″を入力すると、Jupyter Labをブラウザで表示できます。

localhost:8888

必要なライブラリをインポートします。

from unsloth import FastLanguageModel, is_bfloat16_supported
from unsloth.chat_templates import get_chat_template
import torch
from datasets import load_dataset
from trl import SFTTrainer
from transformers import TrainingArguments, AutoTokenizer
from peft import AutoPeftModelForCausalLM
from huggingface_hub import login
import wandb

モデルをダウンロードするために、HuggingFaceにログインします。

token = "*************"
login(token)

token = “*************”には、HuggingFaceで発行したアクセストークンが入ります。

HuggingFaceでアクセストークンを発行する方法は以下の記事で解説しています。

あわせて読みたい

Hugging Faceの使い方！アクセストークン作成からログインまで この記事では、Hugging Faceアカウントの作り方やコマンドラインからのログイン方法を紹介しています。 Hugging Faceを使うことで、生成AIに必要なモデルやデータセット…

学習ログの管理をするために、WandBにログインします。（WandBを使用しない場合は省略してください。）

API_KEY = "*************"
wandb.login(key=API_KEY)

API_KEY = “*************”には、Wandbで発行したAPIキーが入ります。

WandbでAPIキーを発行する方法を以下の記事で解説しています。

あわせて読みたい

WandBの使い方！ログイン方法からAPIキー発行、ログ取得まで紹介 WandBは機械学習のログ管理に便利なツールです。 WandBを利用することで、モデル学習のログ収集やグラフ表示が簡単にできるようになります。 この記事では、WandBのアカ…

モデルとトークナイザーのダウンロードをして読み込みます。

max_seq_length = 2048
dtype = torch.bfloat16
load_in_4bit = True

model, tokenizer = FastLanguageModel.from_pretrained(
    model_name = "unsloth/llama-3-70b-bnb-4bit",
    max_seq_length = max_seq_length,
    dtype = dtype,
    load_in_4bit = load_in_4bit,
)

モデルの読み込み時点でGPUメモリを消費します。

ファインチューニングをする前のモデルでテキスト生成のテストをしてみます。

alpaca_prompt = """
### Instruction:
{} 

### Response:
{}"""

FastLanguageModel.for_inference(model) 

inputs = tokenizer(
    [
        alpaca_prompt.format(
        "観葉植物の効果とは？",
        "",
        )
    ], return_tensors = "pt").to("cuda")

terminators = [
    tokenizer.eos_token_id,
    tokenizer.convert_tokens_to_ids("<|eot_id|>")
]

outputs = model.generate(
    **inputs, 
    max_new_tokens = 128,
    use_cache = True,
    eos_token_id=terminators,
    temperature=0.6,
    top_p=0.9
)
decoded_outputs = tokenizer.batch_decode(outputs,skip_special_tokens=True)
print(decoded_outputs[0])

ファインチューニング前のモデルでは、日本語の標準的な口調になっています。

モデルの学習に使用するデータセット「bbz662bbz/databricks-dolly-15k-ja-gozarinnemon」を読み込みます。

dataset = load_dataset("bbz662bbz/databricks-dolly-15k-ja-gozarinnemon", split="train")

データセットから「instruction（指示）」と「output（応答）」を抽出して、チャットテンプレートに変換して、データセットを更新します。

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

dataset = dataset.map(formatting_prompts_func, batched = True,)
print(dataset[246]["text"])

LoRAのパラメータ設定をします。

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

学習パラメータを設定しています。

training_arguments = TrainingArguments(
        bf16 = True,
        group_by_length=True,
        per_device_train_batch_size = 16,
        gradient_accumulation_steps = 4,
        num_train_epochs=3,
        warmup_steps = 100,
        max_steps = 50,
        learning_rate = 2e-4,
        logging_steps = 1,
        output_dir = "outputs",
        optim = "adamw_8bit",
        report_to="wandb" 
    )

SFTTrainerを使って教師ありファインチューニング（Supervised Fine-tuning）を実行します。

trainer = SFTTrainer(
    model = model,
    train_dataset = dataset,
    dataset_text_field = "text",
    max_seq_length = max_seq_length,
    tokenizer = tokenizer,
    args = training_arguments,
)

wandb.init(project="llama3_unsloth")
trainer.train()
model.save_pretrained("llama3_unsloth/new_model")
tokenizer.save_pretrained("llama3_unsloth/new_model")

WandBに保存したファインチューニング実行中のメトリクスを確認します。

Loss（損失）は最初の60ステップで急速に減少し、その後は緩やかに小さくなり収束しています。

GPUメモリは、50%～70％(40GB～56GB)を使用しています。