📝

LLMに新しい情報の入れ方

2024/03/26に公開

個人ブログもやっています、もっと詳しくな説明はそちらに。アニメから会話データセットを作るためにもGithubをみてください。
大規模言語モデルに新しい情報の入れ方
animespeechdataset

新しい情報を取り入れる方法

今回使用したい手法は、モデルの仕組み を活かす方法です。ベースモデルは大規模なデータセットで学習されました。その学習は CausalLM と呼ばれます。そのタスクは次のトークンを予測することです。ラベルなしのデータセットは使用できますので、ベースモデルに 新しい知識 を与えることができます。効率的な学習を行うために、今回は QLoRA の手法を使用します。この方法は、量子化 されたモデル(今回は4ビット)の上にLoRAアダプターを付けます。LoRA学習 はバックボーンの重みを固定しますので、量子化の状態のままで、勾配計算を行わないため、LoRA学習が問題なく行えます。パラメータはランクとアルファを両方とも64に設定し、ドロップアウトは0.1に設定しました。LoRAに付けたモジュールはすべての全結合層でした。

この設定を用いて、最初は 指示チューニング のモデルを試してみました。指示チューニングはベースのモデルとの学習目的、および 損失関数が異なる ため、学習がうまくいかないことがありました、指示チューニングモデルの損失関数は生成するはずの文章の中に、次のトークンを予測。突然、損失が0になることがあり、学習済みLoRAを使用すると、モデルが奇妙 なことを言ったり、文が長くなったりしました。次のトークンを予測するタスクには指示チューニングのモデルは適していないようです。したがって、CausalLMタスクにおいて、最適なモデルはやはりベース のモデルでしょう。指示チューニングの学習方法とは異なり、やりたい学習はベースモデルと同じなので、最も精度が高くなりそうです。

非常に興味深い点は、ベースモデルのLoRA重み が指示チューニングの重みと干渉していないようで、モデルを壊さずに使用できることです。指示チューニングされたモデルの 構造はベースモデルと同じ ですが、ファインチューニングが行われているため、ベースモデルで学習されたLoRAは指示チューニングのモデルに使用できます。この指示に従う重みは、通常の知識とは異なるため、指示チューニングモデルに 新しい知識 を入れることができます。

実践

最初のステップは必要なライブラリをインストールすることです。今回はコラボ環境でファインチューニングを行います。

!pip install -U transformers Peft safetensors datasets
!pip install bitsandbytes
!pip install git+https://github.com/huggingface/accelerate.git

CausalLM学習において、次のトークンを予測するタスク では、指示チューニングのモデルは使用できません。なぜなら、それを使用するとモデルが 壊れてしまう可能性がある からです。そのため、ベースのモデルを使用します。QLoRA学習 においては、モデルを量子化して、4ビットバージョンを取得します。HuggingFace を使えば、量子化を簡単に行うことができます。今回は、CyberAgent が公開したモデルである Calm2 を使用します。最初の試行では、既にモデルが存在する場合は、そのモデルを削除してメモリを解放します。

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
import bitsandbytes as bnb
import accelerate
from transformers.trainer_utils import set_seed

# 乱数シードを42に固定
set_seed(42)

try:
  del model
  torch.cuda.empty_cache()
except:
  "No model"

# 量子化の設定
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.float16,
    bnb_4bit_use_double_quant=True,
)
model = AutoModelForCausalLM.from_pretrained(
    "cyberagent/calm2-7b",
    quantization_config=bnb_config,
    device_map={"":0},
    trust_remote_code=True,
    local_files_only=False,
    torch_dtype=torch.float16,
)

その後、トークナイザをロードし、paddingトークンを"end of sentence"トークンに変更します。これは、DataCollator がpaddingを追加する際に使用されますが、eosトークンは損失に影響を与えない ため、便利です。今回は同じ長さの文章を使用する予定ですが、DataCollatorのmlm引数はマスクを作成することを示しています。これは、Masked Language Modeling に使用され、エンコーダーモデルによく用いられる手法です。BertやRobertaなどがこれに該当します。

from transformers import DataCollatorForLanguageModeling, AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained("cyberagent/calm2-7b")

tokenizer.pad_token = tokenizer.eos_token
data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)

個人的にモデルを確認する方が良いと思います。今回は私のお気に入りのアニメに関する情報を入れたいので、情報なしの状態を確認しましょう。

torch.cuda.empty_cache()
model.eval()

text = """" 乙女ゲーム「Revolution」の世界に、"""

inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=512).to("cuda")
outputs = model.generate(**inputs,
                            max_new_tokens=256,
                            pad_token_id=tokenizer.pad_token_id,
                            eos_token_id=tokenizer.eos_token_id,
                            do_sample=True,
                            top_k=40,
                            temperature=0.7,
                            top_p=0.9,
                            repetition_penalty = 1.1,
                            ).cpu()
