こちらは「LLM・LLM活用 Advent Calendar 2024」の９日目の記事です!
https://qiita.com/advent-calendar/2024/large-language-model

追記: 🎉v4.48.0でリリースされました！

https://github.com/huggingface/transformers/releases/tag/v4.48.0#:~:text=in %2334979-,DiffLlama,-Differential Transformer combines

概要

transformersにPull Requestを作成し、自分のモデルアーキテクチャを追加した。

はじめに

皆さんはあのモデルがtransformersにあればいいのに、もしくは画期的なアーキテクチャを開発して、みんなに使って欲しい、そんなことはありませんか？そういうわけでtransformersに新たなモデルを追加しましたので、記録に残します。

Transformersを読み解く

まずは自分のモデルを追加する前にtransformersを理解しましょう。とりあえずmodelに関わる部分だけを読めばいいので、例としてLlama Modelを見ていきましょうか。

llamaのモデル定義はsrc/transformers/models/llamaにあります。

src/transformers/models/llama/
├── __init__.py
├── convert_llama_weights_to_hf.py
├── configuration_llama.py
├── modeling_flax_llama.py
├── modeling_llama.py
├── tokenization_llama_fast.py
└── tokenization_llama.py

このうち必須なのはconfiguration_llama.pyとmodeling_llama.pyです。そのほかはMeta公式重みからの変換ツールやJax/Flaxのモデルファイルです。（transformersは実はPytorchだけではなくTFやJax/Flaxにも部分対応しています）

configuration_llama.py

このファイルではLlamaConfigを定義しています。そもそもtransformersはPretrainedConfigやPreTrainedModelなどのGeneral Model Classがあります。これらのクラスはfrom_pretrained()などのメソッドがすでに実装してあり、容易に保存、読み込みができます。そのため、これらを継承して実装するのがtransformersの基本です。

https://github.com/huggingface/transformers/blob/329f5dbf97a5cb2473914c88c05aa3dcb242e19a/src/transformers/models/llama/configuration_llama.py#L155-L216

いつもモデルの事前学習している時に設定しているconfigと同じですね。モデルに必要なハイパーパラメータを用意しているだけです。

modeling_llama.py

(間に合わなかったので後日)

モデルの追加

How to add a model to 🤗 Transformers?に従って進めていきます。目標はLlamaモデルのAttentionをDifferential TransformerのAttentionに変更したDiffLlamaです。

環境構築

まずtransformersをForkしてローカルにクローンします。

# Forkした後
git clone https://github.com/[your Github handle]/transformers.git
cd transformers
git remote add upstream https://github.com/huggingface/transformers.git

次にPythonの仮想環境をセットアップします。

python -m venv .env
source .env/bin/activate
#おそらくエラーが出るので下のコマンドで代用
#pip install -e ".[dev]"
pip install -e ".[quality]"

モデルコードの追加

次に、ついに新しいコードを🤗 Transformersに追加できます。🤗 Transformersのフォークのクローンに移動してください

transformers-cli add-new-model-like

What is the model you would like to duplicate? Please provide the lowercase `model_type` (e.g. roberta): llama
What is the name (with no special casing) for your new model in the paper (e.g. RoBERTa)? DiffLlama
What identifier would you like to use for the `model_type` of this model? [diffllama] 
What lowercase name would you like to use for the module (folder) of this model? [diffllama] 
What prefix (camel-cased) would you like to use for the model classes of this model (e.g. Roberta)? [DiffLlama] 
What prefix (upper-cased) would you like to use for the constants relative to this model? [DIFFLLAMA] 
What will be the name of the config class for this model? [DiffLlamaConfig] 
Please give a checkpoint identifier (on the model Hub) for this new model (e.g. facebook/FacebookAI/roberta-base): 
Will your new model use the same processing class as llama (LlamaTokenizerFast) (yes/no)? yes
Should we add # Copied from statements when creating the new modeling file (yes/no)? [yes] 
Should we add a version of your new model in all the frameworks implemented by llama (['pt']) (yes/no)? [yes] 
The tests for symbolic tracing with torch.fx were disabled, you can add those once symbolic tracing works for your new model.

transformers リポジトリでプルリクエストを開く

