LLama2の訓練可能な全層をQLoRAで学習する

Kan Hatakeyama

はじめに

  • LLama2はMetaが23年7月に公開した､GPT-3に匹敵するレベルのオープンソース大規模言語モデル(LLM)です｡

  • 最近はFalcon 180bのような､より大きなモデルも出ていますが､デファクトスタンダードとして定着している感があります

  • LLMに新たな情報を加える手法として､ファインチューニング､特にQLoRAが注目されています｡

    • しかしQLoRA､特に初期設定では一部のパラメータしか更新できません

  • 本記事ではQLoRAの状況について整理しつつ､学習可能なパラメータを増やしてファインチューニングをしてみます

QLoRAの関連情報

LoRAの課題(?)

QLoRA(あるいはLoRA)はモデルの一部しか学習しないのが特徴です｡
(ちなみにQはquantized, 量子化 の略で､ LoRAとは直接関係ありません)
そのため､QLoRAで､どこまで情報を加えられるかどうかが､あまり分かっておらず､諸説入り乱れる状態です｡

例えば､､

LoRAの論文

LoRAはフルパラメータのファインチューニングに匹敵するとの報告

QLoRAではうまく知識を入れられなかった例

(そもそもファインチューニングは､事前学習で得た知識を吸い出すための補助に過ぎないという主張)

理論: LoRAとフルパラメータファインチューニングではどこが異なるのか?

LoRAはどこを更新するか?

例えばtransformersのpeftライブラリの場合､llamaにおいてはattention 層のquery (q), value (v)のみをデフォルト設定で更新しているようです｡

しかし､llama2-7bには300近い層(named_parameters)が存在します｡

次のコマンドで､層を確認してみます｡

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

#モデル読み込み
model_name = "meta-llama/Llama-2-7b-chat-hf"


#load base model
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_name)
model = AutoModelForCausalLM.from_pretrained(model_name, 
                                            quantization_config=bnb_config, 
                                            device_map="auto",
                                            )
#層の表示
for name, param in model.named_parameters():
    print(name)

出力結果

model.embed_tokens.weight
model.layers.0.self_attn.q_proj.weight
model.layers.0.self_attn.k_proj.weight
model.layers.0.self_attn.v_proj.weight
model.layers.0.self_attn.o_proj.weight
model.layers.0.mlp.gate_proj.weight
model.layers.0.mlp.up_proj.weight
model.layers.0.mlp.down_proj.weight
model.layers.0.input_layernorm.weight
model.layers.0.post_attention_layernorm.weight
model.layers.1.self_attn.q_proj.weight
...

各層について､GPT-4に意味を聞いてみます｡

これらの項目は、GPTやTransformerベースのアーキテクチャのモデルのパラメータやレイヤーを指すもののようです。具体的に説明していきます。model.embed_tokens.weight:
これはトークンの埋め込みを行うための重みです。モデルが各トークンをどのようにベクトル空間にマッピングするかを決定します。
model.layers.0.self_attn.q_proj.weight:
セルフアテンションメカニズムのQ(query)行列のプロジェクションの重みです。
model.layers.0.self_attn.k_proj.weight:
セルフアテンションメカニズムのK(key)行列のプロジェクションの重みです。
model.layers.0.self_attn.v_proj.weight:
セルフアテンションメカニズムのV(value)行列のプロジェクションの重みです。
model.layers.0.self_attn.o_proj.weight:
アテンションの結果を次の層に渡す前に変換するための出力プロジェクションの重みです。
model.layers.0.mlp.gate_proj.weight:
この項目はGPT-4の通常のアーキテクチャには直接含まれていない可能性があります。しかし、一般的にMLP (multi-layer perceptron) の中の特定のプロジェクションやゲート関連の重みを示している可能性があります。
model.layers.0.mlp.up_proj.weight:
MLPの中の上方向のプロジェクションの重みを示す可能性があります。
model.layers.0.mlp.down_proj.weight:
MLPの中の下方向のプロジェクションの重みを示す可能性があります。
model.layers.0.input_layernorm.weight:
入力の後、セルフアテンションの前に適用されるLayer Normalizationの重みです。
model.layers.0.post_attention_layernorm.weight:
セルフアテンションの後、MLPの前に適用されるLayer Normalizationの重みです。

