Формат контрольной точки PEFT

В этом документе описано, как устроены файлы контрольных точек PEFT, и как выполнять преобразование между форматом PEFT и другими форматами.

Файлы PEFT

Методы PEFT (parameter-efficient fine-tuning) обновляют лишь небольшую часть параметров модели, а не все. Это удобно, поскольку файлы контрольных точек обычно значительно меньше исходных файлов модели, что упрощает их хранение и обмен. Однако это также означает, что для загрузки модели PEFT необходимо иметь доступ к исходной модели.

При вызове save_pretrained() для модели PEFT она сохраняет три файла, описанные ниже:

1. adapter_model.safetensors или adapter_model.bin

По умолчанию модель сохраняется в формате safetensors — безопасной альтернативе формату bin, который известен своей подверженностью уязвимостям безопасности, поскольку использует утилиту pickle. При этом оба формата хранят один и тот же state_dict и взаимозаменяемы.

state_dict содержит только параметры модуля адаптера, а не базовой модели. Для наглядности: обычная модель BERT занимает ~420 МБ на диске, тогда как адаптер IA³ поверх этой модели BERT требует лишь ~260 КБ.

2. adapter_config.json

Файл adapter_config.json содержит конфигурацию модуля адаптера, необходимую для загрузки модели. Ниже приведен пример adapter_config.json для IA³-адаптера со стандартными настройками, примененного к модели BERT:

{
  "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, две последние записи будут 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

В коде __init__, используемом всеми классами LoraLayer в PEFT, есть целый набор параметров для инициализации модели, но для файла чекпоинта релевантны лишь некоторые из них: lora_A, lora_B, lora_embedding_A и lora_embedding_B. Эти параметры перечислены в атрибуте класса adapter_layer_names и содержат обучаемые параметры, поэтому их обязательно нужно включать в файл чекпоинта. Все остальные параметры, например ранг r, выводятся из adapter_config.json и должны быть указаны там (если не используется значение по умолчанию).

Проверим state_dict модели PEFT LoRA, применённой к BERT. При выводе первых пяти ключей с настройками LoRA по умолчанию (остальные ключи аналогичны, но с другими номерами слоёв) получаем:

• 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 применяется к слоям query и value модуля внимания. Именно поэтому в именах ключей для каждого слоя вы видите attention.self.query и attention.self.value.
• LoRA разлагает веса на две матрицы низкого ранга, 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). По этой причине к ключам добавляются эти два префикса. При конвертации в формат PEFT необходимо добавить эти префиксы.

Это последнее утверждение не относится к методам prefix tuning, таким как prompt tuning. В них дополнительные эмбеддинги хранятся непосредственно в 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, вам следует создать контрольную точку LoRA с дополнительными записями для DoRA. Вы можете увидеть это в __init__ предыдущего кода LoraLayer:

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

Это указывает на то, что для DoRA существует необязательный дополнительный параметр на каждый слой.

adapter_config

Вся остальная информация, необходимая для загрузки модели PEFT, содержится в файле adapter_config.json. Проверим этот файл для модели LoRA, применённой к BERT:

{
  "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 или base_model_name_or_path (если это модель Transformers). Эта информация помогает другим лучше понять модель и легче ею поделиться. Чтобы узнать, какие ключи и значения ожидаются, посмотрите файл config.py (например, это файл конфигурации для LoRA) в исходном коде PEFT.

Хранение моделей

В некоторых случаях может потребоваться сохранить всю модель PEFT, включая базовые веса. Это необходимо, например, если базовая модель недоступна пользователям, которые пытаются загрузить модель PEFT. Сначала можно объединить веса или преобразовать модель в модель Transformers.

Слияние весов

Самый простой способ сохранить всю модель PEFT — это объединить веса адаптера с базовыми весами:

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

Однако у этого подхода есть некоторые недостатки:

• После вызова merge_and_unload() вы получаете базовую модель без какой-либо функциональности, специфичной для 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>)