LightChatAssistant 2x7B を再現する - ローカルLLM自由帳

前回Chat Vectorについて簡単に予習したので、とりあえず「LightChatAssistant 2x7B」の作成手順を再現してみたいと思います。
作者さんがモデルカードで丁寧に説明してくださっているので、基本的にそれをなぞるだけです。まずはまったく同じモデルを作ってみます。

Chat Vectorによる処理

Google ColabのCPU（ハイメモリ）で試します。3つ分のモデルファイルを扱えるだけのRAM容量さえあればよく、GPU無しで完結するようです。
まずは、Chat Vector処理を行うための依存関係をインストールします。ColabだとTransformersはプリインストールされてます。

# 依存関係のインストール
!pip install accelerate protobuf

さてChat Vectorは、追加学習したチャットモデルの重みから、対応するベースモデルの重みを差し引くことで得られる差分です。
ここでは、ベースモデルは"Mistral-7B-v0.1"、チャットモデルに"Mistral-7B-Instruct-v0.2"を使っています。いずれもMistralの公式モデルで、元論文を踏襲した処理です。
それぞれのモデルをHuggingFaceからロードします。

from transformers import AutoModelForCausalLM
import torch

# ベースモデル＝"Mistral-7B-v0.1"
base_model = AutoModelForCausalLM.from_pretrained(
    "mistralai/Mistral-7B-v0.1",
    torch_dtype=torch.bfloat16,
    device_map="cpu",
)
# 差分を取り出すチャットモデル＝"Mistral-7B-Instruct-v0.2"
inst_model = AutoModelForCausalLM.from_pretrained(
    "mistralai/Mistral-7B-Instruct-v0.2",
    torch_dtype=torch.bfloat16,
    device_map="cpu",
)

つづいて、取り出したChat Vectorを統合するための日本語モデルを用意します。
1つ目は"chatntq-ja-7b-v1.0"です。Mistral7B系モデルの中では相対的に高い日本語チャット性能を持っているモデルです。
すでに日本語でチャットファインチューンされたモデルですが、英語モデルのChat Vectorを適用してさらに強化を図る形になります。

# 差分を張り付ける日本語モデル（１）＝"chatntq-ja-7b-v1.0"
cp_model = AutoModelForCausalLM.from_pretrained(
    "NTQAI/chatntq-ja-7b-v1.0",
    torch_dtype=torch.bfloat16,
    device_map="cpu",
)

あとは以下のようにChat Vectorを抽出します。この手順は、jovyanさんのこちらの記事で詳しく紹介されているものです。大変ありがたいです。

# Chat Vectorを取り出して"chatntq-ja-7b-v1.0"へコピー
for k, v in cp_model.state_dict().items():
    chat_vector = inst_model.state_dict()[k] - base_model.state_dict()[k]
    new_v = v + ( 0.8 * chat_vector.to(v.device) )
    v.copy_(new_v)

# Chat Vectorを張り付けた"chatntq-ja-7b-v1.0"を保存
cp_model.save_pretrained("./chatntq-ja-7b-v1.0-chatvector")

なお、元論文ではChat Vectorを掛ける強度として1.0または0.5を試しています。1.0だと差分の影響が強すぎて英語が混ざりやすくなり、逆に0.5だと差分の影響が表れにくくなる可能性があるようです。
- おそらくそうした理由からLightChatAssistant 2x7Bを含むいくつかのChatVectorモデルで0.8掛けという数値を採用しているのだと思います。
さて、もう一つの日本語モデルは小説生成モデルの"Antler-7B"です。このモデルの学習データにNSFWが含まれているのでご注意ください。
手順は先ほどと同じで、今度はAntler-7BにChatVectorを掛けます。こちらはチャットファインチューンされていないモデルにチャット能力を付与するという元論文の意図に近い形になっています。

# 差分を張り付ける日本語モデル（２）＝"Antler-7B"
cp_model2 = AutoModelForCausalLM.from_pretrained(
    "Elizezen/Antler-7B",
    torch_dtype=torch.bfloat16,
    device_map="cpu",
)

# Chat Vectorを取り出して"Antler-7B"へコピー
for k, v in cp_model2.state_dict().items():
    chat_vector = inst_model.state_dict()[k] - base_model.state_dict()[k]
    new_v = v + ( 0.8 * chat_vector.to(v.device) )
    v.copy_(new_v)

# Chat Vectorを張り付けた"Antler-7B"を保存
cp_model2.save_pretrained("./Antler-7B-chatvector")

これで2つのChat Vector適用済みモデルが用意できました。
Tokenizerはそれぞれの日本語モデルから以下のように流用します。

from transformers import AutoTokenizer

cp_tokenizer = AutoTokenizer.from_pretrained("NTQAI/chatntq-ja-7b-v1.0")
cp_tokenizer.save_pretrained("./chatntq-ja-7b-v1.0-chatvector")

