【Llama3】SFTTrainerで簡単ファインチューニング(QLoRA)

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

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

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

mkdir llama3_sftqlora
cd llama3_sftqlora
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のインストール
RUN /app/.venv/bin/pip install torch==2.2.2 torchvision==0.17.2 torchaudio==2.2.2 --index-url https://download.pytorch.org/whl/cu121

# Transformer関連のインストール
RUN /app/.venv/bin/pip install \
    transformers==4.41.2 \
    ninja \
    packaging \
    accelerate \
    datasets \
    bitsandbytes \
    evaluate \
    trl \
    peft \
    wheel

# Flash attentionをインストール
RUN /app/.venv/bin/pip install flash-attn --no-build-isolation

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

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

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

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

nano docker-compose.yml

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

services:
  llama3_sftqlora:
    build:
      context: .
      dockerfile: Dockerfile
    image: llama3_sftqlora
    runtime: nvidia
    container_name: llama3_sftqlora
    ports:
      - "8888:8888"
    volumes:
      - .:/app/llama3_sftqlora
    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

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

import torch
from torch import cuda, bfloat16
from transformers import (
    AutoTokenizer,
    AutoModelForCausalLM,
    BitsAndBytesConfig,
    HfArgumentParser,
    TrainingArguments,
    pipeline,
    logging
)
from datasets import load_dataset
from peft import LoraConfig, PeftModel
from trl import SFTTrainer
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キーを発行する方法を以下の記事で解説しています。

EdgeHUB

WandBの使い方！ログイン方法からAPIキー発行、ログ取得まで紹介 | EdgeHUB この記事では、WandBのアカウント作成から使い方まで解説していきます。 WandBは機械学習のログ管理に便利なツールです。 WandBを利用することで、モデル学習のログ収集や…

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

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16
)
model_id="meta-llama/Meta-Llama-3-8B-Instruct"
model = AutoModelForCausalLM.from_pretrained(
    model_id, 
    trust_remote_code=True,
    token=token,
    quantization_config=bnb_config,
    device_map='auto',
    torch_dtype=torch.bfloat16,
    attn_implementation="flash_attention_2"
)
tokenizer = AutoTokenizer.from_pretrained(
    model_id,
    padding_side="right",
    add_eos_token=True
)
if tokenizer.pad_token_id is None:
  tokenizer.pad_token_id = tokenizer.eos_token_id

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

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

messages = [
    {"role": "system", "content": "あなたは日本語で回答するアシスタントです。"},
    {"role": "user", "content": "GDPとは何ですか？"},
]

input_ids = tokenizer.apply_chat_template(
    messages,
    add_generation_prompt=True,
    return_tensors="pt"
).to(model.device)

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

outputs = model.generate(
    input_ids,
    max_new_tokens=256,
    eos_token_id=terminators,
    do_sample=True,
    temperature=0.6,
    top_p=0.9,
)
response = outputs[0][input_ids.shape[-1]:]

print(tokenizer.decode(response, skip_special_tokens=True))

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

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

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

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

def formatting_func(example):
        messages = [
            {'role': "system",'content': "あなたは日本語で回答するアシスタントです。"},
            {'role': "user",'content': example["instruction"]},
            {'role': "assistant",'content': example["output"]}
        ]
        return tokenizer.apply_chat_template(messages, tokenize=False)

def update_dataset(example):
    example["text"] = formatting_func(example)
    for field in ["index", "category", "instruction", "input", "output"]:
        example.pop(field, None)
    return example

dataset = dataset.map(update_dataset)

print(dataset[293]["text"])

モデル学習の前には、データセットをチャットテンプレートに変換しておく必要があります。

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

peft_config = LoraConfig(
    r=8,
    lora_alpha=16,
    lora_dropout=0.05,
    target_modules = ["q_proj", "k_proj", "v_proj", "o_proj","gate_proj", "up_proj", "down_proj",],
    bias="none",
    task_type="CAUSAL_LM",
    modules_to_save=["embed_tokens"],
)

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

training_arguments = TrainingArguments(
    output_dir="./results",
    bf16=True,
    per_device_train_batch_size=4,
    gradient_accumulation_steps=16,
    num_train_epochs=3,
    optim="adamw_torch_fused",
    learning_rate=2e-4, 
    lr_scheduler_type="cosine",
    weight_decay=0.01,
    warmup_steps=100,
    logging_steps=10,
    group_by_length=True,
    report_to="wandb" 
)

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

new_model="llama-3-8b-ft"

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

wandb.init(project="llama3_sftqlora")
trainer.train()
trainer.model.save_pretrained(new_model)

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

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

GPUメモリは、常に88%(70GB)を使用しています。