QLoRAで遊ぶメモ
このメモを読むと
・QLoRAを導入できる
・でかいローカルLLMモデルをファインチューニングできる
検証環境
・Windows11
・VRAM24GB
・ローカル（Anaconda）
・2023/6/M時点
事前準備
Anacondaを使うメモ｜おれっち (note.com)
Gitを使うメモ｜おれっち (note.com)
QLoRAとは
VRAMを節約することで容量の大きいローカルLLMモデルをファインチューニングできるようになる技術。
実装してデカモデルをファインチューニングしてみます。
環境構築
とても簡単です！
1. 仮想環境を作成し、環境切替
conda create -n qloratest python=3.10
activate qloratest2. 追加パッケージのインストール
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
pip install https://github.com/jllllll/bitsandbytes-windows-webui/raw/main/bitsandbytes-0.39.0-py3-none-any.whl
pip install git+https://github.com/huggingface/accelerate.git
pip install git+https://github.com/huggingface/transformers.git
pip install git+https://github.com/huggingface/peft.git
pip install datasets sentencepiece protobuf==3.20.03. 以下からCUDA Toolkit11.8を導入します。（未導入の場合）
完了です！
現時点ではaccelerateとtransformersはdev版でなければ動作しません。
うまくいかない場合は新規仮想環境で再度インストールしてください。
QLoRAを使ってみる
QLoRAでファインチューニングします。
用意するものは以下の二つです。
・ベースモデル
・データセット（学習データ）
それぞれこちらをお借りしました。
ベースモデル：cyberagent/open-calm-7b
データセット：bbz662bbz/databricks-dolly-15k-ja-gozaru
では実際に学習とテストをしてみましょう。
学習
好きな名前で下記スクリプトを作成し実行します。
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, DataCollatorForLanguageModeling, Trainer, TrainingArguments
from peft import get_peft_model, LoraConfig, TaskType, prepare_model_for_kbit_training
from datasets import load_dataset
# 基本パラメータ
base_model = "cyberagent/open-calm-7b"
dataset = "bbz662bbz/databricks-dolly-15k-ja-gozaru"
peft_name = "test-Ocalm-7b"
output_dir = "test-Ocalm-7b-result"
# トレーニング用パラメータ
eval_steps = 100
save_steps = 200
logging_steps = 20
epochs = 3
max_steps = 150 # 0にするとepochsに応じて自動設定
# LoRA用パラメータ
lora_r = 8
lora_alpha = 16
lora_dropout = 0.1
# 他パラメータ
VAL_SET_SIZE = 0.2 # 検証分割比率
CUTOFF_LEN = 512 # コンテキスト長の上限
# ベースモデル量子化パラ設定
bnb_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_use_double_quant=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch.bfloat16,
)
# ベースモデルの読み込み
model = AutoModelForCausalLM.from_pretrained(
base_model,
quantization_config=bnb_config,
device_map="auto"
)
# Rinnaのトークナイザーでは、「use_fast=False」も必要になる
tokenizer = AutoTokenizer.from_pretrained(base_model)
model.gradient_checkpointing_enable()
model = prepare_model_for_kbit_training(model)
# PEFT(LoRA)の設定
config = LoraConfig(r=lora_r,
lora_alpha=lora_alpha,
lora_dropout=lora_dropout,
inference_mode=False,
task_type=TaskType.CAUSAL_LM,
)
model = get_peft_model(model, config)
model.print_trainable_parameters() # 学習可能パラメータの確認
# トークナイズ関数
def tokenize(prompt, tokenizer):
result = tokenizer(prompt,
truncation=True,
max_length=CUTOFF_LEN,
padding=False,
)
return {"input_ids": result["input_ids"],
"attention_mask": result["attention_mask"],
}
# データセットの準備
data = load_dataset(dataset)
# プロンプトテンプレートの準備
# 回答の最後にベースモデル特有のEOSトークンを挿入
eos_token = tokenizer.decode([tokenizer.eos_token_id])
def generate_prompt(data_point):
if data_point["input"]:
result = f"""ユーザー:{data_point["instruction"]}
入力:{data_point["input"]}
システム:{data_point["output"]}{eos_token}"""
else:
result = f"""ユーザー:{data_point["instruction"]}
システム:{data_point["output"]}{eos_token}"""
result = result.replace('\n', '<NL>') # 改行→<NL>
return result
# 学習データと検証データの準備
train_val = data["train"].train_test_split(test_size=VAL_SET_SIZE, shuffle=True, seed=42)
train_data = train_val["train"]
val_data = train_val["test"]
train_data = train_data.shuffle().map(lambda x: tokenize(generate_prompt(x), tokenizer))
val_data = val_data.shuffle().map(lambda x: tokenize(generate_prompt(x), tokenizer))
# トレーナーの定義
trainer = Trainer(
model = model,
train_dataset = train_data,
eval_dataset = val_data,
args = TrainingArguments(
num_train_epochs=epochs,
learning_rate=3e-4,
logging_steps=logging_steps,
evaluation_strategy="steps",
save_strategy="steps",
max_steps=max_steps,
eval_steps=eval_steps,
save_steps=save_steps,
output_dir=output_dir,
report_to="none",
save_total_limit=3,
push_to_hub=False,
auto_find_batch_size=True
),
data_collator= DataCollatorForLanguageModeling(tokenizer, mlm=False),
)
model.config.use_cache = False # 警告を黙らせます
trainer.train()
model.config.use_cache = True # 推論のために再度有効化
trainer.model.save_pretrained(peft_name) # LoRAモデルの保存
print("Done!")お試しなので訓練ステップ"max_steps"を150に設定しています。
0にするとepochs = 3を元に本来の回数に自動設定されます。
推論
学習結果をテストします。
好きな名前で下記スクリプトを作成し実行します。
import time
import torch
from peft import PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
# 基本パラメータ
base_model = "cyberagent/open-calm-7b"
peft_name = "test-Ocalm-7b"
# 入力文章
input_text = "カレーにジャガイモは入れるべき？"
# ベースモデル量子化パラ設定
bnb_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_use_double_quant=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch.bfloat16
)
# ベースモデルの読み込み
model = AutoModelForCausalLM.from_pretrained(
base_model,
quantization_config=bnb_config,
device_map='auto',
)
# Rinnaのトークナイザーでは、「use_fast=False」も必要になる
tokenizer = AutoTokenizer.from_pretrained(base_model)
# PEFT(LoRA)の読み込み
model = PeftModel.from_pretrained(model, peft_name)
model.eval()# 評価モード
# プロンプトテンプレートの準備
def generate_prompt(data_point):
result = f"""ユーザー:{data_point["instruction"]}
システム:"""
result = result.replace('\n', '<NL>') # 改行→<NL>
return result
# テキスト生成関数の定義
def textgen(instruction,input=None,maxTokens=512):
prompt = generate_prompt({'instruction':instruction,'input':input})
input_ids = tokenizer(prompt,
return_tensors="pt",
truncation=True,
add_special_tokens=False).input_ids.cuda()
outputs = model.generate(input_ids=input_ids,
max_new_tokens=maxTokens,
do_sample=True,
temperature=0.7,
top_p=0.75,
top_k=40,
no_repeat_ngram_size=2,
)
outputs = outputs[0].tolist()
#print(tokenizer.decode(outputs))
# EOSトークンにヒットしたらデコード完了
if tokenizer.eos_token_id in outputs:
eos_index = outputs.index(tokenizer.eos_token_id)
decoded = tokenizer.decode(outputs[:eos_index])
# レスポンス内容のみ抽出
sentinel = "システム:"
sentinelLoc = decoded.find(sentinel)
if sentinelLoc >= 0:
result = decoded[sentinelLoc+len(sentinel):]
return(result.replace("<NL>", "\n")) # <NL>→改行
else:
return('Warning: Expected prompt template to be emitted. Ignoring output.')
else:
return('Warning: no <eos> detected ignoring output')
start_time = time.time()
print(f"Q: {input_text}\nA: {textgen(input_text)}")
end_time = time.time()
elapsed_time = end_time - start_time
print(f"Elapsed time: {elapsed_time} seconds")Q: カレーにジャガイモは入れるべき？
A: はいでござる。で、ござります。
Elapsed time: 2.060875654220581 seconds
VRAM比較
LoRA （過去記事の方法）
QLoRA (本記事の方法)
VRAM節約できてそうですね。
おわり
VRAMを節約してデカモデルをファインチューニングできました。
チューニングの成果物は下記記事でマージできます。
参考資料
・OpenCALM-7BをLoRAでinstruction tuningするための実装解説 / QLoRAの実装も紹介 - Qiita