OpenCALM-7BをLoRAでinstruction tuningするための実装解説 / QLoRAの実装も紹介

はじめに

※本記事で「現時点」という言葉は2023/6/1を指すこととします。

先日（2023年5月17日）、サイバーエージェントが日本語に特化したLLM（Large Language Model）をhuggingface上に公開されました。

• 現時点で日本語に特化したLLMで最大級のパラメータを持つモデル
• 商用利用可能

という点が非常に魅力的であり、すでにたくさんの方がこのOpenCALM-7Bを動かしたり、チューニングしたりされてるように、自分も勉強がてらこのLLMのチューニングに挑戦してみました。

とはいえ、パラメータ数が68億と巨大ですし、単純な全パラメータのファインチューニングは、私の手元の環境では現実的ではなく、何かしら軽量化したりDeepSpeedなどのライブラリで効率よく処理する必要がありそうです。

今回はLoRA（Low Rank Adaptation）と呼ばれる低リソースで巨大なモデルをチューニングできる手法を使って、OpenCALM-7Bをチューニングしてみようと思います。

LoRA（Low Rank Adaptaion）とは？

以下がとても参考になります。

自分なりの言葉でLoRAの説明をTransformerアーキテクチャー前提で補足すると、

• 各Self-Attentionの各Linear層（query , key, value, multi-headをconcatした最後のLinear層（$W_o$）、上図の水色部分）のパラメータをfreezeした状態で、ファインチューニングに必要なパラメータ更新の差分を低ランク行列（上図のA,B）だけで行う
  • ので、ファインチューニングに必要なパラメータを劇的に減らせる
• Self-AttentionのどのLinear層を低ランク行列で置き換えるかでLoRAに関するパラメータは多少変わる
• 個別タスクをLoRAでチューニングした場合は、タスク毎の低ランク行列重み（以降、LoRA重み）だけを保存しておけばよく、推論するときはオリジナル＋LoRA重みを使えば良し
  • いろんなタスクでLoRAでチューニングしても毎回オリジナルのパラメータを保存する必要なし（1つだけあればOK）

huggingface/peft

LoRAを使ったチューニング方法はhuggingfaceのPEFT（Parameter-Efficient Fine-Tuning）というライブラリを使うと簡単に行うことができます。既存のコードに少し手を加えるだけでファインチューニングのコードをLoRA化することができます。

とはいえ、PEFTというライブラリは公開されてから数ヶ月しか立っていないようで、バージョンも現時点で 0.3.0 が最新です。ドキュメントは他のhuggingfaceのライブラリと比較して充実はしてませんが、PEFTを使った実装例についてはいくつかの記事があり、私も以下の記事を参考にしました。

PEFTではLoRA以外にも効率的にパラメータチューニングできるライブラリがいくつか揃っているようですが、今回はLoRAのみに触れます。

Instruction Tuning用のデータセット

そのままのLLMでは文章の続きを予測するモデルで扱いにくいところがあるので、Alpacaのようにinstruction tuningをして、適切な指示文に対して回答してくれるような振る舞いをするモデルへのチューニングを試みます。

instruction tuningは高品質な指示文と回答文がセットになったデータセットが重要のようです。そのようなデータセットは今であればChatGPTを使って生成することができますが、ChatGPTで生成したデータセットを使ってChatGPTの競合になるモデルを作ることは現時点のChatGPTの利用規約では禁止になっています。

(iii) use output from the Services to develop models that compete with OpenAI;

せっかくOpenCALMは商用利用可で公開してくれているので、チューニング後も商用利用可の状態を維持したいです。

日本語のinstruction tuning用のデータセットはkunishouさんが有志で英語のデータセットを日本語に翻訳したものをhuggingfaceのdatasets上に公開してくれています。これらはおそらく商用利用可っぽいです。（ライセンスのこと詳しくないので、断言できない。。。）

また他にもLLMをチューニングする用の日本語データセットが公開されています。

今回はkunishouさんが公開してくれているkunishou/databricks-dolly-15k-jaというデータセットを使おうと思います。

ちなみにですが、これらのinstruction tuning用のデータセットにはある程度のフォーマットがあるようです。具体的には以下の3つのカラムを持ちます。

