はじめに
以前に、文章生成 AI モデルのファインチューニングしてみました。
文章生成 AI のファインチューニングしてみた #LLM - Qiita
これの続きです。
生成 AI モデルのファインチューニングする
実行環境を用意する
AI プログラムの実行環境は、高速な計算するために大きなメモリや GPU を使います。そのため高額なマシンが必要になります。
Google Colab を使ってみたり、Google Compute Engine で仮想マシンを用意したりしました。
用意したマシンのスペックは、以下の通りです。
GPU：NVIDIA L4
GPU RAM：23GB
システム RAM：53GB
L4 プロセッサは、性能とメモリサイズに対してコストパフォマンスがいいですね。
AI モデルをファインチューニングする
以前に文章生成 AI モデルをファインチューニングしたときのコードです。↓
import transformers
# モデルとトークナイザの準備
model = transformers.AutoModelForCausalLM.from_pretrained(
"cyberagent/open-calm-large"
)
tokenizer = transformers.AutoTokenizer.from_pretrained(
"cyberagent/open-calm-large"
)
import datasets
# データセットを準備する
datadic = datasets.load_dataset("federerjiang/dialect.osaka")
dataset = datadic['train']
# データセットのトークン化
tokenized_dataset = dataset.map(
lambda example:
tokenizer(example["sentence"], truncation=True),
batched=True
)
# トレイナの準備
trainer = transformers.Trainer(
model=model,
data_collator=transformers.DataCollatorForLanguageModeling(
tokenizer,
mlm=False
),
args=transformers.TrainingArguments(
output_dir="./output",
num_train_epochs=5,
per_device_eval_batch_size=1,
),
train_dataset=tokenized_dataset
)
# トレーニングする
trainer.train()
# モデルを保存する
trainer.save_model("./trained_model")
モデルは cyberagent/open-calm-large にしました。モデルの本体を見ると 1.8 GB ありました。
モデルのサイズが性能の全てでありませんが、もう少し大きいサイズのモデルにするとどうでしょうか。
例えば line-corporation/japanese-large-lm-3.6b にしてみます。モデルの本体は 7.2 GB でした。
上記のコードのモデルを変更して実行してみます。
trainer.train() でエラーになりました。↓
OutOfMemoryError: CUDA out of memory.
Tried to allocate 18.00 MiB. GPU 0 has a total capacity of 22.17 GiB of which 16.88 MiB is free.
Process 27687 has 22.14 GiB memory in use. Of the allocated memory 21.71 GiB is allocated by PyTorch, and 203.24 MiB is reserved by PyTorch but unallocated.
GPU が持っているメモリを使い尽くしてしまったようです。
量子化と LoRA で対応する
調べてみると、量子化 と PEFT で対応できるようです。
- 量子化（quantization）
モデルのパラメータは浮動小数点数で記録されているが、そのビット数を減らすことでメモリ使用量を抑える。
参考：LLMを読み込む際にメモリを節約する対策「量子化」について #GGUF - Qiita
- PEFT（Parameter-Efficient Fine-Tuning）
モデルのパラメータの全てでなく、一部だけチューニングすることで計算コストを削減する。
PEFT 手法は幾つかあるが、LoRA を使うことにする。
- LoRA（low-rank adaptation（低ランク適応）
低ランク行列分解を使って調整するパラメータを削減する
参考：LLMを効率的に再学習する手法(PEFT)を解説 | DOORS DX
- QLoRA
量子化されたモデルを、LoRA でチューニングする
参考：GPUメモリが小さくてもパラメーター数が大きい言語モデルをトレーニング可能になる手法「QLoRA」が登場、一体どんな手法なのか？ - GIGAZINE
QLoRA でファインチューニングする
モデルを量子化する
モデルを読込するとき量子化します。
必要なライブラリをインストールする
モデルの量子化するのに bitsandbyts ライブラリを使います。
$ pip install accelerate
$ pip install bitsandbytes
モデルを量子化して読込する
モデルを読込するとき量子化する指定します。
import accelerate
import bitsandbytes
quantized_model = transformers.AutoModelForCausalLM.from_pretrained(
"line-corporation/japanese-large-lm-3.6b",
device_map="auto",
torch_dtype=torch.float16,
quantization_config=transformers.BitsAndBytesConfig( # 量子化する指定
load_in_4bit=True,
bnb_4bit_use_double_quant=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch.float16,
)
)
LoRA でモデルを準備し直す
必要なライブラリをインストールする
LoRA で処理するのに peft ライブラリを使います。
$ pip install peft
パラメータを追加する場所を調べる
LoRA は既存のモデルにパラメータを追加して、これをチューニングする。
モデルのどこにパラメータを追加できるか、以下のコードで分かるそうです。
import re
linear_layers = list(set(re.findall(
r'\((\w+)\): Linear', str(model.modules)
)))
print(linear_layers)
line-corporation/japanese-large-lm-3.6b なら、以下の通り。
['dense_4h_to_h', 'dense', 'dense_h_to_4h', 'embed_out', 'query_key_value']
PEFT できるようにモデルを加工する
PEFT で実行できるように get_peft_model で加工します。
import peft
peft_model = peft.get_peft_model(
quantizerd_model,
peft_confg=peft.LoraConfig( # LoRA の指定
r=4,
lora_alpha=8,
lora_dropout=0.05,
bias="none",
fan_in_fan_out=False,
target_modules = ['dense_4h_to_h', 'dense', 'dense_h_to_4h', 'embed_out', 'query_key_value'],
task_type="CAUSAL_LM"
)
)
チューニングする
上記のモデルを Trainer に渡すようにします。
# トレイナの準備
trainer = transformers.Trainer(
model=peft_model,
(以下略)
メモリ不足にならないで処理できました。