日本語が話せるLlamaモデルをDIYする

2023/7/27追記: 項目2.1でrun_clm_pt_with_peft.pyを編集する必要がありますが書き忘れていました。編集箇所と方法を追記しました。一部情報が正しくない部分もあったので修正しました。

この記事ではMetaがリリースしたLlama2-7Bを元に日本語が話せるLlamaモデルを作る方法を解説しています。

日本語Llamaを作ってみようと思ったきっかけはこちらのgithubリポジトリです。このリポジトリは中国の研究者たちがLlamaモデルを中国語に対応させるために使ったもので、流れとしてはまずLlamaのトークナイザーを拡張し中国語に対応させて、LoRAで二次事前学習を行い、その後インストラクト用に学習させます。(LoRAはQLoRAではなく普通のLoRAです。)
これ自体はLlamaが中国語を話せるようにするという目的で存在していますが、おそらく中国語以外の言語にも使えると思われます。
実際に私もLlama2-7Bを元に日本語wiki40bを学習させてみましたが、学習前と比べて確かに日本語が話せるようになっていると思いました。ただ素人が夏休みの自由研究程度にやってみただけなので、本当に日本語が話せるだけで話す内容はめちゃくちゃです。
今回私が作ったモデルはHuggingfaceにfp16版とggml版をアップロードしてあります。

作成した日本語Llamaの出力例

プロンプト: 江戸幕府は
結果: 江戸幕府は、徳川家康の養子・千代松（後の嫡男・徳川秀忠）に「千代松」から改名するという要望から、千代松の養母が「千代松」から「千代田」へと改名して、そこから「千代田」と命名されたとも、千代田の名は、千代松の養母である北条氏真、或いは、秀忠の娘婿であるお登（おたけ）に因むとする説がある。

プロンプト: 周期表において
結果: 周期表において、赤色矮星は青色矮星と同様に、1958年の1月1日の恒食以来、156日間の一定周期で起きる。また、赤色矮星が出現する際にも、赤色矮星以外の恒星とは色質が異なるにもかかわらず、赤色矮星と同時代の天体でも明るく見えない場合もある。例として、1月7日から12月30日までは、南天の赤色矮星であるアリーヤ（Aldebaran）が見えなくなる。これは、アリーヤ以外に赤色派の恒星があるかという問題があるが、アリーヤからその周期で出現することとなる、赤色惑星が存在するはずという主張もある。赤色矮星は、赤色族を形成している恒星の中でも最も光度が高いため、最も明るい恒星である。これは、アリーヤ（Aldebaran）やアリーア（Alpha Tauri）のように、赤色族が最も多くある星座のような赤色矮星で、その大きさからも明るさからも最も明るい恒星となる。また、赤色族では光度が高いの赤色矮星に限らず、光度が小さい赤色派の恒星もある。

内容はさておき同じ単語を繰り返している感がありますね。

モデルを完成させるまでの流れは次のようになります。

1 元のLlama2のトークナイザーを日本語用に拡張する。
2 ベースのLlama2(chatではない方)を日本語のプレーンテキストで二次事前学習させ、日本語のベースLlama2モデルを作る。
3 日本語ベースLlama2をインストラクト用に学習させる。

最初は3のインストラクト用学習までやろうと思っていましたが、2でできたベースモデルの出来が期待外れで頓挫しましたorz

1. トークナイザーの拡張

そのままのトークナイザーは一部の日本語常用漢字を認識できません。
Llamaのトークナイザーに日本語を変換させると以下のようになります。いくつかの漢字が正しく変換されていないことがわかります。また、平仮名も含めて一字ずつしか処理されていません。
与えられた文章を一字ずつ処理するよりも、頻出する文字列をまとめて一つのものとして処理した方が言語モデルの性能は良くなります。