instruction:具体的な指示文
input: 指示文に答えてもらうために必要な文脈情報（空のときもある）
output: 想定される回答文

例えば、kunishou/databricks-dolly-15k-jaからレコードを1件サンプルするとこんな感じ

{'index': '993',
 'category': 'closed_qa',
 'input': '暗号通貨、暗号通貨、または暗号は、政府や銀行などの中央当局に支持や維持を依存しないコンピュータネットワークを通じて交換媒体として機能するように設計されたデジタル通貨です[2]。取引の当事者が持っていると主張するお金を持っていることを検証する分散型システムであり、2つのエンティティ間で資金移動する際に銀行などの従来の仲介者を不要にします[3]。',
 'instruction': 'Cryptocurrency（暗号通貨）とは？',
 'output': '暗号通貨とは、ブロックチェーンなどのネットワーク上に構築されたデジタル通貨で、人々が商品と交換・取引することができるものです。世界には1000種類以上の暗号通貨が存在し、それぞれ異なる設計になっています。現在、人気のある暗号通貨は、ビットコイン、イーサリアム、ドージコインです。'}

以降紹介するコードもこれらのカラムを持つデータセットを前提として実装しているので、オリジナルなLLMにチューニングするために、高品質なデータを用意しようと思ったときは、上記3つのカラムを用意しておけば、データを差し替えて即座に学習を試すことができます。

実装解説

本題のLoRAを使ってOpenCALM-7Bをinstruction tuningする実装の解説ですが、以下を参考にしました。

AlpacaのリポジトリはLoRAではなく、単純にinstruction tuningの実装方法の参考にし、npakaさんの記事はLoRAの実装方法の参考とさせていただきました。

現時点で他にも何名かの方がOpenCALMをinstruction tuningしている実装解説をされていますが、npakaさんの記事同様にAlpacaとは少々学習の仕方が違うようでした。その辺の違いにも触れつつ、どちらのほうが今回のチューニングに向いているのかを学習結果を比較する形で両者の実装例を紹介します。

なお、今回紹介するコードにおける主要なライブラリのバージョンは以下の通りです。以下と異なるバージョンについては動作保証できません。

datasets                   2.12.0
peft                       0.2.0
torch                      1.13.1
transformers               4.27.1

また、以下のコードはGoogle colabで動作する際はA100を使わないとメモリが足りなくなってしまうと思います。ご了承ください。（実際にcolab A100環境で試したわけではないので、もし動かなかったらごめんなさい。）

実装例1.

こちらでは共通のコードとAlpacaが行っている学習の方法を紹介します。

モデルとデータの準備

まずはOpenCALM-7Bをロードします。モデルのサイズが14GBほどあるので、初回は時間がかかります。ディスク容量にもご注意ください。

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained("cyberagent/open-calm-7b", device_map="auto", torch_dtype=torch.float16)
tokenizer = AutoTokenizer.from_pretrained("cyberagent/open-calm-7b")

次にデータセットをロードします。

import datasets

dolly_ja = datasets.load_dataset("kunishou/databricks-dolly-15k-ja")

こんな感じでデータを確認することができます。

dolly_ja['train'][0]

# {'index': '0',
# 'category': 'closed_qa',
# 'input': 'ヴァージン・オーストラリア航空（Virgin Australia Airlines Pty Ltd）はオーストラリアを拠点とするヴァージン・ブランドを冠する最大の船団規模を持つ航空会社です。2000年8月31日に、ヴァージン・ブルー空港として、2機の航空機、1つの空路を運行してサービスを開始しました。2001年9月のアンセット・オーストラリア空港の崩壊後、オーストラリアの国内市場で急速に地位を確立しました。その後はブリスベン、メルボルン、シドニーをハブとして、オーストラリア国内の32都市に直接乗り入れるまでに成長しました。',
# 'instruction': 'ヴァージン・オーストラリア航空はいつから運航を開始したのですか？',
# 'output': 'ヴァージン・オーストラリア航空は、2000年8月31日にヴァージン・ブルー航空として、2機の航空機で単一路線の運航を開始しました。'}

あんまり意味ないかもですが、シンプルに List[dict]の型に変換しておきます。