自動生成されたコードを適応し始める前に、🤗 Transformers に「作業中（WIP）」プルリクエストを開くタイミングです。 例：「[WIP] Add DiffLlama」などです。 これにより、ユーザーと Hugging Face チームが🤗 Transformers にモデルを統合する作業を並行して行うことができます。

以下の手順を実行してください：
1. メインブランチから分かりやすい名前のブランチを作成します。

git checkout -b add_diffllama

2. 自動生成されたコードをコミットしてください:

git add .
git commit

3. 現在の main ブランチにフェッチしてリベース

git fetch upstream
git rebase upstream/main

4. 変更をあなたのアカウントにプッシュするには、次のコマンドを使用します：

git push -u origin add_diffllama

5. 満足したら、GitHub上のフォークのウェブページに移動します。[プルリクエスト]をクリックします。将来の変更に備えて、Hugging Face チームのメンバーのGitHubハンドルをレビュアーとして追加してください。

6. GitHubのプルリクエストウェブページの右側にある「ドラフトに変換」をクリックして、PRをドラフトに変更します。

PR内でHFの方が丁寧に教えてくれるので、困ったらメンションしてみましょう。(どこかにモデルの追加の時は@ArthurZuckerさんにとかリストがあったはずです)

コードゴリゴリタイム

src/transformers/models/diffllama/にモデルファイルがありますが、そちらは編集せずにmodular_diffllama.pyを作成してください。Modular Transformersを使います。

使い方

モデルのクラスがほとんど同じ場合は(configのクラスだけの違いなど)次のように継承&passだけでOKです。

from ..gemma.modeling_gemma import GemmaForCausalLM
from ..llama.modeling_llama import (
 LlamaDecoderLayer,
 LlamaForQuestionAnswering,
 LlamaForSequenceClassification,
 LlamaForTokenClassification,
 LlamaModel,
 LlamaPreTrainedModel,
 LlamaRMSNorm,
 LlamaRotaryEmbedding,
 apply_rotary_pos_emb,
 repeat_kv,
)
from ..mistral.modeling_mistral import MistralMLP
from .configuration_diffllama import DiffLlamaConfig


logger = logging.get_logger(__name__)

_CONFIG_FOR_DOC = "DiffLlamaConfig"


class DiffLlamaRMSNorm(LlamaRMSNorm):
 pass


ALL_LAYERNORM_LAYERS.append(DiffLlamaRMSNorm)


class DiffLlamaRotaryEmbedding(LlamaRotaryEmbedding):
 pass


class DiffLlamaMLP(MistralMLP):
 pass

次のように実装する場合は普通に書けます。

modular_diffllama.pyを完成させた後、python utils/modular_model_converter.py --files_to_parse src/transformers/models/diffllama/modular_diffllama.pyでmodeling_diffllama.py自動生成できます。

モデルのテスト

この時点で、新しいモデルが正常に追加されました。 ただし、モデルがまだ必要な設計に完全に準拠していない可能性が非常に高いです。 🤗 Transformersと完全に互換性があることを確認するために、すべての一般的なテストがパスする必要があります。 Cookiecutterはおそらくモデル用のテストファイルを自動的に追加しているはずで、おそらく同じディレクトリにtests/models/diffllama/test_modeling_diffllama.pyとして存在します。 このテストファイルを実行して、すべての一般的なテストがパスすることを確認してください：
sh pytest tests/models/diffllama/test_modeling_diffllama.py

ドキュメントの追加

よくわからないので11. ドキュメントの追加を参考

コードのリファクタリング

最後にコードを綺麗にします。

make style
make quality

ここまでできたらあとはPR内の指示に従いましょう。

終わりに

あなたはコミュニティの誰でも簡単にアクセスできる別のモデルを作成しました！ 🤯
私はDifferental Transformerの論文に見聞きしてからPRを作成したのが10月11日で、別のこともやりながらようやく数日前に最終レビュー(であって欲しい)に辿り着きました。モデルを一から開発するよりは簡単でしょうが、すごく大変でした。ただ、モデルを追加しているうちに普段使っているtransformersを理解することができました。ぜひ、皆さんも自作のモデルアーキテクチャを追加してみてください。