llama2のファインチューニング(QLORA)のメモ

Kan Hatakeyama

2023年7月20日 10:12

2023/11/13追記

以下の記事は､Llama2が公開されて数日後に書いた内容です｡
公開から数ヶ月経った23年11月時点では､諸々の洗練された方法が出てきていますので､そちらも参照されることをおすすめします｡

(以下､元記事です)
話題のLamma2をファインチューニングします｡
QLoRAライブラリを使うパターンと､公式推奨の2つを試しました｡前者が個人的にはオススメです｡

前提

Hugging faceで配布されている公式のモデルが必要です｡以下を参考に､ダウンロードしておきます｡

手段1. QLoRAライブラリを使用

こちらを参考に､QLoRA周りをセッティングします｡

データセット作成
(7/20 15:20追記設定ミスってたので修正しました)

test.jsonを適当に作ります｡

[
    {
        "input": "",
        "output": "### Human: 富士山といえば?### Assistant: なすび"
    },
    {
        "input": "",
        "output": "### Human: 明日の天気は?### Assistant: 雪"
    },
    {
        "input": "",
        "output": "### Human: 東京といえば?### Assistant: 神田川"
    },
    {
        "input": "",
        "output": "### Human: AIといえば?### Assistant: Llama2"
    }
]

学習コードを実行します｡パラメータは適当です｡

model_name=meta-llama/Llama-2-7b-chat-hf
dataset_name=dataset/json/test.json

python qlora.py \
    --model_name $model_name \
    --output_dir ./output/test\
    --dataset $dataset_name \
    --dataset_format input-output\
    --max_steps 1000 \
    --use_auth \
    --logging_steps 10 \
    --save_strategy steps \
    --data_seed 42 \
    --save_steps 50 \
    --save_total_limit 40 \
    --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 4096 \
    --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

重要なパラメータは以下の通りです｡
model_name : モデル名
output_dir : LoRA重みの保存先
dataset : データセットのパス

VRAM使用量の概算
1バッチあたり､7b / 2 (by 量子化) =3.5 GB
今回はbatch size=4なので､4倍して､3.5x4=14GB程度となります｡

推論

元のモデルの読み込み

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

#load base model

model_id = "meta-llama/Llama-2-7b-chat-hf"
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16
)

tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id, quantization_config=bnb_config, device_map={"":0})

LoRA重みの追加


#LoRAのパスを指定
peft_name = "/output/test/checkpoint-50/adapter_model/"
model = PeftModel.from_pretrained(
    model, 
    peft_name, 
    device_map={"":0}
)
model.eval()

テキスト生成

device = "cuda:0"
def ask(text):
  inputs = tokenizer(text, return_tensors="pt").to(device)

  with torch.no_grad():
    outputs = model.generate(**inputs, max_new_tokens=100)
  print(tokenizer.decode(outputs[0], skip_special_tokens=True))

text = "### Human: 富士山といえば ### Assistant: "
ask(text)
text = "### Human: 明日の天気は ### Assistant: "
ask(text)
text = "### Human: AIといえば ### Assistant: "
ask(text)

出力の確認

学習前

学習後

13, 70bの訓練(7/21追記)

追加設定が必要です｡
qlora.pyの647行目の下に､次のコードを加えればOKです｡

    model.config.pretraining_tp = 1 #追加!

それ以外は基本､同じです(コードを追記したファイルをqlora_4bit.pyとしました)｡

model_name=meta-llama/Llama-2-70b-chat-hf
dataset_name=dataset/json/test.json
out_dir=output/test
python qlora_4bit.py \
    --model_name $model_name \
    --output_dir $out_dir\
    --dataset $dataset_name \
    --dataset_format input-output\
    --max_steps 300 \
    --use_auth \
    --logging_steps 10 \
    --save_strategy steps \
    --data_seed 42 \
    --save_steps 10 \
    --save_total_limit 40 \
    --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 4096 \
    --per_device_train_batch_size 1 \
    --gradient_accumulation_steps 16 \
    --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

70bの学習の様子

(メモ: 24GB x2のVRAMでは､target_max_len = 768 with batch_size=4あたりがVRAMの限界でした｡max_lenがそれ以上大きいと､memory errorが出やすいです｡今回のサンプルはtarget_max_len = 4096となってますが､文章が短いので､実質的にはtarget_max_len = 20くらいです)

保存されるアダプターモデルのサイズは1.6GBでした｡流石の70bです｡

70bでの推論
こちらもモデル読み込みに追加設定が必要です｡

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


model_id = "meta-llama/Llama-2-70b-chat-hf"
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16
)
config = AutoConfig.from_pretrained(model_id)
config.pretraining_tp = 1
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id, quantization_config=bnb_config, 
                                             device_map="auto")

残りは今までと同じです｡checkpoint-2を読み込んでみます｡
今回の環境では､以下のコードの実行に5分以上を要しました｡

peft_name="output/0720test70b/checkpoint-2/adapter_model"
model = PeftModel.from_pretrained(
    model, 
    peft_name, 
    device_map="auto"
)
model.eval()

推論

device = "cuda:0"
def ask(text):
  inputs = tokenizer(text, return_tensors="pt").to(device)

  with torch.no_grad():
    outputs = model.generate(**inputs, max_new_tokens=100)
  print(tokenizer.decode(outputs[0], skip_special_tokens=True))

text = "### Human: 富士山といえば ### Assistant: "
ask(text)
text = "### Human: 明日の天気は ### Assistant: "
ask(text)
text = "### Human: AIといえば ### Assistant: "
ask(text)

10step後の学習結果はこちら｡70bモデルもうまく洗脳できました(4問目以外)｡

手段2. 公式のやり方

公式のとおりにやってみます｡

以下､7b-chatの例となります｡
ライブラリのクローン

pip install trl
git clone https://github.com/lvwerra/trl
cd trl
pip install -e .

サンプルスクリプトの実行
(学習はwandbの使用が前提のようなので､適当に登録します)


python examples/scripts/sft_trainer.py \
    --model_name meta-llama/Llama-2-7b-chat-hf \
    --dataset_name timdettmers/openassistant-guanaco \
    --load_in_4bit \
    --use_peft \
    --batch_size 4 \
    --gradient_accumulation_steps 2

RTX3090x2 (140W)の計算時間｡

メモリ使用量は以下のとおりです｡

sft_trainer.pyを少し修正します｡
sft_trainer.pyをコピーして､sft_trainer_original.pyを作りました｡88行目を変えて､jsonlファイルを読み込むようにします


#dataset = load_dataset(script_args.dataset_name, split="train")
dataset = load_dataset('json',data_files=script_args.dataset_name, split="train")

作成したpyとjsonlを指定して､実行します｡

python examples/scripts/sft_trainer_original.py \
    --model_name meta-llama/Llama-2-7b-chat-hf \
    --dataset_name dataset/test.jsonl\
    --load_in_4bit \
    --use_peft \
    --batch_size 4 \
    --gradient_accumulation_steps 2 \
    --output_dir peft_test

あまりにも学習データが少ないと､学習重みが保存されず､output_dir が空のままになる模様です｡データが保存されなかったので､loraモデルの読み込みなども試していません｡

この記事が気に入ったらサポートをしてみませんか？