print(tokenizer.decode(outputs[0], skip_special_tokens=True))
print(torch.cuda.max_memory_allocated())
#乙女ゲーム「Revolution」の世界に、あなた自身が主人公として入り込んでしまう体験型コンテンツです。
#アニメイトタイムズ内の『RE:VOICE』特設ページでは、『Revolution』の世界観をイラストや小説でお楽しみいただけます。
#また、2015年9月6日より東京・池袋パルコにて開催される「オトメイト夏の市2015 ~桃色書店へようこそ~」内にて、本作品の世界観を再現したコラボカフェがオープンします!
#さらに、8月28日(金)には本作の主題歌を担当する『Rouge』『Mia Regina』の2組のヴォーカルユニットによるトークイベントも開催されますので、ぜひ遊びにいらしてください♪
#【開催期間】2015年7月30日(木)~2015年9月6日(日)
#【会場】東京・池袋パルコ 本館7F パークプラザ
#【営業時間】10:00~21:00 (フードLO 20:30/ドリンクLO 20:45)
#【料金】入場無料 ※フード、ドリンクをご利用の場合は別途必要となります。
#【公式HP】http://www.animatetimes.com/event/touring-project/revolution/
#4810984448

出力はこんな感じになりますが、あまり良くないですよね。別のアニメの話になってしまいます。最後の"max_memory_allocated()"はメモリの確認をしたいからです。注意点はその モデルのeval です。これは検証モードです。理由は変わらないけど、Peftモデルを変えた時に、"prepare_for_kbit_training()"はLoRA以外の重みを全部凍結 するはずですが、実際にそんなことはありませんでしたので、メモリ使用量が非常に高くなりました。Google Colabでは"out of memory"になりました。

7bモデルは大体5gbで使用できるため、今は 奇妙な勾配計算 が行われていません。モデルの確認をした後、LoRAモジュールを作成します。QLoRAのスクリプト から"find_all_linear_names"の関数を使用します。今は非常に便利だと思います。Peftのライブラリーでは"all-linear"を使用すると同じ効果が期待できます。今回は rankとalphaを両方とも64 に設定します。この数は大きいかもしれません。
"EFFICIENT AND EFFECTIVE TEXT ENCODING FOR CHINESE LLAMA AND ALPACA"
の論文では、Llama が中国語を理解できるようにするために、ファインチューニングが行われます。彼らは事前学習の時に、rank=8とalpha=32 を使用します。その後、指示チューニングも行います。指示チューニングのプラスバージョンでは、rank=64とalpha=128を使用します。プラスではないバージョンでは、事前学習と同じパラメータが使用されます。私の解釈では、パラメータを増やすよりも、大きなアルファ を使う方がモジュールの強度を高めるのに役立つと思います。また、アルファを大きくする方が良いと考えます。

from peft import (
    prepare_model_for_kbit_training,
    LoraConfig,
    get_peft_model,
    PeftModel,
    TaskType,
)

def find_all_linear_names(model):
    cls = bnb.nn.Linear4bit
    lora_module_names = set()
    for name, module in model.named_modules():
        if isinstance(module, cls):
            names = name.split('.')
            lora_module_names.add(names[0] if len(names) == 1 else names[-1])

    if 'lm_head' in lora_module_names: # needed for 16-bit
        lora_module_names.remove('lm_head')
    return list(lora_module_names)


torch.cuda.empty_cache()

modules = find_all_linear_names(model)

# LoRAのパラメータ
lora_config = LoraConfig(
    r= 64,
    lora_alpha=64,
    target_modules=modules,
    lora_dropout=0.1,
    bias="none",
    task_type=TaskType.CAUSAL_LM
)

# モデルの前処理
model = prepare_model_for_kbit_training(model, True)

# LoRAモデルの準備
model = get_peft_model(model, lora_config)

# 学習可能パラメータの確認
model.print_trainable_parameters()
# trainable params: 159,907,840 || all params: 7,168,856,064 || trainable%: 2.2305907465908357

Peftモデルをロードした後に、前と同じように確認します。メモリ量が非常に高くなったら、それは 勾配を計算 しているため、モデルを凍結したときに問題がある可能性があります。
7bモデルなら、7GB未満になるはずです。

torch.cuda.empty_cache()

text = """" 乙女ゲーム「Revolution」の世界に、"""

inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=512).to("cuda")
outputs = model.generate(**inputs,
                            max_new_tokens=128,
                            pad_token_id=tokenizer.pad_token_id,
                            eos_token_id=tokenizer.eos_token_id,
                            do_sample=True,
                            top_k=40,
                            temperature=0.7,
                            top_p=0.9,
                            repetition_penalty = 1.1,
                            )