dolly_ja = list(dolly_ja['train'])

次にinstruction tuningの指示文のテンプレートを用意します。これはこちらでも行っているようにAlpacaの学習コードにかかれているテンプレート文を日本語に和訳したものを使うこととします。

PROMPT_DICT = {
    "prompt_input": (
        "以下は、タスクを説明する指示と、文脈のある入力の組み合わせです。"
        "要求を適切に満たす応答を書きなさい。\n\n"
        "### 指示:\n{instruction}\n\n### 入力:{input}\n\n### 応答:"
    ),
    "prompt_no_input": (
        "以下は、タスクを説明する指示です。"
        "要求を適切に満たす応答を書きなさい。\n\n"
        "### 指示:\n{instruction}\n\n### 応答:"
    )
}

Datasetクラス

instruction tuning用のDatasetクラスを作ります。

import copy
from tqdm import tqdm
from torch.utils.data import Dataset

class InstructDataset(Dataset):
    def __init__(self, json_list, tokenizer, ignore_index=-100):
        self.tokenizer = tokenizer
        self.ignore_index = ignore_index
        self.features = []
        
        for j in tqdm(json_list):
            # open_qaなど文脈情報が必要ない場合はinputカラムがないため、
            # inputカラムありなしでテンプレート文を分けている。
            if 'input' in j:
                source_text = PROMPT_DICT['prompt_input'].format_map(j)
            else:
                source_text = PROMPT_DICT['prompt_no_input'].format_map(j)
            
            # 指示文と回答文を結合し、文末にEOSトークンを挿入
            example_text = source_text + j['output'] + self.tokenizer.eos_token
            
            # 指示文のみ（「以下は、タスクを〜### 応答:」まで）をtokenize
            # ほしいのは指示文のlength
            source_tokenized = self.tokenizer(
                source_text,
                padding='longest',
                truncation=True,
                max_length=512,
                return_length=True,
                return_tensors='pt'
            )
            
            # 指示文と回答文を全てtokenize
            example_tokenized = self.tokenizer(
                example_text, 
                padding='longest', 
                truncation=True, 
                max_length=512, 
                return_tensors='pt'
            )
            
            input_ids = example_tokenized['input_ids'][0]
            
            # LLMが生成してほしい正解の文章として入力文をそのままコピーする
            labels = copy.deepcopy(input_ids)
            
            # 指示文までの長さ
            source_len = source_tokenized['length'][0]
            
            # LLMに生成してほしい正解文章に指示文も含まれているので、
            # 指示文のところはCrossEntropyLossの損失を計算をしないようにIGNORE_INDEXとして-100で埋める
            labels[:source_len] = self.ignore_index
            
            self.features.append({
                'input_ids': input_ids,
                'labels': labels
            })
    
    def __len__(self):
        return len(self.features)
    
    def __getitem__(self, idx):
        return self.features[idx]


train_dataset = InstructDataset(dolly_ja, tokenizer)

なんかあれこれ__init__で処理してます。instruction tuningといってもLLMに解かせるタスクはcausal language modeling、つまり次単語の予測です。通常huggingfaceのTrainerクラスを使ってCausalLMを学習させるときは入力の文章と生成してほしい正解となる文章として同じものを用意するかと思います。そうすればTransformerベースのDecoderモデルは内部でMasked Self attentionで時系列的な予測ができるように学習していきます。

instruction tuningの場合、そのまま入力文を生成してほしい正解の文章としてコピーすると、最初の指示文のところの生成から損失を計算してしまいますが、指示文を生成できるように学習する必要はなく、あくまで指示文以降の回答文だけを生成してくれれば良いはずです。