元の文章: ショーペンハウアーによれば、世界とは「わたしの表象」であり、現象にほかならない。つまり、主観である意志に対応する客観としてのみ、世界が存在する。時間・空間・因果関係においてある現象に対して、カントのたてた物自体とは実は意志そのものにほかならない。
変換後: ['▁', '<0x0A>', 'シ', 'ョ', 'ー', 'ペ', 'ン', 'ハ', 'ウ', 'ア', 'ー', 'に', 'よ', 'れ', 'ば', '、', '世', '界', 'と', 'は', '「', 'わ', 'た', 'し', 'の', '表', '象', '」', 'で', 'あ', 'り', '、', '現', '象', 'に', 'ほ', 'か', 'な', 'ら', 'な', 'い', '。', 'つ', 'ま', 'り', '、', '主', '<0xE8>', '<0xA6>', '<0xB3>', 'で', 'あ', 'る', '意', '志', 'に', '<0xE5>', '<0xAF>', '<0xBE>', '<0xE5>', '<0xBF>', '<0x9C>', 'す', 'る', '客', '<0xE8>', '<0xA6>', '<0xB3>', 'と', 'し', 'て', 'の', 'み', '、', '世', '界', 'が', '存', '在', 'す', 'る', '。', '時', '間', '・', '空', '間', '・', '因', '果', '関', '<0xE4>', '<0xBF>', '<0x82>', 'に', 'お', 'い', 'て', 'あ', 'る', '現', '象', 'に', '<0xE5>', '<0xAF>', '<0xBE>', 'し', 'て', '、', 'カ', 'ン', 'ト', 'の', 'た', 'て', 'た', '物', '自', '体', 'と', 'は', '実', 'は', '意', '志', 'そ', 'の', 'も', 'の', 'に', 'ほ', 'か', 'な', 'ら', 'な', 'い', '。', '<0x0A>']

以下は日本語を認識できるように拡張したトークナイザーを使って同じ文章を変換させたものです。漢字を認識できており、「世界」などの文字列が一つのものとして処理されています。

['▁', '<0x0A>', 'ショー', 'ペン', 'ハ', 'ウ', 'アー', 'に', 'よ', 'れば', '、', '世界', 'とは', '「', 'わ', 'た', 'し', 'の', '表', '象', '」', 'であり', '、', '現', '象', 'に', 'ほか', 'なら', 'ない', '。', 'つ', 'まり', '、', '主', '観', 'である', '意', '志', 'に', '対応', 'する', '客', '観', 'として', 'のみ', '、', '世界', 'が', '存在', 'する', '。', '時間', '・', '空', '間', '・', '因', '果', '関係', 'に', 'お', 'いて', 'ある', '現', '象', 'に', '対', 'して', '、', 'カン', 'ト', 'の', 'た', 'て', 'た', '物', '自', '体', 'とは', '実', 'は', '意', '志', 'その', 'も', 'の', 'に', 'ほか', 'なら', 'ない', '。', '<0x0A>']

2. 二次事前学習

インストラクトやチャット用にあらかじめ学習させられたモデルをベース(completion)モデルとして学習させようとしても、うまくいきません。また、モデルをいきなり日本語でインストラクト・チャット用に学習させるよりも、最初にプレーンテキストで学習させてある程度基礎的な日本語能力を向上させてからのほうが良いと思われます。
中国語Llamaのリポジトリでは二次事前学習にファインチューニングではなくLoRAを使っています。7BモデルはcolabのA100(40GB)でも学習できるようです。
今回は学習させるデータとしてwiki40bを採用しました。前処理済みデータでサイズも1GB程度とちょうどよかったのが理由です。こちらからダウンロードできるparquetファイルを元にtxt形式のデータを用意しました。txtファイルに直したものはこちらにアップロードしました。

3.インストラクト用学習

自分もここまではやっていませんが、もともと比較的クオリティが良いと思った日本語Dollyデータセットを使ってqloraで学習させるつもりでした。
インストラクトやチャット用のファインチューニング、LoRA学習に使うデータセットはさほど多くなくて良いと言われています。例えば質の良いQAペアのデータを1000個ほど用意して事前学習済みモデルを学習させた場合でも、大量のデータを学習させた場合と結果はさほど変わらないようです。これはLIMAの論文で解説されています。

ここから具体的な方法を解説していきます。

1.1 日本語データからトークナイザーを作る

まずトークナイザーの拡張から始めますが、最初に日本語専用のトークナイザーを作ります。後でこの日本語専用トークナイザーとLlamaのトークナイザーを合体させてから新しいLlama用のトークナイザーを作り、これを学習に使います。

日本語のテキストファイルを用意します。学習させるデータセットそのものでも何でもいいですが、できるだけ多くの日本語のボキャブラリーを認識できるようにすることを念頭に置きましょう。
テキストファイルを用意したら次のコードを実行します。vocab_sizeはテキストファイルの規模などに合わせて調整してください。

make_tokenizer.py

import sentencepiece as spm

# テキストデータのパス
input_file = "path/to/txt"

# 学習済みモデルの出力パス
model_prefix = "tokenizer"

# ハイパーパラメータの設定
vocab_size = 15000#トークンの数
model_type = "unigram"  # モデルのタイプ
character_coverage = 1.0 # 文字のカバレッジ

