アダプター注入

PEFTを使用すると、torchモジュールに学習可能なアダプターを注入できるため、PEFTのモデリングクラスに依存せずにアダプター手法を利用できます。これは、プロンプト学習に基づくアダプター（プレフィックスチューニングやp-tuningなど）および、複数のターゲットレイヤー間で状態を共有するアダプターを除く、すべてのアダプターで機能します。

アダプタを注入すべきタイミングを確認するには、以下の表を参照してください。

利点	欠点
モデルはインプレースで変更され、元のすべての属性とメソッドが保持されます	`from_pretrained` と `save_pretrained` ユーティリティ関数を Hugging Face から手動で書き、アダプターを保存および読み込む
`torch`のモジュールやモダリティすべてで動作します	`PeftModel`が提供するアダプターの無効化やマージなどのユーティリティメソッドとは互換性がありません

inject_adapter_in_model は、複数のターゲット層間でアダプタの状態を共有する PEFT メソッドをサポートしていません。これには現在、TinyLoRA、UniLoRA、VeRA、PVeRA、VBLoRA、FRoD が含まれます。これらのメソッドには代わりに get_peft_model() を使用してください。

新しいPEFTモデルの作成

アダプタの注入を実行するには、inject_adapter_in_model() メソッドを使用します。このメソッドは3つの引数を受け取ります。PEFT設定、モデル、および任意のアダプタ名です。異なるアダプタ名を指定して inject_adapter_in_model() を複数回呼び出すことで、モデルに複数のアダプタをアタッチすることもできます。

例えば、DummyModelモジュールのlinearサブモジュールにLoRAアダプターを注入するには:

import torch
from peft import inject_adapter_in_model, LoraConfig

class DummyModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.embedding = torch.nn.Embedding(10, 10)
        self.linear = torch.nn.Linear(10, 10)
        self.lm_head = torch.nn.Linear(10, 10)

    def forward(self, input_ids):
        x = self.embedding(input_ids)
        x = self.linear(x)
        x = self.lm_head(x)
        return x


lora_config = LoraConfig(
    lora_alpha=16,
    lora_dropout=0.1,
    r=64,
    bias="none",
    target_modules=["linear"],
)

model = DummyModel()
model = inject_adapter_in_model(lora_config, model)

dummy_inputs = torch.LongTensor([[0, 1, 2, 3, 4, 5, 6, 7]])
dummy_outputs = model(dummy_inputs)

アダプターが正しく注入されたことを確認するために、モデルを出力します。

DummyModel(
  (embedding): Embedding(10, 10)
  (linear): Linear(
    in_features=10, out_features=10, bias=True
    (lora_dropout): ModuleDict(
      (default): Dropout(p=0.1, inplace=False)
    )
    (lora_A): ModuleDict(
      (default): Linear(in_features=10, out_features=64, bias=False)
    )
    (lora_B): ModuleDict(
      (default): Linear(in_features=64, out_features=10, bias=False)
    )
    (lora_embedding_A): ParameterDict()
    (lora_embedding_B): ParameterDict()
  )
  (lm_head): Linear(in_features=10, out_features=10, bias=True)
)

state_dictに基づくインジェクション

PEFT アダプターのチェックポイントが存在するものの、何らかの理由で対応する PEFT 設定が不明な場合があります。このチェックポイントに PEFT レイヤーを注入するには、通常、チェックポイントの state_dict をもとに、特に target_modules 引数を含む対応する PEFT 設定を逆算する必要があります。これは手間がかかり、エラーが発生しやすい作業です。これを避けるために、inject_adapter_in_model() を呼び出し、読み込んだ state_dict を引数として渡すことも可能です。

from safetensors.torch import load_file

model = ...
state_dict = load_file(<path-to-safetensors-file>)
lora_config = LoraConfig(...)
model = inject_adapter_in_model(lora_config, model, state_dict=state_dict)