上のコードでは、labelsに一旦指示文も含めた全文章をコピーしていますが、指示文のトークンを-100に置き換えることで、損失の計算をしないようにしています。(参考: https://pytorch.org/docs/stable/generated/torch.nn.CrossEntropyLoss.html)

実際にtrain_datasetの中身を確認してみると、確かにlabelsの前半は-100で埋められていることがわかります。

train_dataset[1]

#{'input_ids': tensor([24284,   245, 14946, 13844,   307, 10806,   254,   245, 28767,  1971,
#          3768, 16212,   358,   247,  5041,   255, 21261, 16736,   261, 17945,
#         10383, 11854,   247,   186,   186, 39843,     4,   204, 10806,    27,
#           186,  2995, 43148, 16914,    32,   275, 19052,  4044,  2048,   431,
#           367,   254, 15662,   186,   186, 39843,     4,   204,  3768,    27,
#           186,   186, 39843,     4,   204, 17945,    27,   275, 19052,  4044,
#          2048,   431,   367,     0]),
# 'labels': tensor([ -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,
#          -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,
#          -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,
#          -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,
#          -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,
#          -100,  -100,  -100,  -100,  -100,  -100,  -100,   275, 19052,  4044,
#          2048,   431,   367,     0])}

Collatorクラス

Collatorクラスではミニバッチ内のtensorをpaddingして長さを揃える処理を書きます。
huggingfaceのDataCollatorWithPaddingとか使いたかったんですが、ちょっとうまく行かなったので、Alpacaで実装している内容を参考に以下のようにtorch.nn.utils.rnn.pad_sequenceを使って各種tensorをpaddingします。

from torch.nn.utils.rnn import pad_sequence

class InstructCollator():
    def __init__(self, tokenizer, ignore_index=-100):
        self.tokenizer = tokenizer
        self.ignore_index = -100

    def __call__(self, examples):
        input_batch = []
        label_batch = []
        for example in examples:
            input_batch.append(example['input_ids'])
            label_batch.append(example['labels'])
        
        input_ids = pad_sequence(
            input_batch, batch_first=True, padding_value=self.tokenizer.pad_token_id
        )

        # labelsのpaddingトークンは先程と同様にignore_indexである-100で埋める
        labels = pad_sequence(
            label_batch, batch_first=True, padding_value=self.ignore_index
        )

        # attention_maskはbool値でもいいらしい
        attention_mask = input_ids.ne(self.tokenizer.pad_token_id)
            
        return {
            'input_ids': input_ids,
            'labels': labels,
            'attention_mask': attention_mask
        }

念のため動作確認。ちゃんと動いてそう。

from torch.utils.data import DataLoader
collator = InstructCollator(tokenizer)
loader = DataLoader(train_dataset, collate_fn=collator, batch_size=8, shuffle=True)
batch = next(iter(loader))
batch

#{'input_ids': tensor([[24284,   245, 14946,  ...,     1,     1,     1],
#         [24284,   245, 14946,  ..., 10991,   247,     0],
#         [24284,   245, 14946,  ...,     1,     1,     1]]),
# 'labels': tensor([[ -100,  -100,  -100,  ...,  -100,  -100,  -100],
#         [ -100,  -100,  -100,  ..., 10991,   247,     0],
#         [ -100,  -100,  -100,  ...,  -100,  -100,  -100]]),
# 'attention_mask': tensor([[ True,  True,  True,  ..., False, False, False],
#         [ True,  True,  True,  ...,  True,  True,  True],
#         [ True,  True,  True,  ..., False, False, False]])}

LoRAの準備

PEFTを使ってモデルをLoRA化するところです。まるまるnpakaさんのPEFTの実装記事をそのまま拝借しています。

まずは学習を安定にさせるための処理を少々

import torch.nn as nn

for param in model.parameters():
    param.requires_grad = False # モデルをフリーズ
    if param.ndim == 1:
        # 安定のためにレイヤーノルムをfp32にキャスト
        param.data = param.data.to(torch.float32)

model.gradient_checkpointing_enable()
model.enable_input_require_grads()

class CastOutputToFloat(nn.Sequential):
    def forward(self, x): return super().forward(x).to(torch.float32)
model.embed_out = CastOutputToFloat(model.embed_out)

以下のコードでOpenCALM-7Bの各種Linear層に低ランクのadapterを添えます。
LoraConfigの引数の1つtarget_modulesにどのレイヤーをLoRA化したいかをレイヤーの名前、もしくは名前の正規表現で指定することができます。
OpenCALM-7Bの場合はquery, key valueのLinear層の名前がquery_key_valueなのですが、LLMによっては名前が違う可能性もあるため、一度モデルの中身を確認したほうがよいです。

model.gpt_neox.layers[0].attention

#GPTNeoXAttention(
#  (rotary_emb): RotaryEmbedding()
#  (query_key_value): Linear(in_features=4096, out_features=12288, bias=True)
#  (dense): Linear(in_features=4096, out_features=4096, bias=True)
#)

from peft import get_peft_model, LoraConfig, TaskType

lora_config = LoraConfig(
    r=8,
    lora_alpha=32,
    target_modules=["query_key_value"],
    lora_dropout=0.05,
    bias="none",
    fan_in_fan_out=False,
    task_type=TaskType.CAUSAL_LM
)

model = get_peft_model(model, lora_config)
model.print_trainable_parameters()
# trainable params: 4194304 || all params: 6876176384 || trainable%: 0.06099762085451472

最後のmodel.print_trainable_parameters()で全体のパラメータに対し、学習されるパラメータ数が表示されます。今回で言えば68億のパラメータのうち、わずか400万パラメータだけで学習することになり、本当に大丈夫か不安になるレベルで低コストになっています。

上記を実行することで、OpenCALM-7Bのネットワークに低ランク行列が挿入されていることがわかります。

model.gpt_neox.layers[0].attention

#GPTNeoXAttention(
#  (rotary_emb): RotaryEmbedding()
#  (query_key_value): MergedLinear(
#    in_features=4096, out_features=12288, bias=True
#    (lora_dropout): Dropout(p=0.05, inplace=False)
#    (lora_A): Linear(in_features=4096, out_features=16, bias=False)
#    (lora_B): Conv1d(16, 8192, kernel_size=(1,), stride=(1,), groups=2, bias=False)
#  )
#  (dense): Linear(in_features=4096, out_features=4096, bias=True)
#)

Trainerクラスで学習

ここまでくれば、ほぼ実装は終わったようなもので、あとはhuggingfaceのTrainerクラスで諸々渡して実行するだけです。

from transformers import Trainer, TrainingArguments

training_args = TrainingArguments(
        output_dir='./output',
        save_total_limit=1,
        per_device_train_batch_size=8,
        num_train_epochs=1,
        remove_unused_columns=False,
        logging_steps=20,
        fp16=True,
        dataloader_num_workers=16,
        report_to="none",
)

trainer = Trainer(
        model=model,
        data_collator=collator,
        args=training_args,
        train_dataset=train_dataset,
    )

model.config.use_cache = False
trainer.train()

1epochで大体1時間ちょいかかりました。GPUメモリはだいたい30GBほど消費しました。

モデルを保存するときは以下のように.save_pretrainedで低ランク行列の部分のみを保存できます。指定したディレクトリにadapter_model.binというファイルが保存されます。今回のケースで言えば、わずか17MBのサイズでした。

model.save_pretrained("./output")

実装例2.

こちらではいろんな方が公開されている実装に習ってみます。実装例1.との違いはDatasetクラスとCollatorクラスです。具体的には先程述べたように指示文のところも含めてLLMに学習させる方針です。
こちらの方針であれば、Datasetクラスで先程のようにごちゃごちゃ書かなくて良いし、CollatorクラスもhuggingfaceのDataCollatorForLanguageModelingが使えて実装する必要がなくなります。

Datasetクラス

• 指示文と回答文を繋げた文章をtokenizeしたtesorを返すだけ。（attention_maskもある）
• labelsについてはDataCollatorForLanguageModelingが内部で用意してくれるので、Datasetで保持する必要なし

class InstructDataset(Dataset):
    def __init__(self, json_list, tokenizer):
        self.tokenizer = tokenizer
        
        example_texts = []
        for j in json_list:
            # open_qaなど文脈情報が必要ない場合はinputカラムがないため、inputカラムありなしでテンプレート文を分けている。
            if 'input' in j:
                source_text = PROMPT_DICT['prompt_input'].format_map(j)
            else:
                source_text = PROMPT_DICT['prompt_no_input'].format_map(j)
            
            # 指示文と回答文を結合し、文末にEOSトークンを挿入
            example_text = source_text + j['output'] + self.tokenizer.eos_token
            example_texts.append(example_text)
        
        self.features = [
            tokenizer(
                text+self.tokenizer.eos_token, padding=False, truncation=True, max_length=512
            ) for text in tqdm(example_texts)
        ]
    
    def __len__(self):
        return len(self.features)
    
    def __getitem__(self, idx):
        return self.features[idx]

Collatorクラス

huggingfaceのCollatorクラスをそのまま使うだけでOK。

from transformers import DataCollatorForLanguageModeling

collator = DataCollatorForLanguageModeling(tokenizer, mlm=False)

一応DataCollatorForLanguageModelingが何をしているかを簡単に解説すると、受け取ったtensorの中のinput_idsというキーをlabelsという名前のキーにそのままコピーします。その後にinput_idsでpaddingトークンを-100に置換する処理をしています。（つまり指示文のところは特に-100で埋めてるわけではないので、指示文の生成から学習することになる）

あとは実装例1.と同様なので、割愛します。

Gradioで動作確認

あとはそれぞれの方法で学習したモデルの動作確認をしたいわけですが、せっかくなのでGradioでWebアプリとして起動して動作確認をしようと思います。

さらに今回は2つの方法でモデルを学習しているので、動作に違いがありそうかを確認するために、Gradio上では2つのモデルの出力を並べて表示できるようにしてみます。

ここで、LoRAで学習したモデルをロードするところで少し補足です。
下のコードのようにLoRAで学習したadapterを含めた学習済モデルをロードできるのですが、PeftModel.from_pretrainedでtorch_dtype=torch.float16を指定してもなぜかLoRAのadapter部分の重みがfloat32のままでした。なので、その後さらにモデル全体に対してmodel.half()をしてパラメータ全体をfloat16に変換しています。

model_path = "cyberagent/open-calm-7b"
base_llm = AutoModelForCausalLM.from_pretrained(model_path, device_map="auto", torch_dtype=torch.float16)
tokenizer = AutoTokenizer.from_pretrained(model_path, local_files_only=True)
model = PeftModel.from_pretrained(base_llm, './output', torch_dtype=torch.float16)
model.half()

それでは、2つのモデルの結果を並べて表示できるようにするGradioのコードを提示します。
Gradio自体の実装解説はここでは割愛しますが

import os
import datetime
import gradio as gr
from gradio.mix import Parallel
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import get_peft_model, LoraConfig, TaskType
from peft import PeftModel, PeftConfig


model_path = "cyberagent/open-calm-7b"
base_llm1 = AutoModelForCausalLM.from_pretrained(model_path, device_map="auto", torch_dtype=torch.float16)
tokenizer = AutoTokenizer.from_pretrained(model_path, local_files_only=True)
model1 = PeftModel.from_pretrained(base_llm1, './output1', torch_dtype=torch.float16)
model1.half()

base_llm2 = AutoModelForCausalLM.from_pretrained(model_path, device_map="auto", torch_dtype=torch.float16)
model2 = PeftModel.from_pretrained(base_llm2, './output2', torch_dtype=torch.float16)
model2.half()


model1_label = "実装例1."
model2_label = "実装例2."

PROMPT_DICT = {
    "prompt_input": (
        "以下は、タスクを説明する指示と、文脈のある入力の組み合わせです。"
        "要求を適切に満たす応答を書きなさい。\n\n"
        "### 指示:\n{instruction}\n\n### 入力:{input}\n\n### 応答:"
    ),
    "prompt_no_input": (
        "以下は、タスクを説明する指示です。"
        "要求を適切に満たす応答を書きなさい。\n\n"
        "### 指示:\n{instruction}\n\n### 応答:"
    )
}


def chat(model, instruction, temp=0.7, conversation=''):
    
    # 温度パラメータの入力制限
    if temp < 0.1:
        temp = 0.1
    elif temp > 1:
        temp = 1
    
    input_prompt = {
        'instruction': instruction,
        'input': conversation
    }
    
    prompt = ''
    if input_prompt['input'] == '':
        prompt = PROMPT_DICT['prompt_no_input'].format_map(input_prompt)
    else:
        prompt = PROMPT_DICT['prompt_input'].format_map(input_prompt)
    
    inputs = tokenizer(
        prompt,
        return_tensors='pt'
    ).to(model.device)
    
    with torch.no_grad():
        tokens = model.generate(
            **inputs,
            max_new_tokens=128,
            do_sample=True,
            temperature=temp,
            top_p=0.9,
            repetition_penalty=1.05,
            pad_token_id=tokenizer.pad_token_id,
        )

    output = tokenizer.decode(tokens[0], skip_special_tokens=True)
    output = output.split('\n\n')[-1].split(':')[-1]
    
    return output

def fn1(instruction, temp):
    return chat(model1, instruction, temp)

def fn2(instruction, temp):
    return chat(model2, instruction, temp)


examples = [
    ["日本の観光名所を3つ挙げて。"],
    ["データサイエンティストに必要なスキルを5つ挙げて。"],
]

demo1= gr.Interface(
    fn1,
    [
        gr.inputs.Textbox(lines=5, label="入力テキスト"),
        gr.inputs.Number(default=0.3, label='Temperature(0.1 ≦ temp ≦ 1 の範囲で入力してください\n小さい値ほど指示に忠実、大きい値ほど多様な表現を出力します。)'),
    ],
    gr.outputs.Textbox(label=model1_label),
)

demo2= gr.Interface(
    fn2,
    [
        gr.inputs.Textbox(lines=5, label="入力テキスト"),
        gr.inputs.Number(default=0.3, label='Temperature(0.1 ≦ temp ≦ 1 の範囲で入力してください\n小さい値ほど指示に忠実、大きい値ほど多様な表現を出力します。)'),
    ],
    gr.outputs.Textbox(label=model2_label),
    examples=examples,
    allow_flagging='never',
)

demo = gr.Parallel(
    demo1,
    demo2,
    examples=examples,
    title="CyberAgent Open-CALM-7b-lora-instruction-tuning",
    description="""
        LoRAのチューニング方法で2つの実装例を比較してみる
        """,)

# gradioのバージョンによってはデフォルトでshare=Trueになってるかもなので、
# 外部に公開したくない場合は明示的にFalseを指定することをおすすめします。
demo.launch(server_name="0.0.0.0", server_port=8001, share=False)

URLにアクセスすると、こんな感じの画面が表示されるかと思います。

試しに「データサイエンティストに求められるスキルを5つ挙げて。」と聞いてみたところ、どちらのモデルもいい感じな回答をしてくれてるように感じます。

いくつか質問文を入力してみましたが、どちらのモデルも似たような振る舞いをしており、少なくともOpenCALM-7BをLoRAで今回のデータセットでinstruction tuningする観点では学習方法の差異はないようでした。
であれば実装が楽な実装例2.がおすすめですかね。

モデルとしての振る舞いとしては、わずか1epochでちょっとのパラメータしか更新（差分計算）してないにもかかわらず、それっぽい回答をしてくれてるようです。

もちろんChatGPTには遥か遠く及ばないし、実行のたびに結構回答文が代わります。おまけにtemperatureをいじると振る舞いがガラッと変わる感じもするし、このままだとまた扱いが難しいように感じました。

QLoRAでもやってみる

QLoRAとは、通常では1枚のGPUにモデルが乗らないような巨大なモデルに対しても、以下の3つの技術を使って、メモリ使用量を抑え、LoRAよりもより良くチューニングできる革新的技術です。（以下は論文の内容をChatGPTで翻訳したものをそのまま載せてます。）

4ビットNormalFloat
正規分布データのための情報理論的に最適な量子化データ型で、4ビットの整数と4ビットの浮動小数点数よりも優れた経験的な結果をもたらす。

ダブル量子化
量子化定数を量子化する方法で、平均してパラメータあたり約0.37ビット（65Bモデルで約3GB）を節約する。

ページドオプティマイザー
NVIDIAの統一メモリを使用して、長いシーケンス長のミニバッチを処理する際に発生する勾配チェックポイントメモリのスパイクを避ける。

今回扱っているOpenCALM-7Bをtorch.float16でロードするとGPUメモリを14GBほど消費しますが、4bit NormalFloatやダブル量子化を使ってロードするとわずか5GBほどで済みます。試してませんが、無料のcolabでもLoRAによるチューニングができそうです。

参考文献

上記のhuggingfaceのブログ記事で紹介されているこちらのnotebookがQLoRAを動かす上で参考になります。

1つ注意としては transformers などのライブラリはリポジトリから直接installしたdevバージョンじゃないとまだ動かないっぽいです。上記のブログ記事に書いてある以下のコマンドでQLoRAを動かすために必要な主要ライブラリをインストールしましょう。

pip install -q -U bitsandbytes
pip install -q -U git+https://github.com/huggingface/transformers.git
pip install -q -U git+https://github.com/huggingface/peft.git
pip install -q -U git+https://github.com/huggingface/accelerate.git

私がQLoRAを動かしたときの上記のライブラリのバージョンは以下のとおりでした。

accelerate               0.20.0.dev0
bitsandbytes             0.39.0
peft                     0.4.0.dev0
transformers             4.30.0.dev0

ライブラリが揃えばあとはブログ記事やcolabのtutorialでサクッと動かせます。DatasetクラスやCollatorクラスは本記事で紹介したものがそのまま使えるので下記では記載を省略しています。

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from transformers import DataCollatorForLanguageModeling
from transformers import Trainer, TrainingArguments
from peft import PeftModel, PeftConfig
from peft import get_peft_model, LoraConfig, TaskType
from peft import prepare_model_for_kbit_training

model_id = "cyberagent/open-calm-7b"

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_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id, quantization_config=bnb_config, device_map="auto"
)

