BLOOM-3BモデルとLoRAを用いたファインチューニング

Posted at 2024-09-16

BLOOM-3BモデルとLoRAを用いたファインチューニングチュートリアル

このチュートリアルでは、BLOOM-3Bモデルと**LoRA（Low-Rank Adaptation）**を使用して、JCommonsenseQAデータセットを用いたファインチューニングの手順を詳しく説明します。初心者のエンジニアの方でも理解しやすいよう、ステップバイステップで進めていきます。

1. 前提知識

1.1 LoRAとは

**LoRA（Low-Rank Adaptation）**は、大規模言語モデルの一部のパラメータのみを更新することで、メモリ使用量と計算コストを削減しつつ効率的にファインチューニングを行う手法です。

1.2 BLOOMモデルとは

BLOOMは、多言語対応の大規模言語モデルで、日本語を含む複数の言語に対応しています。BLOOM-3Bは約30億のパラメータを持ち、RTX 3090などのGPUで扱いやすいサイズです。

2. 環境の準備

2.1 必要なソフトウェアとライブラリ

Python 3.8以上
PyTorch
Transformers（Hugging Face）
Datasets（Hugging Face）
Accelerate
PEFT（Parameter-Efficient Fine-Tuning）
bitsandbytes（8bit量子化のため）

2.2 仮想環境の設定

仮想環境の作成とアクティベート

python -m venv bloom_lora_env

Linux/Macの場合:
```
source bloom_lora_env/bin/activate
```
Windowsの場合:
```
bloom_lora_env\Scripts\activate
```

ライブラリのインストール

pip install --upgrade pip
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
pip install transformers datasets accelerate sentencepiece peft bitsandbytes

注意: bitsandbytesは8bit量子化を行うために必要です。

3. データセットの準備

3.1 JCommonsenseQAデータセットの取得

from datasets import load_dataset

dataset = load_dataset("leemeng/jcommonsenseqa-v1.1")

データセットにはtrain、validation、testのスプリットがあります。

3.2 データの確認

print(dataset['train'][0])

出力例:

{
    'q_id': 0,
    'question': '主に子ども向けのもので、イラストのついた物語が書かれているものはどれ？',
    'choice0': '世界',
    'choice1': '写真集',
    'choice2': '絵本',
    'choice3': '論文',
    'choice4': '図鑑',
    'label': 2
}

4. モデルの準備

4.1 Hugging Faceへのログイン

セキュリティ上の注意: トークンは環境変数や.envファイルを使用して安全に管理してください。

import os
from huggingface_hub import login

HUGGINGFACE_TOKEN = os.getenv('HUGGINGFACE_TOKEN')
login(token=HUGGINGFACE_TOKEN)

トークンの取得方法: Hugging Faceのアカウントページからアクセストークンを取得します。

4.2 モデルとトークナイザーのロード

from transformers import AutoTokenizer, AutoModelForCausalLM

model_name = "bigscience/bloom-3b"
tokenizer = AutoTokenizer.from_pretrained(model_name)

model = AutoModelForCausalLM.from_pretrained(
    model_name,
    load_in_8bit=True,        # 8bit量子化でメモリを節約
    device_map='auto'         # デバイスを自動設定
)

注意: load_in_8bit=Trueを使用するため、bitsandbytesライブラリが必要です。

5. データの前処理

5.1 データの整形

データをモデルが理解できる形式に整形します。

def format_example(example):
    question = example['question']
    choices = [example[f'choice{i}'] for i in range(5)]  # choice0からchoice4まで取得
    label_index = int(example['label'])  # 正解のインデックス（0から4）

    # 選択肢を整形
    options = '\n'.join([f"{chr(65 + i)}: {choice}" for i, choice in enumerate(choices)])
    prompt = f"質問: {question}\n{options}\n答えの選択肢を一つ選んでください。"

    # 正解のテキストを取得
    answer_text = choices[label_index]

    return {'prompt': prompt, 'label': answer_text}

formatted_dataset = dataset.map(format_example)

5.2 トークナイズ

def tokenize_function(example):
    # 質問と回答を結合
    full_text = example['prompt'] + "\n答え: " + example['label']
    
    # 全体をトークナイズ
    tokenized = tokenizer(
        full_text,
        max_length=512,
        truncation=True,
        padding='max_length',
    )
    
    # ラベルの作成
    labels = tokenized['input_ids'].copy()
    
    # プロンプト部分をマスク
    prompt_length = len(tokenizer(example['prompt'])['input_ids'])
    labels[:prompt_length] = [-100] * prompt_length  # -100は損失計算から除外されるトークンを示す

    tokenized['labels'] = labels
    return tokenized

tokenized_dataset = formatted_dataset.map(tokenize_function, batched=False)

説明:
- -100を使用して、モデルが回答部分のみを学習するように設定します。

6. LoRAを用いたファインチューニング

6.1 LoRAの設定

from peft import LoraConfig, get_peft_model

lora_config = LoraConfig(
    r=8,
    lora_alpha=32,
    target_modules=["query_key_value"],  # BLOOMモデルの適切なモジュール
    lora_dropout=0.1,
    bias="none",
    task_type="CAUSAL_LM"
)

target_modulesの確認: モデル内のモジュール名を確認することで、適切なターゲットを指定できます。
```
for name, module in model.named_modules():
    if 'query_key_value' in name:
        print(name)
```

6.2 モデルにLoRAを適用

model = get_peft_model(model, lora_config)
model.print_trainable_parameters()