Google Colab：Mergekitによる日本語モデルMoEの作成

0. はじめに

Mistral AIによるMixtral -8x7bモデルの成功により、最近MoE（Mixture of Experts）モデルが流行ってきている気がします。
他にもPhi-2のMoEであるPhixtual-2x2bとかありましたね。
それに倣って、日本語モデルのMoEも作ってみたくなって試してみました。
以下が作ってみたモデルです。

HachiML/youri-2x7b_dev · Hugging Face

1. Configファイルの作成

configファイルにはbase_modelと、expertsを指定する。
positive_promptsには、Expertモデルのそれぞれが得意なプロンプトを入力する。適切なプロンプトが不明だったため、それぞれが得意とする(スコアが高い)JGLUEのタスクのプロンプトを使用した。

base_model: rinna/youri-7b-chat
gate_mode: hidden # one of "hidden", "cheap_embed", or "random"
dtype: bfloat16 # output dtype (float32, float16, or bfloat16)
experts:
  - source_model: rinna/youri-7b-chat
    positive_prompts: 
      - "質問と回答の選択肢を入力として受け取り、選択肢から回答を選択してください。"
      - "前提と仮説の関係を含意、矛盾、中立の中から回答してください。"
      - "以下のテキストを、ポジティブまたはネガティブの感情クラスのいずれかに分類してください。"
      - "以下は、タスクを説明する指示と、文脈のある入力の組み合わせです。要求を適切に満たす応答を書きなさい。"
  - source_model: rinna/youri-7b-instruction
    positive_prompts: 
     - "質問に対する回答を題名と文章から一言で抽出してください。回答は名詞で答えてください。"
     - "与えられたニュース記事を要約してください。"
     - "与えられた文が文法的であるかを回答してください。"

2. Colabでmergekitの実行

# Google Drive　マウント 
from google.colab import drive 
drive.mount('/content/drive')

# 作業フォルダの宣言
base_dir = "/content/drive/MyDrive/moe"

# mergekitのClone
# Mixture of Expertsの作成ができるmixtralブランチを取得する
!git clone -b mixtral https://github.com/cg123/mergekit.git {base_dir}/mergekit

# outputフォルダとconfigフォルダを作成
!mkdir {base_dir}/mergekit/output
!mkdir {base_dir}/mergekit/config

configフォルダに作成済みconfigファイルを配置する。

# ライブラリInstall
!pip install -e {base_dir}/mergekit
!pip install git+https://github.com/huggingface/transformers

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

実行後、作成したoutputにモデルのファイルが出力されている。
マージの実行自体はほとんどモデルの取得の時間で20~30分くらいで実行できた気がする。
メモリがそこそこいるので、ハイメモリで実行しないといけない。

以上です。
出力されたモデルは11.1Bパラメータとなりました。
ほとんどMoEの実装について知らなくても作成できました。
実際に精度はどう変わったかなど、今後調べていければと思っています。

今後、MoEやマージモデルはより一層流行っていく気がするので参考になればと思います。

3. 参考