PEFTチェックポイント形式

このドキュメントでは、PEFTのチェックポイントファイルの構造と、PEFT形式と他の形式との相互変換方法について説明します。

PEFTファイル

PEFT（パラメータ効率の良いファインチューニング）手法では、モデルのすべてのパラメータではなく、その一部のみを更新します。これにより、チェックポイントファイルは元のモデルファイルよりもはるかに小さくなる傾向があり、保存や共有が容易になります。ただし、PEFTモデルを読み込むには、元のモデルも利用可能である必要があります。

PEFTモデルでsave_pretrained()を呼び出すと、PEFTモデルは以下の3つのファイルを保存します。

1. adapter_model.safetensors または adapter_model.bin

デフォルトでは、モデルはsafetensors形式で保存されます。これはbin形式の安全な代替手段であり、bin形式は内部でpickleユーティリティを使用しているため、セキュリティ上の脆弱性に脆弱であることが知られています。ただし、両方の形式は同じstate_dictを保存しており、互換性があります。

state_dictには、ベースモデルのパラメータではなく、アダプターモジュールのパラメータのみが含まれます。サイズの差を説明するために、通常のBERTモデルには約420MBのディスク容量が必要ですが、このBERTモデル上に構築されたIA³アダプターには約260KBしか必要ありません。

2. adapter_config.json

adapter_config.json ファイルには、モデルの読み込みに必要なアダプターモジュールの設定が含まれています。以下は、BERTモデルに標準設定を適用したIA³アダプターのadapter_config.jsonの例です。

{
  "auto_mapping": {
    "base_model_class": "BertModel",
    "parent_library": "transformers.models.bert.modeling_bert"
  },
  "base_model_name_or_path": "bert-base-uncased",
  "fan_in_fan_out": false,
  "feedforward_modules": [
    "output.dense"
  ],
  "inference_mode": true,
  "init_ia3_weights": true,
  "modules_to_save": null,
  "peft_type": "IA3",
  "revision": null,
  "target_modules": [
    "key",
    "value",
    "output.dense"
  ],
  "task_type": null
}

設定ファイルには以下が含まれています:

• 保存されたアダプターモジュールのタイプ、"peft_type": "IA3"
• "base_model_name_or_path": "bert-base-uncased" のようなベースモデルに関する情報
• モデルのリビジョン（ある場合）、"revision": null

ベースモデルが事前学習済みのTransformersモデルでない場合、後者の2つのエントリはnullになります。それ以外については、設定はすべて、モデルのファインチューニングに使用された特定のIA³アダプターに関連しています。

3. README.md

生成されたREADME.mdはPEFTモデルのモデルカードであり、いくつかの項目が事前に入力されています。これは、モデルを他の人と共有しやすくし、モデルに関する基本的な情報を提供することを目的としています。このファイルはモデルの読み込みに必要ありません。

PEFT形式に変換

他の形式からPEFT形式へ変換する場合は、adapter_model.safetensors（またはadapter_model.bin）ファイルとadapter_config.jsonファイルの両方が必要です。

adapter_model

モデルの重みについては、PEFTがファイルをロードするために、パラメータ名から値への正しいマッピングを使用することが重要です。PEFTアダプターには一般的に合意されたフォーマットが存在しないため、このマッピングを正しく設定することは、実装の詳細を確認する作業となります。

幸いなことに、一般的な基本ケースでは、このマッピングを特定することはそれほど複雑ではありません。LoraLayerを具体的な例として見てみましょう。

# showing only part of the code

class LoraLayer(BaseTunerLayer):
    # All names of layers that may contain (trainable) adapter weights
    adapter_layer_names = ("lora_A", "lora_B", "lora_embedding_A", "lora_embedding_B")
    # All names of other parameters that may contain adapter-related parameters
    other_param_names = ("r", "lora_alpha", "scaling", "lora_dropout")

    def __init__(self, base_layer: nn.Module, **kwargs) -> None:
        self.base_layer = base_layer
        self.r = {}
        self.lora_alpha = {}
        self.scaling = {}
        self.lora_dropout = nn.ModuleDict({})
        self.lora_A = nn.ModuleDict({})
        self.lora_B = nn.ModuleDict({})
        # For Embedding layer
        self.lora_embedding_A = nn.ParameterDict({})
        self.lora_embedding_B = nn.ParameterDict({})
        # Mark the weight as unmerged
        self._disable_adapters = False
        self.merged_adapters = []
        self.use_dora: dict[str, bool] = {}
        self.lora_magnitude_vector: Optional[torch.nn.ParameterDict] = None  # for DoRA
        self._caches: dict[str, Any] = {}
        self.kwargs = kwargs

