Google Colab で Llama-2-7B のQLoRA ファインチューニングを試す

「Google Colab」で「Llama-2-7B」のQLoRA ファインチューニングを試したので、まとめました。

前回

1. QLoRA と ござるデータセット

「QLoRA」のファインチューニングのスクリプトと、「ござるデータセット」(bbz662bbz/databricks-dolly-15k-ja-gozarinnemon) を使ってQLoRAファインチューニングに挑戦してみます。

bbz662bbz/databricks-dolly-15k-ja-gozarinnemon · Datasets at Hugging Face

2. Colabでの学習

Google Colabでの学習手順は、次のとおりです。

(1) メニュー「編集→ノートブックの設定」で、「ハードウェアアクセラレータ」で「GPU」を選択。
「T4」でも学習できますが、「A100」の方が時間が半分以下です。

(2) Googleドライブのマウント

# Googleドライブのマウント
from google.colab import drive
drive.mount("/content/drive")

(3) 作業フォルダへの移動

# 作業フォルダへの移動
import os
os.makedirs("/content/drive/My Drive/work", exist_ok=True)
%cd "/content/drive/My Drive/work"

(4) パッケージのインストール。

# パッケージのインストール
!git clone https://github.com/artidoro/qlora
%cd qlora
!pip install -U -r requirements.txt

(5) HuggingFaceのログイン。
リンクからHuggingFace Hubのトークンを取得し、「Token:」に入力してください。

# HuggingFaceのログイン
!huggingface-cli login

    _|    _|  _|    _|    _|_|_|    _|_|_|  _|_|_|  _|      _|    _|_|_|      _|_|_|_|    _|_|      _|_|_|  _|_|_|_|
    _|    _|  _|    _|  _|        _|          _|    _|_|    _|  _|            _|        _|    _|  _|        _|
    _|_|_|_|  _|    _|  _|  _|_|  _|  _|_|    _|    _|  _|  _|  _|  _|_|      _|_|_|    _|_|_|_|  _|        _|_|_|
    _|    _|  _|    _|  _|    _|  _|    _|    _|    _|    _|_|  _|    _|      _|        _|    _|  _|        _|
    _|    _|    _|_|      _|_|_|    _|_|_|  _|_|_|  _|      _|    _|_|_|      _|        _|    _|    _|_|_|  _|_|_|_|
    
    To login, `huggingface_hub` requires a token generated from https://huggingface.co/settings/tokens .
Token: 
Add token as git credential? (Y/n) n
Token is valid (permission: read).
Your token has been saved to /root/.cache/huggingface/token
Login successful

(6) 「qlora.py」の編集。
・642行目あたり
「ござるデータセット」の書式は「Alpaca」なので、Alpacaのダウンロード先を「bbz662bbz/databricks-dolly-15k-ja-gozarinnemon」に書き換えます。

        if dataset_name == "alpaca":
            return load_dataset("bbz662bbz/databricks-dolly-15k-ja-gozarinnemon")  # 書き換え

・679行目あたり
データセットの確認用に、先頭のデータを出力してます。

            dataset = dataset.map(
                extract_alpaca_dataset, remove_columns=["instruction"]
            )
            print("dataset:", dataset["train"][0])  # 追加

dataset: {'category': 'closed_qa', 'index': '0', 'output': '我、りんえもんは思う。 ヴァージン・オーストラリア航空は、2000年8月31日にヴァージン・ブルー航空として、2機の航空機で単一路線の運航を開始しましたでござる。知らんけど。', 'input': 'Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.\n\n### Instruction:\nヴァージン・オーストラリア航空はいつから運航を開始したのですか？\n\n### Input:\nヴァージン・オーストラリア航空（Virgin Australia Airlines Pty Ltd）はオーストラリアを拠点とするヴァージン・ブランドを冠する最大の船団規模を持つ航空会社です。2000年8月31日に、ヴァージン・ブルー空港として、2機の航空機、1つの空路を運行してサービスを開始しました。2001年9月のアンセット・オーストラリア空港の崩壊後、オーストラリアの国内市場で急速に地位を確立しました。その後はブリスベン、メルボルン、シドニーをハブとして、オーストラリア国内の32都市に直接乗り入れるまでに成長しました。\n\n### Response: '}