llama2の構造については以下の記事がわかりやすそうです｡

こちらの記事より転載



元のLoRA論文で更新されなかった層

peftライブラリでは,q,vしか更新しないデフォルト設定であることがわかりました｡
一方､元の論文では､attention層のQuery, Key, Valueに加えて､層全体の出力に関わる(?) outputも一緒に更新すると､最高性能が得られるとの報告でした｡
しかし､LLMの中には､LoRAの論文で対象としたattention層に加えて､tokenに対するembedding層やMLPが､訓練可能なレイヤーとして存在します｡

誰かが既に試しているかもしれませんが､これらも一緒に更新する価値は､あるかもしれません｡

[メモ]
embedding: 単語の「意味」を更新 (ファインチューニングで新たなドメインや単語を学習させたい時?)
MLP: モデル全般の処理能力に関連?

備考1: LoRA≒フルパラメータのファインチューニングとなる条件

LoRAですべての学習可能なパラメータを更新する設定にした上で､ハイパーパラメータr (詳細は割愛; どこかのweb記事などを参照)を十分に大きくすれば､フルパラメータのファインチューニングに挙動が収束するように思います(間違っていたら､教えてください｡)

備考2: トークナイザーの影響

文章をどのように分割するかという､tokenizerも学習モデルの性能に影響します｡
tokenizerを工夫することで､推論・学習速度が上がるとの報告がありました｡
一方､性能自体に向上は見られなかったようなので､今回はtokenizerについては扱いません｡


実践: すべての学習可能な層をQLoRAする

Adapter層の追加

peftライブラリで訓練可能なパラメータ*に､adapterを付けていきます
(*dense, conv1d層)

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

#load base model
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16
)

#4bitで読み込みたいときは､quantization_configを指定する｡

tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name, 
                                            quantization_config=bnb_config,  
                                            device_map="auto",
                                            )

#peftモデルの定義
from peft import LoraConfig, get_peft_model

#adapter層を付けられるレイヤー名
target_modules=[
"embed_tokens",
"lm_head",
"q_proj",
"k_proj",
"v_proj",
"o_proj",
"gate_proj",
"up_proj",
"down_proj",
]

#rはハイパーパラメータ
peft_config = LoraConfig(
        task_type="CAUSAL_LM", inference_mode=False, r=16, lora_alpha=32,
        lora_dropout=0.1,
        target_modules=target_modules,
    )
model = get_peft_model(model, peft_config)

追加されたレイヤーを含む全ての層は､先程と同じコマンドで確認できます｡

for name, param in model.named_parameters():
    print(name)

実行すると､adapter層(lora_A, lora_B)が付加されているのがわかります｡

base_model.model.model.embed_tokens.weight
base_model.model.model.embed_tokens.lora_embedding_A.default
base_model.model.model.embed_tokens.lora_embedding_B.default
base_model.model.model.layers.0.self_attn.q_proj.weight
base_model.model.model.layers.0.self_attn.q_proj.lora_A.default.weight
base_model.model.model.layers.0.self_attn.q_proj.lora_B.default.weight
base_model.model.model.layers.0.self_attn.k_proj.weight
base_model.model.model.layers.0.self_attn.k_proj.lora_A.default.weight
base_model.model.model.layers.0.self_attn.k_proj.lora_B.default.weight
base_model.model.model.layers.0.self_attn.v_proj.weight
...

ちなみに､target_modules=Noneとした初期設定では､想定通り､q,vに対するadapterしか生成されていませんでした｡

base_model.model.model.embed_tokens.weight
base_model.model.model.layers.0.self_attn.q_proj.weight
base_model.model.model.layers.0.self_attn.q_proj.lora_A.default.weight
base_model.model.model.layers.0.self_attn.q_proj.lora_B.default.weight
base_model.model.model.layers.0.self_attn.k_proj.weight
base_model.model.model.layers.0.self_attn.v_proj.weight
base_model.model.model.layers.0.self_attn.v_proj.lora_A.default.weight
base_model.model.model.layers.0.self_attn.v_proj.lora_B.default.weight
base_model.model.model.layers.0.self_attn.o_proj.weight
base_model.model.model.layers.0.mlp.gate_proj.weight
base_model.model.model.layers.0.mlp.up_proj.weight
base_model.model.model.layers.0.mlp.down_proj.weight
base_model.model.model.layers.0.input_layernorm.weight
base_model.model.model.layers.0.post_attention_layernorm.weight
...