PEFT のすべての LoraLayer クラスで使用される __init__ コードには、モデルを初期化するために使用される多くのパラメータがありますが、チェックポイントファイルに関連するのは lora_A、lora_B、lora_embedding_A、および lora_embedding_B の数個だけです。これらのパラメータはクラス属性 adapter_layer_names にリストされており、学習可能なパラメータを含むため、チェックポイントファイルに含める必要があります。ランク r などのその他のすべてのパラメータは、adapter_config.json から導出されるため、そこに含める必要があります（デフォルト値が使用されている場合は除く）。

BERTに適用されたPEFT LoRAモデルのstate_dictを確認しましょう。デフォルトのLoRA設定で最初の5つのキーを表示すると（残りのキーは同じで、レイヤー番号だけが異なります）、次のようになります。

• base_model.model.encoder.layer.0.attention.self.query.lora_A.weight
• base_model.model.encoder.layer.0.attention.self.query.lora_B.weight
• base_model.model.encoder.layer.0.attention.self.value.lora_A.weight
• base_model.model.encoder.layer.0.attention.self.value.lora_B.weight
• base_model.model.encoder.layer.1.attention.self.query.lora_A.weight
• など

これを分解して見てみましょう：

• デフォルトでは、BERTモデルに対してLoRAは、attentionモジュールのquery層とvalue層に適用されます。そのため、各層のキー名にattention.self.queryとattention.self.valueが表示されます。
• LoRAは重みを2つの低ランク行列lora_Aとlora_Bに分解します。キー名におけるlora_Aとlora_Bは、この分解に由来しています。
• これらの LoRA 行列は nn.Linear レイヤーとして実装されているため、パラメータは .weight 属性（lora_A.weight、lora_B.weight）に保存されます。
• デフォルトでは、LoRAはBERTの埋め込み層には適用されないため、lora_A_embeddingおよびlora_B_embeddingのエントリはありません。
• state_dict のキーは常に "base_model.model." で始まります。その理由は、PEFT ではベースモデルをチューナー固有のモデル（この場合は LoraModel）でラップし、さらにそのモデルを汎用的な PEFT モデル（PeftModel）でラップするためです。そのため、これらの 2 つのプレフィックスがキーに追加されます。PEFT 形式に変換する際には、これらのプレフィックスを追加する必要があります。

この最後の点は、prompt tuningのようなprefix tuning手法には当てはまりません。これらの手法では、追加のembeddingがキーにプレフィックスを追加することなく、state_dictに直接保存されます。

読み込んだモデルのパラメータ名を確認すると、base_model.model.encoder.layer.0.attention.self.query.lora_A.default.weight のように、少し異なる見た目になっていることに驚くかもしれません。違いは、最後のセグメントの一つ前の部分にある .default です。この部分が存在するのは、PEFT が一般的に複数のアダプタを同時に追加することを許可しているためです（nn.ModuleDict または nn.ParameterDict を使用してそれらを保存します）。例えば、“other” という名前の別のアダプタを追加した場合、そのアダプタのキーは base_model.model.encoder.layer.0.attention.self.query.lora_A.other.weight になります。

save_pretrained()を呼び出すと、キーからアダプター名が除去されます。その理由は、アダプター名はモデルアーキテクチャにとって重要な部分ではなく、単に任意の名前だからです。アダプターを読み込む際に全く異なる名前を選択しても、モデルは同じように動作します。そのため、アダプター名はチェックポイントファイルに保存されません。

save_pretrained("some/path") を呼び出し、アダプター名が "default" でない場合、アダプターはアダプター名と同じ名前のサブディレクトリに保存されます。つまり、名が “other” の場合、some/path/other の中に保存されます。

状況によっては、チェックポイントファイルにどの値を追加するかを判断するのがやや複雑になる場合があります。例えば、PEFTでは、DoRAはLoRAの特殊ケースとして実装されています。DoRAモデルをPEFTに変換したい場合は、DoRA用の追加エントリを含むLoRAチェックポイントを作成する必要があります。これは、前述のLoraLayerコードの__init__で確認できます。

self.lora_magnitude_vector: Optional[torch.nn.ParameterDict] = None  # for DoRA

これは、DoRAではレイヤーごとに任意の追加パラメータが存在することを示しています。

adapter_config

