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

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

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

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

mkdir llama3_torchtune
cd llama3_torchtune
nano Dockerfile

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

# ベースイメージ(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

# 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

# Torchtuneのインストール
RUN /app/.venv/bin/pip install --pre torchtune --extra-index-url https://download.pytorch.org/whl/nightly/cu121

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

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

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

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

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

nano docker-compose.yml

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

services:
  llama3_torchtune:
    build:
      context: .
      dockerfile: Dockerfile
    image: llama3_torchtune
    runtime: nvidia
    container_name: llama3_torchtune
    volumes:
      - .:/app/torchtune
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: 1
              capabilities: [gpu]

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

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

docker compose run llama3_torchtune

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

tune download meta-llama/Meta-Llama-3-8B-Instruct \
--output-dir /app/torchtune/models \
--hf-token '************'

Hugging Faceのアクセストークン作成方法は以下の記事で解説しています。

あわせて読みたい

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

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

テキスト生成の推論に使用するyamlファイルを作成します。

nano torchtune/inference_confg.yaml

# Model Arguments
model:
  _component_: torchtune.models.llama3.llama3_8b

checkpointer:
  _component_: torchtune.utils.FullModelMetaCheckpointer
  checkpoint_dir: /app/torchtune/models
  checkpoint_files: [/app/torchtune/models/original/consolidated.00.pth]
  recipe_checkpoint: null
  output_dir: /app/torchtune/output
  model_type: LLAMA3

device: cuda
dtype: bf16
seed: null

# Tokenizer arguments
tokenizer:
  _component_: torchtune.models.llama3.llama3_tokenizer
  path: /app/torchtune/models/original/tokenizer.model

# Generation arguments
prompt: "AirPodsとは何ですか"
max_new_tokens: 128
temperature: 0.8
top_k: 40
stop_tokens: "<|eot_id|>"

quantizer: null

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

テキスト生成のテストを実行します。

tune run generate --config torchtune/inference_confg.yaml

ファインチューニング前のモデルでは、英語での応答になっています。

_custom.pyを作成します。

nano .venv/lib/python3.10/site-packages/torchtune/datasets/_custom.py

from typing import Any, List, Mapping
from torchtune.data import Message
from torchtune.modules.tokenizers import Tokenizer
from torchtune.datasets._chat import ChatDataset

def message_converter(
    sample: Mapping[str, Any],
    train_on_input: bool = True,
    ) -> List[Message]:

    messages = [
        Message(role="user", content=sample["instruction"]),
        Message(role="assistant", content=sample["output"]),
    ]
    return messages

def custom_dataset(
    *,
    tokenizer: Tokenizer,
    source: str,
    chat_format: str = None,
    max_seq_len: int,
    train_on_input: bool = True,
    split: str,
) -> ChatDataset:

    return ChatDataset(
        tokenizer=tokenizer,
        source=source,
        convert_to_messages=message_converter,
        chat_format=chat_format,
        max_seq_len=max_seq_len,
        train_on_input=train_on_input,
        split=split,
    )

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

データセットの「__init__.py」ファイルを編集します。

nano .venv/lib/python3.10/site-packages/torchtune/datasets/__init__.py

「__init__.py」ファイルに次の２つのコードを追加します。

from torchtune.datasets._alpaca import alpaca_cleaned_dataset, alpaca_dataset
from torchtune.datasets._chat import chat_dataset, ChatDataset
from torchtune.datasets._cnn_dailymail import cnn_dailymail_articles_dataset
from torchtune.datasets._concat import ConcatDataset
from torchtune.datasets._grammar import grammar_dataset
from torchtune.datasets._instruct import instruct_dataset, InstructDataset
from torchtune.datasets._packed import PackedDataset
from torchtune.datasets._preference import PreferenceDataset
from torchtune.datasets._samsum import samsum_dataset
from torchtune.datasets._slimorca import slimorca_dataset
from torchtune.datasets._stack_exchanged_paired import stack_exchanged_paired_dataset
from torchtune.datasets._text_completion import (
    text_completion_dataset,
    TextCompletionDataset,
)
from torchtune.datasets._wikitext import wikitext_dataset
from torchtune.datasets._custom import custom_dataset #追加

__all__ = [
    "alpaca_dataset",
    "alpaca_cleaned_dataset",
    "grammar_dataset",
    "samsum_dataset",
    "stack_exchanged_paired_dataset",
    "InstructDataset",
    "slimorca_dataset",
    "ChatDataset",
    "instruct_dataset",
    "chat_dataset",
    "text_completion_dataset",
    "TextCompletionDataset",
    "cnn_dailymail_articles_dataset",
    "PackedDataset",
    "ConcatDataset",
    "wikitext_dataset",
    "PreferenceDataset",
    "custom_dataset",  #追加
]

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

ファインチューニングに使用するYAMLファイルを作成します。

nano torchtune/finetune_gozaru.yaml

以下の記述をコピーしてYAMLファイルに貼り付けます。

# Model Arguments
model:
  _component_: torchtune.models.llama3.qlora_llama3_8b
  lora_attn_modules: ["q_proj", "k_proj", "v_proj", "o_proj","gate_proj", "up_proj", "down_proj",]
  apply_lora_to_mlp: True
  apply_lora_to_output: True
  lora_rank: 8
  lora_alpha: 16

# Tokenizer
tokenizer:
  _component_: torchtune.models.llama3.llama3_tokenizer
  path: /app/torchtune/models/original/tokenizer.model

checkpointer:
  _component_: torchtune.utils.FullModelMetaCheckpointer
  checkpoint_dir: /app/torchtune/models/original/
  checkpoint_files: [
    consolidated.00.pth
  ]
  recipe_checkpoint: null
  output_dir: /app/torchtune/models/
  model_type: LLAMA3
resume_from_checkpoint: False

# Dataset and Sampler
dataset:
  _component_: torchtune.datasets.custom_dataset
  source: "bbz662bbz/databricks-dolly-15k-ja-gozaru"
  max_seq_len: 2048
  train_on_input: True
  split: train
seed: null
shuffle: True
batch_size: 8

# Optimizer and Scheduler
optimizer:
  _component_: torch.optim.AdamW
  weight_decay: 0.01
  lr: 2e-4
lr_scheduler:
  _component_: torchtune.modules.get_cosine_schedule_with_warmup
  num_warmup_steps: 100

loss:
  _component_: torch.nn.CrossEntropyLoss

# Training
epochs: 2
max_steps_per_epoch: null
gradient_accumulation_steps: 8
compile: False

# Logging
output_dir: /app/torchtune/output/
metric_logger:
  _component_: torchtune.utils.metric_logging.WandBLogger
  project: llama3_torchtune
log_every_n_steps: 1
log_peak_memory_stats: False

# Environment
device: cuda
dtype: bf16
enable_activation_checkpointing: True

# Profiler (disabled)
profiler:
  _component_: torchtune.utils.profiler
  enabled: False

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

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

wandb login ********

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

EdgeHUB

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

ファインチューニングを実行します。

tune run lora_finetune_single_device --config torchtune/finetune_gozaru.yaml

ファインチューニングには60分程度の時間がかかりました。

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

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

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