# 学習の実行
spm.SentencePieceTrainer.train(
    input=input_file,
    model_prefix=model_prefix,
    vocab_size=vocab_size,
    model_type=model_type,
    character_coverage=character_coverage
)

1.2　トークナイザーを合体させる

冒頭のリポジトリを持ってきて必要なライブラリをインストールします。

git clone https://github.com/ymcui/Chinese-LLaMA-Alpaca
cd Chinese-LLaMA-Alpaca
pip install -r requirements.txt

scripts/merge_tokenizer直下に1.1で作ったトークナイザーを置き、llama_tokenizerフォルダを作り、Llamaのtokenizer.modelとtokenizer_config.jsonを置きます。
scripts/merge_tokenizerフォルダに移動します。

merge_tokenizer.pyを実行します。

python merge_tokenizer.py --llama_tokenizer_dir llama_tokenizer --chinese_sp_model_file tokenizer.model

merged_tokenizer_hfフォルダに拡張されたトークナイザーができているはずです。

2.1 二次事前学習

scripts/trainingフォルダに移動します。run_pt.shを編集します。
pretrained_modelに元となるLlama2-7Bモデルのパスを指定します。huggingfaceのリポジトリ名でも良いです。
chinese_tokenizer_pathには1.2で作った拡張済みトークナイザーのパスを指定し、dataset_dirに学習させるtxtファイルが置かれているフォルダのパスを指定します。フォルダに複数のtxtファイルを含めても良いようです。

script/training/run_clm_pt_with_peft.pyの550行目あたりに次のようなコードがあると思います。ここの49953という数字を、1.2で作ったトークナイザーのボキャブラリー数に書き換えます。

run_clm_pt_with_peft.py

    if not (
       (model_vocab_size==32000 and len(tokenizer)==49953) or \ #ここの49953を書き換える
       (model_vocab_size==32000 and len(tokenizer)==32000) or \
       (model_vocab_size==49953 and len(tokenizer)==49953) or \
       (model_vocab_size==49954 and len(tokenizer)==49954)
    ):

トークナイザーのボキャブラリー数はこのコードで確認できます。

check_vocab_size.py

import sentencepiece as spm
tokenizer = spm.SentencePieceProcessor()
tokenizer.load("path/to/your/tokenizer.model")
vocab_size = len(tokenizer)
print(vocab_size)

後はbash run_pt.shを実行するだけです。学習が進めばLoRAモデルが保存されていきます。

2.2　LLamaとLoRAをマージする

scriptフォルダのmerge_llama_with_chinese_lora.pyを実行します。作成したLoRAを普通にLlamaに適用するだけでは効果がなく、このスクリプトでマージする必要があります。

python merge_llama_with_chinese_lora.py \
    --base_model path/to/llama/model \
    --lora_model path/to/lora/model \
    --output_type huggingface \
    --output_dir path/to/output/dir

これで日本語Llamaができました！完成したモデルをggmlやGPTQに変換して自宅のPCで試してみるのもいいかも？
自分がやったときはLoRAである程度学習させてから同じデータでQLoRAに学習させるということをやりましたが正直意味があることだったのかわかりませんw(lossが下がらなかったです)
日本語を学習できるか確認したいだけだったので、自分のモデルは3時間程度学習させてから切り上げました。colab proを購入していてコンピューティングユニットが50個くらいあるならA100で同じことができると思います。

終わりに

中国語Llamaでは二次事前学習に最大120GBの中国語データが使われています。データを増やしたら当然性能は良くなりますがその分学習時間も延びるのでとにかく費用がかかりそう。
元の論文ではA100を最大で48個使って学習させたと書いてありました。GPU代も馬鹿になりませんね...

現状、言語モデルは英語・中国語モデルが独走している状態です。これは、英中双方とも話者数の規模が圧倒的に大きく、データを集めるのに苦労しないという理由もあると思います。
事前学習に使われるデータセットも、英語ではC4やcc100などのデータセットから、PileやRedPajamaなどの量だけでなくクオリティも重視したデータセットに移行しています。そしてこれらはすべてオープンソースで、企業、研究者から個人まで誰でも使えます。日本語の言語モデルを作る時にハードルとなるのは、必要な（GPUやサーバーなどの）ハードウェアリソースを用意するための予算や技術力はもちろん、すぐに使えるオープンソースでクオリティの高いデータセットがないことが、英語モデルを開発する場合に比べて大きな壁となっていると思います。