この場合、PEFT は PEFT 設定ではなく、state_dict を参照して対象とするレイヤーを決定します。ユーザーとして、この動作のために PEFT 設定の正確な target_modules を設定する必要はありません。ただし、この例では LoraConfig のように、適切なタイプの PEFT 設定を渡す必要があります。target_modules は None のままで構いません。

この処理では、未初期化のPEFTレイヤーのみが作成されることに注意してください。state_dictの値は、モデルの重みを設定するために使用されません。重みを設定するには、以下で説明するset_peft_model_state_dict()を呼び出すことで続行してください。

⚠️ PEFT の設定とstate_dictの内容に不一致がある場合、PEFT は警告を表示します。PEFT の設定が正しく指定されていないことがわかっている場合は、この警告を無視できます。

元の PEFT アダプターがtarget_modulesではなくtarget_parametersを使用していた場合、state_dictからの注入は正しく動作しません。この場合、注入には正しい PEFT 設定を使用する必要があります。

モデルの保存

アダプターのみを保存するには、get_peft_model_state_dict()関数を使用します:

from peft import get_peft_model_state_dict

peft_state_dict = get_peft_model_state_dict(model)
print(peft_state_dict)

それ以外の場合、model.state_dict()はモデルの完全なstate dictを返します。

モデルの読み込み

保存されたstate_dictを読み込んだ後、set_peft_model_state_dict()関数を使用して適用できます。

from peft import set_peft_model_state_dict

model = DummyModel()
model = inject_adapter_in_model(lora_config, model)
outcome = set_peft_model_state_dict(model, peft_state_dict)
# check that there were no wrong keys
print(outcome.unexpected_keys)

アダプターの注入が遅い場合や、多数のアダプターを読み込む必要がある場合は、メタデバイス上で「空」のアダプターを作成し、set_peft_model_state_dict() が呼び出された時点で実際の重みで重みを埋めることを可能にする最適化を使用できます。これを行うには、inject_adapter_in_model() と set_peft_model_state_dict() の両方に low_cpu_mem_usage=True を渡します。

model = DummyModel()
model = inject_adapter_in_model(lora_config, model, low_cpu_mem_usage=True)

print(model.linear.lora_A["default"].weight.device.type == "meta")  # should be True
set_peft_model_state_dict(model, peft_state_dict, low_cpu_mem_usage=True)
print(model.linear.lora_A["default"].weight.device.type == "cpu")  # should be True

ハブから重みを読み込むには、低レベルのload_peft_weights()関数があります：

state_dict = load_peft_weights("my-account/my-adapter-repo")

ベース重みの設定と読み込み

上記の関数はadapterのstate dictを扱います。また、PEFTラッパーを介してbase modelの重みを読み書きする必要がある場面もあります。これは完全に自明ではありません。PEFTは対象モジュールのパラメータ名を変更する（例: q_proj.weightがq_proj.base_layer.weightになる）ためです。さらに、base modelには存在しないadapterのパラメータも追加されます。2つの命名規則の間で変換するには、get_base_model_state_dict()とset_base_model_state_dict()を使用します。

from peft import LoraConfig, get_peft_model, get_base_model_state_dict, set_base_model_state_dict
from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("facebook/opt-125m")
model = get_peft_model(model, LoraConfig(target_modules="all-linear"))

# state dict with the original (pre-PEFT) key names, adapter parameters excluded
base_state_dict = get_base_model_state_dict(model)

# load a state dict with original key names into the PEFT-wrapped model
outcome = set_base_model_state_dict(model, base_state_dict, strict=False)
# check that there were no wrong keys
print(outcome.missing_keys, outcome.unexpected_keys)

この機能の主なユースケースは、モデルがすでに PEFT でラップされた後にベースの重みをロードすることです。例えば、TorchTitan などの FSDP トレーニングセットアップでは、モデルをメタデバイス上で初期化し、PEFT を適用し、その結果をシャーディングします。実際のメモリはシャーディング後にのみ存在するため、キーが元の名称を使用しているチェックポイントは、すでにラップされたモデルにロードする必要があります。なお、get_base_model_state_dict() はモデルのライブなテンソルを返し、コピーを返すわけではありません。

Update on GitHub