PEFTモデルの読み込みに必要なその他のすべての情報は、adapter_config.jsonファイルに含まれています。BERTに適用されたLoRAモデルについて、このファイルを確認してみましょう。

{
  "alpha_pattern": {},
  "auto_mapping": {
    "base_model_class": "BertModel",
    "parent_library": "transformers.models.bert.modeling_bert"
  },
  "base_model_name_or_path": "bert-base-uncased",
  "bias": "none",
  "fan_in_fan_out": false,
  "inference_mode": true,
  "init_lora_weights": true,
  "layer_replication": null,
  "layers_pattern": null,
  "layers_to_transform": null,
  "loftq_config": {},
  "lora_alpha": 8,
  "lora_dropout": 0.0,
  "megatron_config": null,
  "megatron_core": "megatron.core",
  "modules_to_save": null,
  "peft_type": "LORA",
  "r": 8,
  "rank_pattern": {},
  "revision": null,
  "target_modules": [
    "query",
    "value"
  ],
  "task_type": null,
  "use_dora": false,
  "use_rslora": false
}

このファイルには多くのエントリが含まれており、一見すると、そこに設定すべき正しい値をすべて見つけるのが圧倒的に感じられるかもしれません。しかし、モデルの読み込みに必要なエントリはごく一部です。これは、デフォルト値が使用されるため追加する必要がない場合や、LoRA重みの初期化にのみ影響を与えるため、モデルの読み込みには無関係な場合があるためです。特定のパラメータの役割がわからない場合、例えば"use_rslora",などは、追加しないでください。そうすれば問題ありません。また、オプションが追加されるにつれて、このファイルには将来的により多くのエントリが追加されますが、後方互換性は保たれるはずです。

最低限、以下のエントリを含める必要があります。

{
  "target_modules": ["query", "value"],
  "peft_type": "LORA"
}

ただし、r（ランク）や、Transformers モデルの場合は base_model_name_or_path など、可能な限り多くのエントリを追加することをお勧めします。この情報により、他の人がモデルをよりよく理解し、より簡単に共有できるようになります。期待されるキーと値を確認するには、PEFT のソースコード内の config.py ファイル（LoRA の設定ファイルの例）を参照してください。

モデルの保存

状況によっては、ベースの重みを含むPEFTモデル全体を保存したい場合があります。例えば、PEFTモデルを読み込もうとするユーザーがベースモデルにアクセスできない場合など、これが必須になることがあります。まず重みをマージするか、Transformerモデルに変換することができます。

重みをマージする

PEFTモデル全体を保存する最もシンプルな方法は、アダプター重みをベース重みにマージすることです：

merged_model = model.merge_and_unload()
merged_model.save_pretrained(...)

ただし、このアプローチにはいくつかのデメリットがあります：

• merge_and_unload() が呼び出されると、PEFT 固有の機能を持たない基本的なモデルが得られます。つまり、PEFT 固有のメソッドはこれ以上使用できません。
• 重みのマージ解除、複数のアダプタの同時読み込み、アダプタの無効化などはできません。
• すべてのPEFTメソッドが重みのマージをサポートしているわけではありません。
• 一部のPEFT手法は一般的にマージを許可しますが、特定の設定では許可されない場合があります（例：特定の量子化技術を使用している場合）。
• PEFTモデルにはベースの重みもすべて含まれるため、モデル全体はPEFTモデルよりもはるかに大きくなります。

ただし、マージされたモデルでの推論は、わずかに高速になるはずです。

Transformersモデルに変換する

ベースモデルがTransformersモデルであるという前提のもとで、モデル全体を保存する別の方法として、PEFTの重みをベースモデルに直接挿入して保存する、このハック的なアプローチがあります。ただし、これはTransformersにPEFTモデルがPEFTモデルではないと「誤認させる」場合のみ機能します。この方法はLoRAでのみ機能し、他のアダプターはTransformersに実装されていないため使用できません。

model = ...  # the PEFT model
...
# after you finish training the model, save it in a temporary location
model.save_pretrained(<temp_location>)
# now load this model directly into a transformers model, without the PEFT wrapper
# the PEFT weights are directly injected into the base model
model_loaded = AutoModel.from_pretrained(<temp_location>)
# now make the loaded model believe that it is _not_ a PEFT model
model_loaded._hf_peft_config_loaded = False
# now when we save it, it will save the whole model
model_loaded.save_pretrained(<final_location>)
# or upload to Hugging Face Hub
model_loaded.push_to_hub(<final_location>)