Llama3.1のファインチューニング！UnslothでGPUメモリ削減！

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

mkdir llama31_unsloth
cd llama31_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]キーで編集モードから抜けます。

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

GitHub

GitHub – unslothai/unsloth: Finetune Llama 3.2, Mistral, Phi, Qwen 2.5 & Gemma LLMs 2-5x faster with… Finetune Llama 3.2, Mistral, Phi, Qwen 2.5 & Gemma LLMs 2-5x faster with 80% less memory – unslothai/unsloth

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

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

nano docker-compose.yml

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

services:
  llama31_unsloth:
    build:
      context: .
      dockerfile: Dockerfile
    image: llama31_unsloth
    runtime: nvidia
    container_name: llama31_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/Meta-Llama-3.1-70B-bnb-4bit",
    max_seq_length = max_seq_length,
    dtype = dtype,
    load_in_4bit = load_in_4bit,
)

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

モデルの学習に使用するデータセットを読み込みます。

dataset = load_dataset("fujiki/llm-japanese-dataset_wikinews", split="train")

データセットを変換するチャットテンプレートを定義します。

prompt_template = """
### Instruction:
{} 

### Input:
{} 

### Output:
{}"""

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

sample = dataset[0]
instruction = sample["instruction"]
input_text = sample["input"]
output = ""

FastLanguageModel.for_inference(model) 

inputs = tokenizer(
    [
        prompt_template.format(instruction, input_text, output)
    ], 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])

ファインチューニング前の応答結果

ファインチューニング前は、ニュースタイトルが生成できていないことが分かります！

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

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

dataset = dataset.map(formatting_prompts_func, batched = True,)
print(dataset[0]["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="llama31_unsloth")
trainer.train()
model.save_pretrained("llama31_unsloth/new_model")
tokenizer.save_pretrained("llama31_unsloth/new_model")

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

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

GPUメモリは、83％(64GB)を使用しています。

ファインチューニングには１時間ほどかかりました！