model.gradient_checkpointing_enable()

# prepare_model_for_kbit_trainingではパラメータのfreezeを行いながら以下の3つの設定も行う
# 1. レイヤーノームをfp32にキャスト
# 2. 出力埋め込みレイヤーが勾配を必要とするように設定
# 3. 言語モデルのヘッドをfp32にアップキャスト
model = prepare_model_for_kbit_training(model)

# 以降はLoRAのときとほとんど同じ
lora_config = LoraConfig(
    r=8,
    lora_alpha=32,
    target_modules=["query_key_value"],
    lora_dropout=0.05,
    bias="none",
    fan_in_fan_out=False,
    task_type=TaskType.CAUSAL_LM
)

model = get_peft_model(model, lora_config)
model.print_trainable_parameters()
# trainable params: 4194304 || all params: 3654950912 || trainable%: 0.11475678062403483

training_args = TrainingArguments(
        output_dir='./output/qlora',
        save_total_limit=1,
        per_device_train_batch_size=8,
        num_train_epochs=1,
        remove_unused_columns=False,
        logging_steps=20,
        bf16=True,
        dataloader_num_workers=16,
        report_to="none",
)

trainer = Trainer(
        model=model,
        data_collator=collator,
        args=training_args,
        train_dataset=train_dataset,
    )