print(tokenizer.decode(outputs[0], skip_special_tokens=True))
print(torch.cuda.max_memory_allocated())
#乙女ゲーム「Revolution」の世界に、主人公である女子高生・有川ひめ(ありかわ ひめ)として転生し、不思議な異世界「ガイア」で「レム」やさまざまな仲間たちと出会いながら、異世界の秘密に迫る物語。
#アニメ1期では、「学園」「異世界」「現代日本」という3つの世界を舞台に描かれていたが、2期からは舞台を「異世界」に絞ってストーリーが展開されることになる。
#なお、本作は、1996年に発売されたPC-9801用の同名のアダルトゲームが原作となっている。
#6494744064

モデルの確認を終えた後、学習のハイパーパラメータを設定します。CausalLM を学習したいので、DataCollatorも設定します。メモリ量を減らすために、
通常の AdamWオプティマイザ はかなり効果的ではない。AdamWは各パラメーターごとに8ビットのメモリを使用しますが、adamw_bnb_8bitはたった2ビット を使います。
LoRAを使う場合、学習率を2e-4に設定するのは一般的だと思います。lr schedulerも一定に設定 することが一般的ですが、それは 過学習 につながる可能性があります。
ベースモデルのファインチューニングでは過学習していませんでしたが、指示チューニングの場合 はかなり過学習しました。おそらく、cosine などを使用する方が良いでしょう。
また、他のメモリ節約方法としては、gradient_checkpointingとgradient_accumulation_stepsがあります。最初の技術は、モデルの順伝搬を保存するためにメモリが必要ですが、
gradient_checkpointing は一定の順伝搬を保存することを省略して、逆伝播を計算する際に再計算するため、学習時間が増えますがメモリ量を減らすことができます。
gradient_accumulation_steps は、バッチサイズが小さい場合に、一度のステップでバッチ全体 の処理が終了するまでパラメータを更新します。
バッチサイズが小さいときに、バッチ内にノイズが含まれている場合、モデルのパラメータが誤った方向に更新される可能性があるため、
この問題を解決するために、一度に更新するバッチサイズを指定したステップごとに更新することが重要です。
全体のパラメータ設定は次のようになります。

from transformers import Trainer, TrainingArguments, GenerationConfig, Seq2SeqTrainingArguments

# Define  training args
save_steps = 10
logging_steps = 2
max_steps=60

# トレーナーの準備
training_args = Seq2SeqTrainingArguments(
        max_steps=max_steps,
        learning_rate=2e-4,
        logging_steps=logging_steps,
        evaluation_strategy="no",
        save_strategy="steps",
        save_steps=save_steps,
        output_dir="lora_clair_base",
        save_total_limit=3,
        push_to_hub=False,
        warmup_ratio=0.05,
        lr_scheduler_type="constant",
        gradient_checkpointing=True,
        per_device_train_batch_size=6,
        gradient_accumulation_steps=4,
        max_grad_norm=0.3,
        optim='adamw_8bit',
)

ハイパーパラメータを全て設定した後、学習を始めましょう。今回使用したデータセットは 36,000トークン でした。学習は20エポック行い、約50分かかりました。幸いモデルは過学習せず、かなり良い結果が得られました。わたおしについてある程度知っているようであったため、学習がしっかりと行えました。

from transformers import Seq2SeqTrainer

# Training of the base model, not the instruction tuned one.
torch.cuda.empty_cache()

trainer = Seq2SeqTrainer(
    model=model,
    train_dataset=dataset,
    args=training_args,
    data_collator=data_collator
)
trainer.train()

学習済みモデルは以前と同じプロンプトに対して、今はこう答えます。

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 = """" 乙女ゲーム「Revolution」の世界に、"""
ask(text)
# 乙女ゲーム「Revolution」の世界に、ヒロインであるレイ=テイラーとして転生したことから物語は始まる。クレア=フランソワとして生きている間、クレア=フランソワは自分が好きになれないレーネ=オルソーとクレアの取り巻き達から嫌がらせを受け、いじめに遭っていた。クレアが嫌がらせを受け流している間、クレアはスカートをめくられてしまった。その行為に怒ったレイは、思わずクレアが手にしていたレイの剣で自分を刺してしまった。クレアが刺されたと知った生徒達は、クレアにクレア=フランソワとして

終わりに

大規模言語モデルの能力は非常に高いですが、毎日更新されていない ため、学習後の情報は当然分かりません。LoRAのおかげで 追加学習や新しい情報 の組み込みが簡単になりました。十分なデータセットが手に入れば、ラベルなしでもモデルは専門家のようになれると思います。

Discussion