モデル学習後のポスト処理（加工・変換・アップロード）について【Team Kuma】

はじめに

GENIAC 松尾研 LLM開発プロジェクト　Team Kumaの加藤　純です。
今回、Team Kumaで構築したモデルについて、学習後のポスト処理のスクリプト構築を担当しました。
本記事では、その担当作業にあたっての取り組み内容を報告します。

また、チーム全体の取り組み内容については、以下Paper & Hackのイベントにて報告していますので、ぜひご覧ください。

目的

LLaMa2-Accessoryを利用して学習したMegaBlocksのMoE（dropless MoE）モデルを、
Transformerライブラリで利用できるHuggingFace形式のMoE（Mixtral）モデルに変換して公開することがチームの目的になります。
また、個人の目的として、以下を目指して活動しました。

• すべての処理が自動化されること（夜間に実行したい）
• CPUのみで実行できること（限りあるGPUリソースを学習に費やしたい）
• 作成者（加藤）以外が処理を実行できること（手が空いている方に実行してもらいたい）

問題

HuggingFace形式のMixtralモデルへの変換にあたり、以下3点の問題が発生しました。

• dropless MoEにおけるExpert層の重み構造が特殊であること
• LLaMa2-Accessoryにおけるレイヤー名称が特殊であること
• モデル並列での学習により、チェックポイントが分割保存されていること

ポスト処理の内容

• 処理全体のイメージは以下の通りです。（処理時間は合計約40分）
  ※本処理はすべてCPU処理での実行となり、bfloat16に対応したアーキテクチャが必要です。
  

• また、最終的な変換後の学習パラメータ数（training_parameter_size）および推論パラメータ数（active_parameter_size）は、以下の表の通りです。
  

以下にて各処理のポイントを解説します。

加工（マージ）処理

加工処理では、merge_model_weights.pyを作成しました。

公式の分割コードを参考に、以下のような逆（マージ）の処理をおこなっています。

1. Expert層の重みの配置を以下のように変更しました。
   • MegaBlocksの仕様によりGPUごとに分割された状態（＝モデル並列）
     ↓（下図参照）
   • 分割されたExpertの重みをマージして配置（＝エキスパート並列）
     ↓
   • マージした各Expertの重みを1つのファイルに集約（並列なしの状態）

2. マージ後のExpert層のレイヤー名称を以下のように変更しました。（一例）
   • ～.feed_forward.w1（LLaMa2-Accessory形式）
     ↓
   • ～.feed_forward.expert.0.w1.weight（Metaオリジナル形式）
   • ～.feed_forward.expert.1.w1.weight（同上）
   • ～.feed_forward.expert.2.w1.weight
   • ～.feed_forward.expert.3.w1.weight
   • ～.feed_forward.expert.4.w1.weight
   • ～.feed_forward.expert.5.w1.weight
   • ～.feed_forward.expert.6.w1.weight
   • ～.feed_forward.expert.7.w1.weight
   • ～.feed_forward.expert.8.w1.weight
     ※w2、w3も同様

変換処理

変換処理では、convert_tokenizer_and_model_to_hugginface.pyを作成しました。

こちらの変換コードを一部変更し、その変更点は以下になります。

1. 変換処理に不要なLLaMa2-Accessoryのライブラリを利用しないように変更しました。
   • 以下のimport処理、関数を削除

from accessory.util.tensor_parallel import (
    infer_checkpoint_format_and_mp_size,
    load_tensor_parallel_shard_state_dict,
    ShardedTensorLoader,
)
～～
def load_and_merge_tensor_parallel_weights

2. 変換元のレイヤー名称を以下のように変更しました。（一例）

   • llma.layers.～
     ↓
   • layers.～

3. 以下のMixtralのconfigの値を学習したパラメータから反映できるように変更しました。

   • max_position_embeddings　: 32768 ⇒　params["max_seq_len"]
   • num_experts_per_tok　　　: 2 　　　⇒　params["moe"]["num_experts_per_tok"]

アップロード処理

アップロード処理では、upload_tokenizer_and_model_to_huggingface_hub.pyを作成しました。

松尾・岩澤研から提供された標準コードを一部変更し、その変更点は以下になります。

1. アップロード時にHuggingFaceの組織名を選択可能にしました。
   ※Hiroto Nishiyamaさんにご助力いただきました。

#ユーザではなく、組織のリポジトリにアップロード
parser.add_argument("--use_orgs", action='store_true')
#組織のリポジトリ名を指定した場合は、その組織のリポジトリにアップロード
parser.add_argument("--org_name", type=str, default="")
～～
    huggingface_workname =  HfApi().whoami()["orgs"][0]["name"] if args.use_orgs else HfApi().whoami()["name"]
    if args.use_orgs and args.org_name != "":
        # Force to use the specified organization
        huggingface_workname = args.org_name
    repository_name = os.path.join(huggingface_workname, args.output_model_name)

2. スペシャルトークンを追加しました。

def test_tokenizer_and_model(tokenizer, model, prompt_text: str) -> str:
    encoded_prompt_text = tokenizer.encode(prompt_text, add_special_tokens=True, return_tensors="pt").to(model.device)
    with torch.no_grad():
        encoded_generation_text = model.generate(encoded_prompt_text, max_new_tokens=50)
    decoded_generation_text = tokenizer.decode(encoded_generation_text[0], skip_special_tokens=True)
    return decoded_generation_text

3. アップロード後に推論テストを実行するようにしました。

        # Loads and tests the remote tokenizer and the remote model.
        remote_tokenizer, remote_model = load_tokenizer_and_model(repository_name)
        remote_decoded_generation_text = test_tokenizer_and_model(remote_tokenizer, remote_model, args.test_prompt_text)

        # Checks the generated text briefly.
        print(f'|--check : model-XX-of-YY.safetensors')
        print(f"|--{args.test_prompt_text = }")
        print(f"|--{remote_decoded_generation_text = }")
        print()
        if len(remote_decoded_generation_text) <= len(args.test_prompt_text):
            print("Error: The generated text should not be shorter than the prompt text."
                  " Something went wrong, so please check either the remote tokenizer or the remote model.")
            return

検証

1. 加工処理の検証

   • すでに分割されている公式の重み情報を利用して、加工処理⇒分割処理を実施した結果が同じである（復元できる）ことを確認しました。
     ※咸　毅成さんにご助力いただきました。

2. 変換処理の検証

   • 変換前後の推論結果が同じであることを確認しました。
     ※kumagai soichiroさんにご助力いただきました。

3. アップロード処理の検証

   • アップロード前後の重みの値と、推論結果が同じであることを確認しました。

4. 処理全体

   • 作成者（加藤）以外が実行した場合でも、正常に動作することを確認しました。
     ※Susumu Otaさんにご助力いただきました。

おわりに

• 日本での事例が少ないと思われる、LLaMa2-Accessoryフレームワークとdropless MoEアーキテクチャを採用しましたが、無事にチームと個人の目的を達成することができました。
• 再現性担保のため、本スクリプトの検証にご協力いただいたチームメンバーのみなさま、ありがとうございました。
• また、本プロジェクトを通じて貴重な経験を積ませていただき、GENIAC関係者の皆様に心より感謝申し上げます。

---

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