model.config.use_cache = False

trainer.train()
model.save_pretrained('./output/qlora')

kunishou/databricks-dolly-15k-ja を1epoch回すのにかかった時間は約1時間と、LoRAのときとあまりかわりませんでしたが、GPUメモリは13GBほどしか消費しませんでした。バッチサイズを落とせば無料colabでも動きそう。

学習されたアダプターを読み込むときは以下のように実行してロードできました。

model_id = "cyberagent/open-calm-7b"
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_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id, quantization_config=bnb_config, device_map="auto"
)
model = PeftModel.from_pretrained(model, './output/qlora')

肝心の振る舞いについては、正直LoRAで学習したときと大差ないように感じました。この規模のLLMだとGPUリソースを削減できる、くらいしか恩恵がないのかな？もう少しデータを変えてみたりして検証してみたいです。

おわりに

今流行りのOpenCALMのinstruction tuningをLoRAとQLoRAで試してみました。
LoRAの実装は簡単だし、学習効率もいいし、他のデータセットでも色々試してみたくなりました。

今後LoRAを使ってOpenCALM-7Bをより良いモデルにするために、以下をするとどうなるのか気になります。

• もっとデータを増やすとどうか？
• 逆にデータ量をもっと減らしてでも高品質なデータにしたらどんな振る舞いをするか？
• ハイパラ調整したらどうなる？
• 指示文のテンプレートを変えると振る舞いはどう変わる？
• 特定のドメインに特化したinstruction tuning用のデータセットで学習させてみるとどうか？

おわり