Sarashinaをファインチューニング【ソフトバンクの日本語LLM】

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

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

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

mkdir sarashina_instruct
cd sarashina_instruct
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 \
    sentencepiece

# 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:
  sarashina_instruct:
    build:
      context: .
      dockerfile: Dockerfile
    image: sarashina_instruct
    runtime: nvidia
    container_name: sarashina_instruct
    ports:
      - "8888:8888"
    volumes:
      - .:/app/sarashina_instruct
    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
import transformers
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline, logging
from datasets import load_dataset
from trl import SFTTrainer, SFTConfig
import wandb

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

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

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

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

EdgeHUB

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

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

model_id = "sbintuitions/sarashina2-7b"

model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.bfloat16,
    device_map="auto",
    attn_implementation="flash_attention_2"
)

tokenizer = AutoTokenizer.from_pretrained(
    model_id,
    padding_side="right",
    add_eos_token=True
)

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

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

prompt = "香川県のご当地グルメを教えて下さい"

pipeline = transformers.pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer
)

output = pipeline(
    prompt,
    max_length=256,
    truncation=True,
    do_sample=True,
    pad_token_id=tokenizer.pad_token_id,
    num_return_sequences=3,
)

print(output[0]['generated_text'])

Sarashinaの事前学習モデルでは、チャット形式の応答ができていないことが分かります。
事前学習モデルは人間の意図に沿った応答ができないため、指示チューニングが必要です。

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

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

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

template = "{% for message in messages %}{{'<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' + '\n'}}{% endfor %}"

tokenizer.chat_template = template

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[312]["text"])

チャットテンプレート変換後のデータセットの例

<|im_start|>system
あなたは日本語で回答するアシスタントです。<|im_end|>
<|im_start|>user
良い写真を撮るためのポイントとは？<|im_end|>
<|im_start|>assistant
良い写真には、光、被写体、構図という3つの重要な要素があります。優れた写真は、ピントが合っていて、シャープで、露出（光）がよく、構図が決まっています。  完璧な写真を撮るためには、正しい焦点距離で撮影し、フレーム内の被写体に適した絞りを設定し、撮影するアクションに合わせてシャッタースピードを設定し（スポーツでは高いシャッタースピード、ポートレートでは低いシャッタースピード）、シャッタースピードと絞りの両方に対応するようにISOが正しく設定されていることを確認する必要があります。  ISO、絞り、シャッタースピードがどのように連動しているかを理解することが、完璧な写真を撮るための鍵となります。<|im_end|>

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

<|im_start|>system...<|im_end|>:システムメッセージが入ります。
<|im_start|>user...<|im_end|>:プロンプト（指示）が入ります。
<|im_start|>assistant...<|im_end|>：レスポンス(応答)が入ります。  

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

sft_config = SFTConfig(
    output_dir="./results",
    bf16=True,
    per_device_train_batch_size=4,
    gradient_accumulation_steps=16,
    num_train_epochs=2,
    optim="adamw_torch_fused",
    learning_rate=2e-4,
    lr_scheduler_type="cosine",
    warmup_steps=100,
    weight_decay=0.01,
    logging_steps=10,
    save_strategy= "epoch", 
    group_by_length=True,
    dataset_text_field="text",
    max_seq_length=512,
    packing=True,
    report_to="wandb" 
)

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

new_model="sarashina2-7b-instruct"

trainer = SFTTrainer(
    model=model,
    tokenizer=tokenizer,
    train_dataset=dataset,
    args=sft_config,
)

wandb.init(project="sarashina2-7b-instruct")
trainer.train()
trainer.model.save_pretrained(new_model)

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

Loss（損失）は学習が進むにつれて小さくなっています。

GPUメモリは、ピーク91%(72GB)を使用しています。

学習は30分程度で完了しました