Zenn
東大松尾・岩澤研究室 | LLM開発 プロジェクト[GENIAC]
🐻

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

jaked97
2024/10/17に公開
LLM
大規模言語モデル
GENIAC
松尾研
tech

はじめに

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つのファイルに集約（並列なしの状態）

  1. マージ後の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

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

    • llma.layers.～
      ↓
    • layers.～

  2. 以下の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)
  1. スペシャルトークンを追加しました。
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
  1. アップロード後に推論テストを実行するようにしました。
        # 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関係者の皆様に心より感謝申し上げます。
東大松尾・岩澤研究室 | LLM開発 プロジェクト[GENIAC]
東大松尾・岩澤研究室 | LLM開発 プロジェクト[GENIAC]

東京大学 松尾・岩澤研究室が運営する松尾研LLMコミュニティのLLM開発プロジェクト[GENIAC] の開発記録、情報発信になります。 各種リンクはこちら linktr.ee/matsuolab_community

Discussion