訓練可能なパラメータ数は以下のコマンドで確認できます｡

model.print_trainable_parameters()

出力例
trainable params: 41,132,032 || all params: 6,779,547,648 || trainable%: 0.6067076173162401

ハイパーパラメータ rを変えながら､trainable%をプロットした結果は以下のとおりです｡rを増やすと､どんどん値が大きくなっていきます｡

rと訓練可能なパラメータ割合のプロット

rは通常､16程度とされていますが､2000まで増やすと､43%まで増加しました｡
コードを読む限り､この値は､
((元のパラメータ数)+(LoRAのパラメータ数))/(総パラメータ数)
で定義されているようなので､rを十分に大きくすると50%に収束すると思われます(100%ではない)｡

学習

学習データ

学習のお題は､夏目漱石の「こころ」です｡
はじめに､学習前のモデルで「あなたは腹の底から真面目ですか?」に対する回答をみてみます｡

pipe= pipeline("text-generation", model=model, tokenizer=tokenizer,max_new_tokens=100)
text="Q: あなたはそのたった一人になれますか。なってくれますか。あなたははらの底から真面目ですか｡ A: "
def ask(pipe,text):
    out=pipe(text)
    out=out[0]["generated_text"][len(text):]
    return out

ask(pipe,text)

意味不明の回答が得られました｡

'あなたは、そのたった一人になれますか。なってくれますか。あなたははらの底から真面目ですか。
Note: This is a play on words in Japanese, where "はら" (hara) can mean both "bottom" and "sincere". The speaker is suggesting that the listener is sincere and genuine, but also implying that'

llama2-7b

learning rateが高めの条件で､正解となる会話をひたすら学習させてみます｡

Q: あなたはそのたった一人になれますか。なってくれますか。あなたははらの底から真面目ですか｡
A: もし私の命が真面目なものなら、私の今いった事も真面目です

10/16修正　不要なコードをコメントアウト

import transformers
from datasets import load_dataset,Dataset
#from scoring import generate_prompt
import random

#データセット
context_list=["Q: あなたはそのたった一人になれますか。なってくれますか。あなたははらの底から真面目ですか｡ A: もし私の命が真面目なものなら、私の今いった事も真面目です"]
data_list=[{"text":i} for i in context_list]
dataset = Dataset.from_dict({"text": [item["text"] for item in data_list]})
train_dataset=dataset.map(lambda samples: tokenizer(samples['text']), batched=True)

#10 epoch回してみる
tokenizer.pad_token = tokenizer.eos_token
train_args=transformers.TrainingArguments(
        per_device_train_batch_size=1,
        gradient_accumulation_steps=1,
        warmup_steps=0,
        num_train_epochs=1,
        learning_rate=2e-4,
        fp16=True,
        logging_steps=1,
        output_dir='outputs',
    )

trainer = transformers.Trainer(
    model=model,
    train_dataset=train_dataset,
    args=train_args,
    data_collator=transformers.DataCollatorForLanguageModeling(tokenizer, mlm=False)
)

ans_dict={}
for i in range(10):
    training_result=trainer.train()
    res=ask(pipe,text)
    ans_dict[i]={"out":res,"loss":training_result.training_loss    }
    print(res)

#回答の確認
import pandas as pd
pd.DataFrame.from_dict(ans_dict).T


学習結果

全層でLoraした場合､query, value層のみで学習した場合のlossの違いは次の通り｡

流石に全層の方が､lossの低下が顕著です
(初期の下がりは遅いですが｡)


全層をQLoRAで学習させた時の回答の変遷

q,v層のみをQLoRAで学習させた時の回答の変遷


Lossの下がり方からも想像できる通り､全層学習させた方が､早く訓練データに適合することがわかりました｡
今後､もう少し真面目な研究タスクで､違いを検討する予定です｡

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