自作アーキテクチャのモデルをHuggingFaceにプッシュする方法

はじめに

現在、私は以下のような試みをしています。

その中で、自作アーキテクチャ（Transformersに実装されていない）モデルをHuggingFaceにpushすると、当たり前ではありますがそのアーキテクチャでは重みをloadできないことに気づきました。
なので、

1. 自作アーキテクチャ（Transformersに実装されていない）モデルをHuggingFaceにpush

2. できる限り簡単にAutoClassを使って自作アーキテクチャでモデルをload

できる様にする方法を試したのでここにまとめます。

1. 全体像

実施すべきことは簡単で、主に以下の2つになります。

1. アーキテクチャ定義ファイルの用意

2. AutoClassへの登録

その後、通常のモデル（Transformersに実装されているアーキテクチャのモデル）と同様、HuggingFaceへpushすれば他の人でもそのアーキテクチャでモデルをloadすることが可能となります。
順番に見ていきます。

2. アーキテクチャ定義ファイルの用意

すでに自作のアーキテクチャのモデルを作れていればファイルの作成はそれほど難しくはないです。重要な点は以下の3点です。

• Transformersが提供しているPreTrainedConfig, PreTrainedModelを自作のconfig, modelに継承させる。

• model_type, config_classの設定

• 設定ファイル名はmodeling_xxx(model_type).pyとする。

2-1. PreTrainedConfig, PreTrainedModelの継承

説明通り、自作のconfig, modelに上記を継承させます。
これらは以下でimport 可能です。

from transformers import PreTrainedConfig, PreTrainedModel

私が作成したbit_llama（LinearをBitLinearに置き換えたLlamaモデル）では、LlamaをベースとしたためLlamaConfig, LlamaForCausalLMを継承していますが、LlamaConfigはPreTrainedConfigを、LlamaForCausalLMはPreTrainedModelを継承しているモデルなので問題ありません。

class BitLlamaConfig(LlamaConfig):
    model_type = "bit_llama"

    def __init__(self, bits=8, **kwargs):
        super().__init__(**kwargs)
        self.bits = bits
:
(中略)
：
class BitLlamaForCausalLM(LlamaForCausalLM):
    config_class = BitLlamaConfig

    def __init__(self, config: BitLlamaConfig):
        super().__init__(config)
        self.model = BitLlamaModel(config)
        self.lm_head = BitLinear(config.hidden_size, config.vocab_size, bias=False, bits=config.bits, flg_before_linear=True)

2-2. model_type, config_classの設定

model_type = "xxx(自作のモデル名)"をconfigに設定します。私の場合はmodel_typeはbit_llamaとしました。

class BitLlamaConfig(LlamaConfig):
    model_type = "bit_llama"

また、model側にはconfig_classとして自作のconfigを指定します。

class BitLlamaForCausalLM(LlamaForCausalLM):
    config_class = BitLlamaConfig

2-3. 設定ファイル名はmodeling_xxx(model_type).py

これはそのままです。私の場合model_type="bit_llama"なのでmodeling_bit_llama.pyとなりました。
私が作ったファイルを共有します。参考にしていただければ幸いです。

📄modeling_bit_llama.py

3. AutoClassへの登録

上記のファイルをローカルかどこかに置き、モデルの学習に使うはずです。
そしてその後にHuggingFaceへpushすると思います。
そのpushの前のどのタイミングでも良いので以下を実施します。

BitLlamaConfig.register_for_auto_class()
BitLlamaForCausalLM.register_for_auto_class("AutoModelForCausalLM")

自作のconfig, modelであるBitLlamaConfig, BitLlamaForCausalLMが、これでTransformersのAutoClassに登録される様です。例えば他にもBitLlamaをAutoModelに登録したい場合には、

BitLlama.register_for_auto_class("AutoModel")

とします。
これによって、

• AutoConfig時にmodel_type="bit_llama"とされているモデルがあれば自動的にBitLlamaConfigをクラスとして使用してくれます。

• 同じ様に、AutoModelForCausalLM時にmodel_type="bit_llama"であればBitLlamaForCausalLMが使用されます。

最後にこのモデルをHuggingFaceにpushすると、AutoClass（AutoConfigやAutoModelForCausalLM）使用時にこれら新しいクラスを使用するために、これらを含む設定ファイルmodeling_xxx.py(私の場合はmodeling_bit_llama.py)も自動的に一緒にpushされます。
例えば、以下の様なタイミングです。

trainer.push_to_hub()

このmodeling_xxx.pyファイルの配置は割とどこでも良さそうでした。
私はこのbitllamaをmybitnetというライブラリをPYPLに配置し、pipでinstallして使っていました。それでも問題なくpushされました。
参考までにここまでの一連の流れのipynbファイルを共有して置きます。

📄01_bit_llama_pretraining.ipynb

4. モデル取得

これまでの流れによって、最終的には誰でも以下の様に自作のアーキテクチャのモデルを一般のモデルと同じやり方でloadできる様になります。

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

model_name = "HachiML/myBit-Llama2-jp-127M-3"

tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16, device_map="auto", trust_remote_code=True)

ただし、この時modeling_xxx.pyの読み込みを許可するかどうか聞かれるためあらかじめtrust_remote_code=Trueを設定しておくと良いと思います。たまにtrust_remote_codeを設定する必要があるモデルがありましたが、こういうことだったんですね。

また、（これは私の場合のみですが）私のmodeling_bit_llama.pyは独自のライブラリであるmybitnetを使っているので上記実施前にこのinstallが必要です。

pip install mybitnet

最後に

自作のアーキテクチャのモデルを配布するのも割と簡単にできることに驚きました。基本的に配布する側が工夫すれば、取得側は何も気にしなくて良いという設計がとても使いやすいですね。
また、途中途中で貼っていたファイルは下記Githubのものです。ご参照ください。

参照

Building custom models