cp_tokenizer2 = AutoTokenizer.from_pretrained("Elizezen/Antler-7B")
cp_tokenizer2.save_pretrained("./Antler-7B-chatvector")

最後に、各モデルのconfig.jsonの修正が必要です。
差分を取り出した"Mistral-7B-Instruct-v0.2"のコンテキスト関係の設定値が他の"Mistral-7B-v0.1"系モデルとは異なっているために必要な処理のようです。
- "Mistral-7B-Instruct-v0.2"は、v0.1系モデルで採用していたSliding Windowという仕組みが無効になっています。
- そのためChat Vector適用したモデルも"Mistral-7B-Instruct-v0.2"のコンテキスト設定に合わせる必要が生じます。
具体的には、"max_position_embeddings" = 32768、"rope_theta" = 1000000.0、"sliding_window" = None、に修正します。
Claudeに投げたところ以下のような修正用コードを作ってくれました。

import json

with open("./chatntq-ja-7b-v1.0-chatvector/config.json", "r", encoding="utf-8") as f:
    config = json.load(f)

config["max_position_embeddings"] = 32768
config["rope_theta"] = 1000000.0
config["sliding_window"] = None

with open("./chatntq-ja-7b-v1.0-chatvector/config.json", "w", encoding="utf-8") as f:
    json.dump(config, f, indent=2, ensure_ascii=False)

with open("./Antler-7B-chatvector/config.json", "r", encoding="utf-8") as f:
    config = json.load(f)

config["max_position_embeddings"] = 32768
config["rope_theta"] = 1000000.0
config["sliding_window"] = None

with open("./Antler-7B-chatvector/config.json", "w", encoding="utf-8") as f:
    json.dump(config, f, indent=2, ensure_ascii=False)

ちなみに、最初に試したときにこのconfig.jsonの処理をスキップしてしまったのですが結果的にモデル出力の質がかなり劣化する形となりました。どうやらロングコンテキストの処理に影響するだけでは留まらないようです。
以上で、Chat Vector関係の処理は終わりです。

MergeKitによるMoEマージ

では先ほど作成した2つのChat Vector適用モデルをMergeKitでマージします。
通常版のMergeKitではMoEマージに対応していないので、Mixtralブランチを指定して導入する必要があります。
MoEマージの手順については、以前も紹介したはちさんの記事で詳しく解説されているものを参考にしています。

# メモリ容量の確保
del cp_model, cp_model2

# MoE対応版のMergeKitの導入
!git clone -b mixtral https://github.com/cg123/mergekit.git
!mkdir ./mergekit/output
!mkdir ./mergekit/config

!pip install -e ./mergekit
!pip install git+https://github.com/huggingface/transformers #セッションは再起動しない

なおColabでは、途中でセッションの再起動を促されますが無視してキャンセルします。
MoEマージのためのConfigは、LightChatAssistant-2x7BのRepoに上がっていますのでここではそのまま借用します。内容的にNSFWです。

# LightChatAssistant-2x7BのRepoからマージ設定をDL（NSFWにつき注意）
!wget https://huggingface.co/Sdff-Ltba/LightChatAssistant-2x7B/resolve/main/mergekit_moe_config.yml -P ./mergekit/config

あとは2モデルのマージを実行するだけです。完了後、outputフォルダに2x7Bのモデルファイルが格納されます。

# MoEマージの実行
!python ./mergekit/mergekit/scripts/mixtral_moe.py ./mergekit/config/mergekit_moe_config.yml ./mergekit/output -v

llama.cppによる量子化

以上で「LightChatAssistant 2x7B」の作成手順の再現が完了しました。
とはいえ実際にローカルでモデルを使用する場合、そのまま使う人はまずいないと思いますのでllama.cppのGGUFフォーマットに変換しておきます。

# llama.cppの導入（依存関係はモデル変換用のみインストール）
!git clone https://github.com/ggerganov/llama.cpp
!pip install -r ./llama.cpp/requirements/requirements-convert.txt

# Q8_0量子化
!python ./llama.cpp/convert.py ./mergekit/output --outfile LightChat-2x7B-Repro-q8.gguf --outtype q8_0

Q8量子化後で、12.8GBほどのファイルサイズになりました。もっと圧縮する場合、f16に変換した後でmakeしてquantizeコマンドから希望の量子化タイプに変換することになります。

再現モデルのテスト

再現したモデルを、元の「LightChatAssistant 2x7B」と比較してみます。
このモデルを、たびたび利用している「ELYZA-tasks-100」ベンチマークにかけてGPT-4に評価させてみました。

名称	LightChatAssistant 2x7B【オリジナル】	LightChatAssistant 2x7B【再現】	ChatNTQ JA 7B
パラメータ	2x7B	2x7B	7B
量子化タイプ	Q8_0	Q8_0	Q8_0
ElyzaTasks100 スコア平均	3.31	3.29	3.06
標準誤差	0.14	0.15	0.15
平均回答字数	273.61	277.00	240.46