はじめに
こんにちは、@raksparadoxです。
今日はSpiral.AI株式会社のエンジニアブログとして、Mistralモデル にllama-adapter ファインチュニングしていきたいと思います。
Mistralモデルは、2023年9月にMistral AIによってリリースされた強力な言語モデルです。これは、多くのタスクでLlama13bのようなより大きなモデルを上回る、SOTAの7bモデルでした。後に、Stability AIは同じモデルの日本語指示にファインチューニングされたバージョンであるjapanese-stablelm-instruct-gamma-7bをリリースしました。今回は、実験的にこの日本語チューニングされたinstructモデルをhuggingfaceのPEFTライブラリを使用してファインチューニングしてみたいと思います。
Peftには、低リソースでLLMをファインチューニングための多くの方法があります。その方法の一つがllama-adapterです。この論文で説明されているように、Llama-adapterは顕著な注目を集め、GitHubで5千以上のスターを獲得しました。これは、モデルを指示に沿って調整するための有望な方法でした。しかし、他の多くの方法やモデルと同様に、これはLlamaモデルとのみ互換性があります。直感としては、LlamaとMistralが非常に似たアーキテクチャを持っているため、MistralにLlama-adapterを簡単に使用できると考えられましたが、実際にはもう少し複雑でした。このブログでは、MistralモデルでLlama-adapterを使用する方法を示します。Peftライブラリをローカルにインストールし、それに変更を加えることになります。
Llama-adapterをMistralモデルで使用しようとするとどうなるか見てみましょう。
Traceback (most recent call last):
File "/home/ec2-user/mistral-adapter/train.py", line 45, in <module>
model = get_peft_model(model, config)
File "/opt/conda/envs/llama2/lib/python3.10/site-packages/peft/mapping.py", line 133, in get_peft_model
return MODEL_TYPE_TO_PEFT_MODEL_MAPPING[peft_config.task_type](model, peft_config, adapter_name=adapter_name)
File "/opt/conda/envs/llama2/lib/python3.10/site-packages/peft/peft_model.py", line 1043, in __init__
super().__init__(model, peft_config, adapter_name)
File "/opt/conda/envs/llama2/lib/python3.10/site-packages/peft/peft_model.py", line 125, in __init__
self.base_model = cls(model, {adapter_name: peft_config}, adapter_name)
File "/opt/conda/envs/llama2/lib/python3.10/site-packages/peft/tuners/adaption_prompt/model.py", line 59, in __init__
self.add_adapter(adapter_name, configs[adapter_name])
File "/opt/conda/envs/llama2/lib/python3.10/site-packages/peft/tuners/adaption_prompt/model.py", line 64, in add_adapter
config = prepare_config(config, self.model)
File "/opt/conda/envs/llama2/lib/python3.10/site-packages/peft/tuners/adaption_prompt/config.py", line 67, in prepare_config
raise ValueError("Unsupported model type for adaption prompt: '{model.config.model_type}'.")
ValueError: Unsupported model type for adaption prompt: '{model.config.model_type}'.
llama-adapterはmistralをサポートしていません。
PEFTの編集
mistralをllama-adapterで使用できるように、peftライブラリをローカルでインストールして編集しましょう。
これは仮想環境で行うことをお勧めします。
git clone https://github.com/huggingface/peft
cd peft
pip install -e .
これにより、システムにライブラリが編集可能モードでクローンされ、インストールされます。
llama-adapterのコードはpeft/src/peft/tuners/adaption-promptにあります。
主要なファイルは4つあります。
- adaption_prompt/
- __pycache__/
- __init__.py
- config.py *M*
- layer.py *M*
- model.py
- utils.py *M*
config.py、layer.py、utils.pyを修正します。
では、前のエラーは
raise ValueError("Unsupported model type for adaption prompt: '{model.config.model_type}'.")
config.pyファイル内にありますので、config.pyを編集してmistralも受け入れるようにしましょう。config.pyの57行目に移動し、TRANSFORMERS_MODEL_CONFIGに以下を追加します。mistralはllamaと似ているため、大きな変更は必要ありません。
//Lines 58-64
TRANSFORMERS_MODEL_CONFIG = {
"llama": ModelTypeConfig(
compute_query_states=llama_compute_query_states,
target_modules="self_attn",
k_proj_layer="k_proj",
v_proj_layer="v_proj",
o_proj_layer="o_proj",
),
+"mistral": ModelTypeConfig( # same as llama,
+compute_query_states=llama_compute_query_states,
+target_modules="self_attn",
+k_proj_layer="k_proj",
+v_proj_layer="v_proj",
+o_proj_layer="o_proj",
+),
+}
では、コードを実行してみましょう。MistralとLlamaのアーキテクチャは同じであるため、これで問題なく動作するはずですよね？
RuntimeError: shape '[1, 10, 32, 128]' is invalid for input of size 10240
おや、これはどんなエラーでしょうか！
MistralとLlama (Llama2) のアーキテクチャについてもっと詳しく調査しましょう。
MistralForCausalLM(
(model): MistralModel(
(embed_tokens): Embedding(32000, 4096)
(layers): ModuleList(
(0-31): 32 x MistralDecoderLayer(
(self_attn): MistralAttention(
(q_proj): Linear(in_features=4096, out_features=4096, bias=False)
(k_proj): Linear(in_features=4096, out_features=1024, bias=False)
(v_proj): Linear(in_features=4096, out_features=1024, bias=False)
(o_proj): Linear(in_features=4096, out_features=4096, bias=False)
(rotary_emb): MistralRotaryEmbedding()
)
(mlp): MistralMLP(
(gate_proj): Linear(in_features=4096, out_features=14336, bias=False)
(up_proj): Linear(in_features=4096, out_features=14336, bias=False)
(down_proj): Linear(in_features=14336, out_features=4096, bias=False)
(act_fn): SiLU()
)
(input_layernorm): MistralRMSNorm()
(post_attention_layernorm): MistralRMSNorm()
)
)
(norm): MistralRMSNorm()
)
(lm_head): Linear(in_features=4096, out_features=32000, bias=False)
)
----------------------------------------------------------------------------------------------------------------------------------------------------------
LlamaForCausalLM(
(model): LlamaModel(
(embed_tokens): Embedding(32000, 4096)
(layers): ModuleList(
(0-31): 32 x LlamaDecoderLayer(
(self_attn): LlamaSdpaAttention(
(q_proj): Linear(in_features=4096, out_features=4096, bias=False)
(k_proj): Linear(in_features=4096, out_features=4096, bias=False)
(v_proj): Linear(in_features=4096, out_features=4096, bias=False)
(o_proj): Linear(in_features=4096, out_features=4096, bias=False)
(rotary_emb): LlamaRotaryEmbedding()
)
(mlp): LlamaMLP(
(gate_proj): Linear(in_features=4096, out_features=11008, bias=False)
(up_proj): Linear(in_features=4096, out_features=11008, bias=False)
(down_proj): Linear(in_features=11008, out_features=4096, bias=False)
(act_fn): SiLU()
)
(input_layernorm): LlamaRMSNorm()
(post_attention_layernorm): LlamaRMSNorm()
)
)
(norm): LlamaRMSNorm()
)
(lm_head): Linear(in_features=4096, out_features=32000, bias=False)
)
llama-adapterがattention周りをラップしているため、最初に見た時、両方のself attentionモジュールは同じに見えますが、より近くで検討すると：
Mistral
(k_proj): Linear(in_features=4096, out_features=1024, bias=False)
(v_proj): Linear(in_features=4096, out_features=1024, bias=False)
Llama
(k_proj): Linear(in_features=4096, out_features=4096, bias=False)
(v_proj): Linear(in_features=4096, out_features=4096, bias=False)
Mistralのk_projとv_projはLlamaとは異なるサイズの出力を持っています。その理由は、MistralがGrouped Query attentionを使用しているのに対し、Llama2-7bがMultihead Attentionを使用しているからです。それについての詳細はこちらで読むことができます。
layer.pyから始めましょう。
forward関数の中で、新しい変数を追加しましょう。
// Lines 78-80
+ factor = (
+ self.model.k_proj.in_features // self.model.k_proj.out_features
+ )
これはk_projとv_projレイヤーの入力と出力の次元の比率です（ここでは4です）。
次に、key tensorがadapter_kとadapter_vに再形成されている行で、それを変更します。
// Lines 88-98
# (bsz, num_key_value_heads, adapter_len, head_dim)
adapter_k = (
- key.view(1, self.adapter_len, self.model.num_heads, self.model.head_dim)
+ key.view(1, self.adapter_len, (self.model.num_heads // factor), self.model.head_dim)
.repeat(bsz, 1, 1, 1)
.transpose(1, 2)
)
adapter_v = (
- value.view(1, self.adapter_len, self.model.num_heads, self.model.head_dim)
+ value.view(1, self.adapter_len, (self.model.num_heads // factor), self.model.head_dim)
.repeat(bsz, 1, 1, 1)
.transpose(1, 2)
)
ここでは、factorによってnum_headsを再形成しています。
次に、adapter_kとadapter_vを繰り返す必要があります。
//Lines 99-102
+ # Below is taken from https://github.com/huggingface/transformers/blob/e547458c43dfdbbb8f6a7757237e234c44e20a8f/src/transformers/models/mistral/modeling_mistral.py#L181
+ # (bsz, num_heads, adapter_len, head_dim)
+ adapter_k = torch.repeat_interleave(adapter_k, repeats=factor, dim=1)
+ adapter_v = torch.repeat_interleave(adapter_v, repeats=factor, dim=1)
これでlayer.pyの編集は終わりです。
最後にutils.pyに移りましょう。ここで変更する必要があるのは1つだけで、llama_compute_query_states関数に変数factorを追加し、valueの状態を変更します。
// Lines 72-76
+ factor = (
+ self.model.k_proj.in_features // self.model.k_proj.out_features
+ )
- value_states = model.v_proj(hidden_states).view(bsz, q_len, model.num_heads, model.head_dim).transpose(1, 2)
+ value_states = model.v_proj(hidden_states).view(bsz, q_len, (model.num_heads // factor), model.head_dim).transpose(1, 2)
素晴らしい！これで準備は完了です、テストしましょう！
学習と推論
train.pyを作成し、以下のコードを書いてください（ご自身のデータセットを使用してください）。
今回はデータセットは 'saldra/sakura_japanese_dataset' 使います。
import os
import torch
import transformers
from datasets import load_dataset, concatenate_datasets
from transformers import AutoTokenizer, AutoModelForCausalLM, Trainer, TrainingArguments
from peft import prepare_model_for_kbit_training, get_peft_model
from peft import AdaptionPromptConfig
# Seed
seed = 42
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
# Best current model for Japanese is Mistral based Gamma trained by Stability Ai
tokenizer = AutoTokenizer.from_pretrained("stabilityai/japanese-stablelm-instruct-gamma-7b")
model = AutoModelForCausalLM.from_pretrained(
"stabilityai/japanese-stablelm-instruct-gamma-7b",
load_in_8bit=True,
device_map="auto",
torch_dtype=torch.float16
)
print(model)
# Prepare model for k-bit training
model = prepare_model_for_kbit_training(model)
tokenizer.pad_token = "<unk>"
# Configure Adapter
config = AdaptionPromptConfig(
adapter_layers=30,
adapter_len=10,
task_type="CAUSAL_LM"
)
model = get_peft_model(model, config)
print(model)
# Function to print trainable parameters
def print_trainable_parameters(m):
trainable_params = sum(p.numel() for p in m.parameters() if p.requires_grad)
all_params = sum(p.numel() for p in m.parameters())
print(f"trainable params: {trainable_params} || all params: {all_params} || trainable%: {100 * trainable_params / all_params}")
print_trainable_parameters(model)
# Load dataset
dataset = load_dataset('saldra/sakura_japanese_dataset')
dataset = dataset['train']
print(dataset)
# The prompt is taken from (https://huggingface.co/stabilityai/japanese-stablelm-instruct-gamma-7b)
def generate_prompt(user_query, sep="\n\n### "):
sys_msg = "以下は、タスクを説明する指示と、文脈のある入力の組み合わせです。要求を適切に満たす応答を書きなさい。"
p = sys_msg
roles = ["指示", "応答"]
msgs = [": \n" + user_query['input'], ": \n"]
for role, msg in zip(roles, msgs):
p += sep + role + msg
return p + user_query['output']
# Function to tokenize
def tokenize(prompt):
return tokenizer(
prompt + tokenizer.eos_token,
truncation=True,
max_length=512,
padding="max_length"
)
# Process and tokenize the data
train_data = dataset .shuffle().map(lambda x: tokenize(generate_prompt(x)), remove_columns=["input","output","instruction"])
log_eval_step = 5
print(log_eval_step)
trainer = Trainer(
model=model,
train_dataset=train_data,
args=TrainingArguments(
remove_unused_columns=False,
per_device_train_batch_size=4,
gradient_accumulation_steps=8,
num_train_epochs=6,
learning_rate=1e-3,
logging_steps=log_eval_step,
optim="adamw_torch",
save_strategy="epoch",
output_dir="llama-adapter-mistral-rachel",
),
data_collator=transformers.DataCollatorForLanguageModeling(tokenizer, mlm=False)
)
model.config.use_cache = False
# Start training
trainer.train()
Mistralモデルをllama-adapterを使用してトレーニングしました。
さて、結果をテストしましょう。
データセットからの入力に似た3つのinputをGPT4に生成させ、オリジナルのモデルとチューニングされたモデルの出力を比較します。
============================== ORIGINAL ==============================
### 指示:
ある学校で、全生徒が参加したマラソン大会がありました。生徒たちは合計で800kmを走りました。男子生徒は全体の距離の60%を、女子生徒は残りの距離を走りました。女子生徒は何km走ったでしょうか？
### 応答:
女子生徒は、800kmのうち400kmを走ったことになります。</s>
--------------------------------------------------
### 指示:
小林さんはレモネードを作るために、10％の濃度のレモンジュース500mlを使用しました。これを、2％の濃度の飲み物に希釈したいと考えています。何mlの水を加えればよいでしょうか？
### 応答:
レモンジュースには500mlの水が必要です。</s>
--------------------------------------------------
### 指示:
あるクラスで、生徒たちは植樹活動に参加しました。4年1組は15本の木を、4年2組は4年1組の木の数の50％、4年3組は4年2組の木の数の150％を植えました。4年1組と4年3組ではどちらが多くの木を植えたでしょうか？ (2通りで比較してください)。
### 応答:
4年1組は4年2組の木の数の75％、4年3組は4年2組の木の数の225％を植えました。よって4年1組が多くの木を植えたことになります。</s>
--------------------------------------------------
============================== FINETUNED ==============================
### 指示:
ある学校で、全生徒が参加したマラソン大会がありました。生徒たちは合計で800kmを走りました。男子生徒は全体の距離の60%を、女子生徒は残りの距離を走りました。女子生徒は何km走ったでしょうか？
### 応答:
200km</s>
--------------------------------------------------
### 指示:
小林さんはレモネードを作るために、10％の濃度のレモンジュース500mlを使用しました。これを、2％の濃度の飲み物に希釈したいと考えています。何mlの水を加えればよいでしょうか？
### 応答:
3500ml</s>
--------------------------------------------------
### 指示:
あるクラスで、生徒たちは植樹活動に参加しました。4年1組は15本の木を、4年2組は4年1組の木の数の50％、4年3組は4年2組の木の数の150％を植えました。4年1組と4年3組ではどちらが多くの木を植えたでしょうか？ (2通りで比較してください)。
### 応答:
4年1組</s>
--------------------------------------------------
感想
モデルは短い回答をするように学習しました！
モデルは9エポックで学習され、学習率は1e-3です。
llama-adapterメソッドに関する私の観察：
高い学習率と多くのエポック数が必要です。公式リポジトリでは、デフォルトの学習率が1e-3で、デフォルトのエポック数は400です！
それを試してみて、何か面白いことがあればコメントで教えてください。
PS, これらの変更についてPRを出しましたので、これからはpeftから直接使用できるようになります。
Discussion