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

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

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

---

## **目次**

1. [前提知識](#前提知識)
2. [環境の準備](#環境の準備)
3. [データセットの準備](#データセットの準備)
4. [モデルの準備](#モデルの準備)
5. [データの前処理](#データの前処理)
6. [LoRAを用いたファインチューニング](#LoRAを用いたファインチューニング)
7. [モデルの保存と評価](#モデルの保存と評価)
8. [まとめ](#まとめ)
9. [参考リンク](#参考リンク)

---

<a name="前提知識"></a>

## **1. 前提知識**

### **1.1 LoRAとは**

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

### **1.2 BLOOMモデルとは**

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

---

<a name="環境の準備"></a>

## **2. 環境の準備**

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

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

### **2.2 仮想環境の設定**

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

   ```bash
   python -m venv bloom_lora_env
   ```

   - **Linux/Macの場合**:

     ```bash
     source bloom_lora_env/bin/activate
     ```

   - **Windowsの場合**:

     ```cmd
     bloom_lora_env\Scripts\activate
     ```

2. **ライブラリのインストール**

   ```bash
   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量子化を行うために必要です。

---

<a name="データセットの準備"></a>

## **3. データセットの準備**

### **3.1 JCommonsenseQAデータセットの取得**

```python
from datasets import load_dataset

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

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

### **3.2 データの確認**

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

**出力例**:

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

---

<a name="モデルの準備"></a>

## **4. モデルの準備**

### **4.1 Hugging Faceへのログイン**

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

```python
import os
from huggingface_hub import login

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

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

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

```python
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`ライブラリが必要です。

---

<a name="データの前処理"></a>

## **5. データの前処理**

### **5.1 データの整形**

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

```python
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 トークナイズ**

```python
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`を使用して、モデルが回答部分のみを学習するように設定します。

---

<a name="LoRAを用いたファインチューニング"></a>

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

### **6.1 LoRAの設定**

```python
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`の確認**: モデル内のモジュール名を確認することで、適切なターゲットを指定できます。

  ```python
  for name, module in model.named_modules():
      if 'query_key_value' in name:
          print(name)
  ```

### **6.2 モデルにLoRAを適用**

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

- **出力例**:

  ```
  trainable params: 2,457,600 || all params: 3,005,015,040 || trainable%: 0.0818
  ```

### **6.3 トレーニングパラメータの設定**

```python
from transformers import Trainer, TrainingArguments

training_args = TrainingArguments(
    output_dir="./bloom-lora-finetuned",
    per_device_train_batch_size=1,
    gradient_accumulation_steps=16,
    num_train_epochs=3,
    learning_rate=1e-4,
    fp16=True,
    save_steps=500,
    logging_steps=100,
    report_to="none"
)
```

### **6.4 データコラトの定義**

```python
import torch

def data_collator(features):
    return {
        'input_ids': torch.tensor([f['input_ids'] for f in features], dtype=torch.long),
        'attention_mask': torch.tensor([f['attention_mask'] for f in features], dtype=torch.long),
        'labels': torch.tensor([f['labels'] for f in features], dtype=torch.long),
    }
```

### **6.5 トレーナーの作成**

```python
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_dataset['train'],
    data_collator=data_collator,
)
```

### **6.6 トレーニングの開始**

```python
trainer.train()
```

---

<a name="モデルの保存と評価"></a>

## **7. モデルの保存と評価**

### **7.1 モデルの保存**

```python
model.save_pretrained("./bloom-lora-finetuned")
tokenizer.save_pretrained("./bloom-lora-finetuned")
```

### **7.2 モデルの評価**

```python
# テストデータで評価
test_results = trainer.evaluate(eval_dataset=tokenized_dataset['validation'])
print(test_results)
```

- **評価指標**: `loss`や`perplexity`など。

---

<a name="まとめ"></a>

## **8. まとめ**

このチュートリアルでは、以下の手順を踏んでBLOOM-3BモデルをLoRAを用いてファインチューニングしました。

1. **環境の準備**: 必要なライブラリと仮想環境を設定。
2. **データセットの準備**: JCommonsenseQAデータセットを取得し、データを確認。
3. **モデルの準備**: BLOOM-3Bモデルとトークナイザーをロード。
4. **データの前処理**: データをモデルが理解できる形式に整形し、トークナイズ。
5. **LoRAを用いたファインチューニング**: LoRAの設定を行い、モデルに適用。トレーナーを作成してトレーニングを実行。
6. **モデルの保存と評価**: ファインチューニング済みモデルを保存し、評価データセットで性能を確認。

---

<a name="参考リンク"></a>

## **9. 参考リンク**

- **LoRA論文**: [LoRA: Low-Rank Adaptation of Large Language Models](https://arxiv.org/abs/2106.09685)
- **Hugging Face PEFTライブラリ**: [https://github.com/huggingface/peft](https://github.com/huggingface/peft)
- **BLOOMモデルカード**: [https://huggingface.co/bigscience/bloom](https://huggingface.co/bigscience/bloom)
- **JCommonsenseQAデータセット**: [https://huggingface.co/datasets/leemeng/jcommonsenseqa-v1.1](https://huggingface.co/datasets/leemeng/jcommonsenseqa-v1.1)

---

## **補足情報**

- **APIトークンの管理**: セキュリティのため、APIトークンはコード内に直接記載せず、環境変数や`.env`ファイルを使用して安全に管理してください。
- **メモリエラーの対処**: メモリエラーが発生した場合、`per_device_train_batch_size`を減らすか、`gradient_accumulation_steps`を増やすことで対処できます。
- **モデルの評価**: 必要に応じて、カスタムの評価関数を実装してモデルの性能を詳細に分析してください。