Tanuki-8x8B-vision-exp学習作業記録・tips【前編】

はじめに

松尾・岩澤研究室のGENIACプロジェクト^[1]において，Tanukiモデルを画像入力にも対応させたマルチモーダルモデル「Tanuki-8B-vision"」^[2]および「Tanuki-8x8B-vision-exp」^[3]の開発に携わり，公開することができた（なお，マルチモーダル開発は本プロジェクトの主目的ではなく，私を含む有志の参加者により実験的に行われたものである）．その中で，私は特にTanuki-8x8B-vision-expの開発を中心として行った．
この記事では，2回に分けてTanuki-8x8BをLLaVA-JPのコードを用いてLLaVA化してVisionモデルを構築・学習した際の作業記録や，注意するべきポイントを公開する．
前編は，Tanuki-8x8B-visionの学習を行う前の準備段階について．

コードリポジトリ

https://github.com/matsuolab/nedo_project_code/tree/team_hatakeyama_phase2/team_hatakeyama_phase2/multimodal/LLaVA-JP

モデル構築・学習の詳細

何が難しかったか？

• 総パラメータ数が約47Bであり，ベースの言語モデルをスケールすることによる学習コストの増加．
• Tanuki-8x8Bはカスタムクラスを採用しており，それに対応できるようにLLaVA-JPのコードを修正する必要があった．
  • ベースとしたリポジトリは https://github.com/tosiyuki/LLaVA-JP をフォークした https://github.com/hibikaze-git/LLaVA-JP
  • 作業用： https://github.com/Osakana7777777/LLaVA-JP （public archive済）
• 実際に学習を始めてみると，思いのほか時間がかかったほか，様々なトラブルが発生した．計算リソースの都合により，8x8Bのモデルは1 epoch学習を終わらせることはできなかった（そのため，チェックポイントの実験的なリリースという意味でモデル名に「-exp」がついている．なお，同じデータを用いて8Bの方は学習が完了できている）．

今回の内容：学習の前準備

Tanuki対応のためコードを書きかえ

上述のように，Tanuki-8x8BはTanukiForCausalLMという独自のクラスが用いられている．
https://huggingface.co/weblab-GENIAC/Tanuki-8x8B-dpo-v1.0/blob/main/modeling_tanuki.py
そのため，こちら^[4]を参考にして，LLaVA関係のクラスが独自クラスを継承できるようにllava_tanuki.pyを作成した．

https://github.com/matsuolab/nedo_project_code/blob/team_hatakeyama_phase2/team_hatakeyama_phase2/multimodal/LLaVA-JP/llavajp/model/llava_tanuki.py

さらに，llava_arch.pyにもTanuki用のembeddingの処理の呼び出しを追加した（メソッド自体はllamaやmixtralと同じであった）．
なお，mixtral（ if001/tiny_mixtral_ja ^[5] を使用させていただいた）にも対応させているのは，小規模なモデルかつMoEで学習がうまくできるかのテストのためである（後編で詳述）．
https://github.com/matsuolab/nedo_project_code/blob/team_hatakeyama_phase2/team_hatakeyama_phase2/multimodal/LLaVA-JP/llavajp/model/llava_arch.py#L96-L106

注意点

llava_tanuki.pyにおいて，TanukiForCausalLM, TanukiModel, TanukiConfigは，リポジトリの中にconfiguration_tanuki.pyとmodeling_tanuki.pyを直接配置することによりインポートエラーを回避した．
https://github.com/matsuolab/nedo_project_code/tree/team_hatakeyama_phase2/team_hatakeyama_phase2/multimodal/LLaVA-JP/llavajp/model/tanuki
customコード（configuration_○○.pyやmodeling_○○.py）を含んだHugging Faceのモデルをロードする際，AutoModelForCausalLMの引数trust_remote_codeをTrueに設定する必要がある^[6]．
https://github.com/matsuolab/nedo_project_code/blob/team_hatakeyama_phase2/team_hatakeyama_phase2/multimodal/LLaVA-JP/train_llava.py#L191-L197
上記のようにconfiguration_tanuki.pyとmodeling_tanuki.pyを入れていないと，Tanuki-8x8Bの（事前学習途中の）チェックポイントを本番環境内からロードする際にモデルがうまく読み込めない事象が発生した．具体的には，modeling_tanuki.py（チェックポイントの中のファイル）で以下のような相対インポートが上手くいかなかった．

modeling_tanuki.py

from .configuration_tanuki import TanukiConfig

ImportError: attempted relative import with no known parent package

この原因の詳細については，正確に把握できていない．

次回

後編では，実際にTanuki-8x8B-visionの学習を詳述する．環境構築法，遭遇したトラブルとその解決策（思い出せる限り）を提示し，学習における気づきとFuture workを考察する．

謝辞

• LLaVA-JPの作者tosiyuki様，およびコードをわかりやすく整備したり機能を追加したりしてくださったhibikaze様に心より感謝申し上げます．

この成果は、NEDO（国立研究開発法人新エネルギー・産業技術総合開発機構）の助成事業
「ポスト５Ｇ情報通信システム基盤強化研究開発事業」（JPNP20017）の結果得られたものです。

脚注

1. https://weblab.t.u-tokyo.ac.jp/geniac_llm/

2. https://huggingface.co/weblab-GENIAC/Tanuki-8B-vision

3. https://huggingface.co/weblab-GENIAC/Tanuki-8x8B-vision-exp

4. https://qiita.com/toshi_456/items/248005a842725f9406e3#llavaのコードを日本語モデル向けに改修する

5. https://huggingface.co/if001/tiny_mixtral_ja

6. https://github.com/matsuolab/nedo_project_code/blob/team_hatakeyama_phase2/team_hatakeyama_phase2/multimodal/LLaVA-JP/train_llava.py#191