・820行目あたり
「model.config.pretraining_tp = 1」を追加します。

model.config.use_cache = False
model.config.pretraining_tp = 1  # 追加

・7B: config.pretraining_tp=1
・13B: config.pretraining_tp=2

(7) 学習の実行。
「A100」では1000ステップを20分で学習できます。
「T4」では「--bf16」未対応のため「--fp16」に変更して50分かかりました。

# 学習の実行
!python qlora.py \
    --model_name meta-llama/Llama-2-7b-hf \
    --output_dir "./output/test_peft" \
    --dataset "alpaca" \
    --max_steps 1000 \
    --use_auth \
    --logging_steps 10 \
    --save_strategy steps \
    --data_seed 42 \
    --save_steps 50 \
    --save_total_limit 40 \
    --max_new_tokens 32 \
    --dataloader_num_workers 1 \
    --group_by_length \
    --logging_strategy steps \
    --remove_unused_columns False \
    --do_train \
    --lora_r 64 \
    --lora_alpha 16 \
    --lora_modules all \
    --double_quant \
    --quant_type nf4 \
    --bf16 \
    --bits 4 \
    --warmup_ratio 0.03 \
    --lr_scheduler_type constant \
    --gradient_checkpointing \
    --source_max_len 16 \
    --target_max_len 512 \
    --per_device_train_batch_size 1 \
    --gradient_accumulation_steps 16 \
    --eval_steps 187 \
    --learning_rate 0.0002 \
    --adam_beta2 0.999 \
    --max_grad_norm 0.3 \
    --lora_dropout 0.1 \
    --weight_decay 0.0 \
    --seed 0 \
    --load_in_4bit \
    --use_peft \
    --batch_size 4 \
    --gradient_accumulation_steps 2

「./output/text_peft」にLoRAのチェックポイントが出力されます。

2. Colabでの推論

Google Colabでの推論手順は、次のとおりです。

(1) トークナイザーとモデルの読み込み。

import torch
from peft import PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

# トークナイザーとモデルの読み込み
tokenizer = AutoTokenizer.from_pretrained(
    "meta-llama/Llama-2-7b-hf"
)
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-2-7b-hf",
    quantization_config=BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_use_double_quant=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype=torch.bfloat16
    ),
    device_map={"":0}
)

(2) LoRAの読み込み。

# LoRAの読み込み
model = PeftModel.from_pretrained(
    model,
    "./output/test_peft/checkpoint-1000/adapter_model/",
    device_map={"":0}
)
model.eval()

(3) 推論の実行。

# プロンプトの準備
prompt = "### Instruction: 富士山とは？\n\n### Response: "

# 推論の実行
inputs = tokenizer(prompt, return_tensors="pt").to("cuda:0")
with torch.no_grad():
    outputs = model.generate(**inputs, max_new_tokens=100)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))

### Instruction: 富士山とは？

### Response: 富士山は、日本の山でござる。知らんけど。

他の質問応答例は、次のとおりです。

### Instruction: 侍とは？

### Response: 侍は、日本の武士のことでござる。知らんけど。

### Instruction: 忍者とは？

### Response: 忍者は、日本の伝統的な武器を使って戦うことができる人でござる。知らんけど。

### Instruction: まどか☆マギカとは？

### Response: まどか☆マギカは、2007年にアニメーション制作会社「A-1 Pictures」が製作したアニメーション作品でござる。

※ 正しくは「A-1 Pictures」→「シャフト」です

関連

llama2 13B LoRA

Fine-Tuning Llama-2: A Comprehensive Case Study for Tailoring Models to Unique Applications

次回

OpenAI GPT-4/ChatGPT/LangChain 人工知能プログラミング実践入門