Google Colab + trl で Falcon-7B のQLoRAファインチューニングを試す
「Google Colab」で「Falcon-7B」のQLoRAファインチューニングを試したのでまとめました。
1. trl v0.4.2
「QLoRA」(bitsandbytesによる4bit量子化)を使用した大規模モデルの学習、新クラス「RewardTrainer」「SFTTrainer」を含むTRLの新バージョンで、「RLHF」をエンドツーエンドで簡単に実行可能になりました。
「trl v0.4.2」を使って「Falcon-7B」をQLoRAファインチューニングしてみます。サンプルColabも提供されてました。
2. Colabでの実行
Colabでの実行手順は、次のとおりです。
(1) メニュー「編集→ノートブックの設定」で、「ハードウェアアクセラレータ」で「GPU」で「A100」を選択。
(2) パッケージのインストール。
「SFTTrainer」を利用するため「accelerate」「peft」「transformers」「datasets」、「4-bit量子化」を利用するため「bitsandbytes」、「Falcon」を読み込むため「einops」をインストールします。
# パッケージのインストール
!pip install -q -U git+https://github.com/lvwerra/trl.git git+https://github.com/huggingface/transformers.git git+https://github.com/huggingface/accelerate.git git+https://github.com/huggingface/peft.git
!pip install -q datasets bitsandbytes einops wandb
(3) データセットの準備。
今回は、汎用チャットAIの学習に適した「OpenAssistant」のサブセットである「Guanaco」を使用します。
from datasets import load_dataset
# データセットの準備
dataset_name = "timdettmers/openassistant-guanaco"
dataset = load_dataset(dataset_name, split="train")
(4) モデルの準備。
「Falcon 7B」を読み込み、4-bit量子化(load_in_4bit=True)し、LoRAアダプタを接続します。
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, AutoTokenizer
# モデルの準備
model_name = "ybelkada/falcon-7b-sharded-bf16"
bnb_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch.float16,
)
model = AutoModelForCausalLM.from_pretrained(
model_name,
quantization_config=bnb_config,
trust_remote_code=True
)
model.config.use_cache = False
(5) トークナイザーの準備。
# トークナイザーの準備
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
tokenizer.pad_token = tokenizer.eos_token
(6) LoRAパラメータの準備。
「QLoRA」の論文によると、パフォーマンスを最大化するには、Transformerブロック内のすべてのlinear層を考慮することが重要です。そこで、mixed query key value層に加えて、ターゲットモジュールにdense、dense_h_to_4_h、dense_4h_to_h層を追加します。
from peft import LoraConfig
# LoRAパラメータの準備
lora_alpha = 16
lora_dropout = 0.1
lora_r = 64
peft_config = LoraConfig(
lora_alpha=lora_alpha,
lora_dropout=lora_dropout,
r=lora_r,
bias="none",
task_type="CAUSAL_LM",
target_modules=[
"query_key_value",
"dense",
"dense_h_to_4h",
"dense_4h_to_h",
]
)
(7) トレーナーパラメータの準備。
from transformers import TrainingArguments
# トレーナーパラメータの準備
output_dir = "./results"
per_device_train_batch_size = 4
gradient_accumulation_steps = 4
optim = "paged_adamw_32bit"
save_steps = 10
logging_steps = 10
learning_rate = 2e-4
max_grad_norm = 0.3
max_steps = 500
warmup_ratio = 0.03
lr_scheduler_type = "constant"
training_arguments = TrainingArguments(
output_dir=output_dir,
per_device_train_batch_size=per_device_train_batch_size,
gradient_accumulation_steps=gradient_accumulation_steps,
optim=optim,
save_steps=save_steps,
logging_steps=logging_steps,
learning_rate=learning_rate,
fp16=True,
max_grad_norm=max_grad_norm,
max_steps=max_steps,
warmup_ratio=warmup_ratio,
group_by_length=True,
lr_scheduler_type=lr_scheduler_type,
)
(8) 学習の実行。
transformersのトレーナーのラッパーを提供する「trl」の「SFTTrainer」を使用します。
from trl import SFTTrainer
# トレーナーの準備
max_seq_length = 512
trainer = SFTTrainer(
model=model,
train_dataset=dataset,
peft_config=peft_config,
dataset_text_field="text",
max_seq_length=max_seq_length,
tokenizer=tokenizer,
args=training_arguments,
)
(9) 「wandb」のAPIキーを求められるので入力します。
学習完了まで35分ほどかかりました。
(10) 「wandb」で、lossが収束されている (正常に学習されているっぽい) ことを確認します。
【おまけ】パラメータ
SFTTrainer のパラメータ
TrainingArguments のパラメータ
LoraConfig のパラメータ
関連
この記事が気に入ったらサポートをしてみませんか?