LoRA のもう一つの大きなメリット、GPUメモリ共有しつつ別のタスク処理モデルへ即時に切り替える方法
低ランク行列を追加することで、大元のモデルを維持しつつ少ないコストで学習できる LoRA(Low-Rank Adaptation of Large Language Models)。先日、日本語でも大規模パラメータモデル cyberagent/open-calm-7b や rinna/japanese-gpt-neox-3.6b 等々がリリースされたり、HuggingFaceからはtransformersでLoRAを簡単に実現できるライブラリ、peft がリリースされたことで、試したことがある方も多いと思います。
ただ、LoRAのメリットについて主に学習の話が殆どで、もう一つの大きなメリットであるLLMのベースモデルのメモリを共有しつつ、複数のタスクをこなす方法の紹介は見かけたことがなかったので、それをpeftで行う方法についてのお話です。
なお、LoRAとは何か？というお話は、輪講資料 LoRA: Low-Rank Adaptation of Large Language Modelsの資料に大変わかりやすくまとまっています。
何が問題なのか・どんな問題が解決されるのか
LLMは名前の通り巨大な言語モデルなので、例えばopen-calm-7bモデルはfp16でGPUにロードすると、それだけで約13GBのメモリを使ってしまいます。そのため、完全なfinetuneした場合、そのタスクをこなすだけで13GBのメモリが必要で、別のタスクをこなすモデルをロードしようとすると、さらに13GBかかってしまいます。26GBのメモリは、とりわけご家庭のGPUではキツいメモリサイズですね。
しかしながら、LoRAで open-calm-7b を低ランク行列をパラメータr=8で追加して学習した場合、追加で必要になるメモリサイズはたったの17MBです。17GBではなく17MBで、特徴を学習させた(=なんらかのタスクが解ける)ニューラルネットワークができあがるのです。
とすると、ベースのLLMモデルの13GBに加え、+17MBで別のタスクをこなすことができます。そしてこれは一つでなく、たとえば10個別々のタスクや学習データから学習させたLoRAのデータを使えば、13GB + 170MB のメモリ使用量だけで、それらのタスクをこなすことができるようになります。めちゃすごい！！！
正直、バッチ処理でたくさんのデータにたいして同じ処理をする場合などは、GPUのメモリロード・アンロードを繰り返してもその時間を待てばよいのですが、リアルタイムに逐次処理したい、たとえばユーザからの入力に対して応対する、みたいなタスクにおいては1GPUでメモリを共有しつつ、複数のタスクをこなせるととてもパフォーマンスが良いですよね。
例えば、ぱっと考えただけでも以下などにに使えそうですね。とりわけGPUはランニングコストがかかるので、できることなら1GPUで色々な処理をさせたい。
- チャットボット応対の表現文章を変えたい
- 記事ホスティングサービスをしているので、ユーザの文章特徴を学習させてユーザによってモデルを切り替えたい
- A/Bテストでどちらの学習が良かったかを、切り替えて評価したい
- LangChain の Agent をローカルで動かすときに、Agent をサクサク切り替えたい
- Agent ごとに色々な機能を持っていて、内容によってAgentを切り替えたい欲求がある。ただそのAgentが巨大なモデルだと、しょっちゅうメモリへのロード・アンロード(メモリ解放)が発生してとても遅くなる
ただ、LLMのベースとなるモデルは同じ必要があるので、その点は注意です。
実際にどう切り替えるのか
peft を使って学習したLoRAモデルを利用する場合、切り替えはとても簡単です。例として、以下の notebook を用意しました。
PeftModelは、adapter という機能により、動作するモデルを切り替えることができます。標準でロードしたものは default という名称で、load_adapter(model_name, adapter_name) で名前をつけて別のモデルをロードできます。
例えば、こんな感じで peft_model をロードします。
from peft import PeftConfig, PeftModel
peft_model_open2ch = "hotchpotch/open-calm-7b_lora_open2ch"
peft_config_open2ch = PeftConfig.from_pretrained(peft_model_open2ch)
model = AutoModelForCausalLM.from_pretrained(peft_config_open2ch.base_model_name_or_path, device_map="auto", torch_dtype=torch.float16)
peft_model = PeftModel.from_pretrained(model, peft_model_open2ch)
別の機能を持ったモデルを追加します。
# https://note.com/masuidrive/n/n0e2a11fc5bfa
peft_model_instruct = "masuidrive/open-calm-instruct-lora-20230525-r4-alpha16-batch32-epoch1"
# instruct という adapter 名として、peft_model にロードさせる
peft_model.load_adapter(peft_model_instruct, "instruct")
あとは、処理に応じて adapter を切り替えるだけです。
# open-calm-7b_lora_open2ch の学習モデル
peft_model.set_adapter("default")
# masuidrive/open-calm-instruct-lora-20230525-r4-alpha16-batch32-epoch1 の学習モデル
peft_model.set_adapter("instruct")
これで、LLMのベースとなる cyberagent/open-calm-7b が約13GBのメモリに展開され、2ch風文章を生成するモデルが adapter の default に、qaの回答を生成できる adapter が instruct というモデルに、両方合わせて34MBのメモリ追加のみで展開できていると思います。
そのため、実行したいタスクに合わせて set_adapter で切り替えるだけで、再び巨大なLLMをロード・メモリ解放することなくうまく使うことができるでしょう。notebook例だと「2ch風文章を作成」と「質問への回答」、2つの機能を持った model を、切り替えながら使っています。
LLM + 様々なアダプターで夢が広がる
今までは full finetune が必要だった巨大パラメータモデルの学習が、LoRAで学習が効率的に保存データサイズも小さく、かつ推論もメモリを共有することで複数のタスクを省メモリでこなすことができるようになってきました。
この辺の分野も日進月歩で日々進化していて、めちゃくちゃ面白いし出ることも広がってきていて、将来が楽しみですね!!1
なお記事タイトルには「GPUメモリ」と書きましたが、別にGPUに限らずメモリ共有できるはずです。
