Formato de punto de control de PEFT

Este documento describe cómo están estructurados los archivos de punto de control de PEFT y cómo convertir entre el formato de PEFT y otros formatos.

Archivos PEFT

Los métodos de PEFT (fine-tuning eficiente en parámetros) solo actualizan un subconjunto pequeño de los parámetros de un modelo, en lugar de todos. Esto es conveniente porque los archivos de checkpoint suelen ser mucho más pequeños que los archivos del modelo original y son más fáciles de almacenar y compartir. Sin embargo, esto también significa que, para cargar un modelo PEFT, es necesario tener disponible el modelo original.

Cuando se llama a save_pretrained() en un modelo PEFT, el modelo PEFT guarda tres archivos, descritos a continuación:

1. adapter_model.safetensors o adapter_model.bin

Por defecto, el modelo se guarda en el formato safetensors, una alternativa segura al formato bin, que se sabe que es susceptible a vulnerabilidades de seguridad porque utiliza la utilidad pickle internamente. Sin embargo, ambos formatos almacenan el mismo state_dict y son intercambiables.

El state_dict solo contiene los parámetros del módulo de adaptador, no los del modelo base. Para ilustrar la diferencia de tamaño, un modelo BERT normal requiere ~420MB de espacio en disco, mientras que un adaptador IA³ sobre este modelo BERT solo requiere ~260KB.

2. adapter_config.json

El archivo adapter_config.json contiene la configuración del módulo de adaptador, que es necesaria para cargar el modelo. A continuación se muestra un ejemplo de un adapter_config.json para un adaptador IA³ con la configuración estándar aplicada a un modelo 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
}

El archivo de configuración contiene:

• el tipo de módulo de adaptador almacenado, "peft_type": "IA3"
• información sobre el modelo base, como "base_model_name_or_path": "bert-base-uncased"
• la revisión del modelo (si la hay), "revision": null

Si el modelo base no es un modelo de Transformers preentrenado, las dos últimas entradas serán null. Aparte de eso, la configuración está relacionada con el adaptador IA³ específico que se utilizó para el ajuste fino del modelo.

3. README.md

El README.md generado es la tarjeta del modelo de un modelo PEFT y contiene algunas entradas predefinidas. El objetivo es facilitar el intercambio del modelo con otras personas y proporcionar información básica sobre el modelo. Este archivo no es necesario para cargar el modelo.

Convertir al formato PEFT

Al convertir desde otro formato al formato PEFT, se requieren tanto el archivo adapter_model.safetensors (o adapter_model.bin) como el archivo adapter_config.json.

adapter_model

Para los pesos del modelo, es importante utilizar el mapeo correcto de nombre de parámetro a valor para que PEFT pueda cargar el archivo. Obtener este mapeo correcto es un ejercicio de verificación de los detalles de la implementación, ya que no existe un formato generalmente aceptado para los adaptadores de PEFT.

Afortunadamente, determinar este mapeo no es demasiado complicado para los casos base comunes. Veamos un ejemplo concreto, el 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

En el código de __init__ utilizado por todas las clases LoraLayer en PEFT, hay un conjunto de parámetros para inicializar el modelo, pero solo unos pocos son relevantes para el archivo de checkpoint: lora_A, lora_B, lora_embedding_A y lora_embedding_B. Estos parámetros se enumeran en el atributo de clase adapter_layer_names y contienen los parámetros aprendibles, por lo que deben incluirse en el archivo de checkpoint. Todos los demás parámetros, como el rango r, se derivan de adapter_config.json y deben incluirse allí (a menos que se utilice el valor predeterminado).

Veamos el state_dict de un modelo PEFT LoRA aplicado a BERT. Al imprimir las primeras cinco claves usando la configuración predeterminada de LoRA (las claves restantes son las mismas, solo con números de capa diferentes), obtenemos:

• 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
• etc.

Desglosemos esto:

• Por defecto, en los modelos BERT, LoRA se aplica a las capas query y value del módulo de atención. Por eso verás attention.self.query y attention.self.value en los nombres de clave de cada capa.
• LoRA descompone los pesos en dos matrices de bajo rango, lora_A y lora_B. De ahí provienen lora_A y lora_B en los nombres de las claves.
• Estas matrices LoRA se implementan como capas nn.Linear, por lo que los parámetros se almacenan en el atributo .weight (lora_A.weight, lora_B.weight).
• Por defecto, LoRA no se aplica a la capa de incrustación de BERT, por lo que no hay entradas para lora_A_embedding y lora_B_embedding.
• Las claves del state_dict siempre comienzan con "base_model.model.". La razón es que, en PEFT, envolvemos el modelo base dentro de un modelo específico del afinador (LoraModel en este caso), que a su vez se envuelve en un modelo PEFT general (PeftModel). Por esta razón, se añaden estos dos prefijos a las claves. Al convertir al formato PEFT, es necesario añadir estos prefijos.

Este último punto no es cierto para las técnicas de ajuste de prefijo como el ajuste de prompt. Allí, los embeddings adicionales se almacenan directamente en el state_dict sin añadir ningún prefijo a las claves.

Al inspeccionar los nombres de los parámetros en el modelo cargado, es posible que te sorprenda descubrir que se ven un poco diferentes, por ejemplo base_model.model.encoder.layer.0.attention.self.query.lora_A.default.weight. La diferencia es la parte .default en el penúltimo segmento. Esta parte existe porque PEFT generalmente permite la adición de múltiples adaptadores a la vez (usando un nn.ModuleDict o un nn.ParameterDict para almacenarlos). Por ejemplo, si agregas otro adaptador llamado “other”, la clave para ese adaptador sería base_model.model.encoder.layer.0.attention.self.query.lora_A.other.weight.

Cuando se llama a save_pretrained(), el nombre del adaptador se elimina de las claves. La razón es que el nombre del adaptador no es una parte importante de la arquitectura del modelo; es solo un nombre arbitrario. Al cargar el adaptador, se podría elegir un nombre totalmente diferente y el modelo seguiría funcionando de la misma manera. Por eso, el nombre del adaptador no se guarda en el archivo de checkpoint.

Si llamas a save_pretrained("some/path") y el nombre del adaptador no es "default", el adaptador se guarda en un subdirectorio con el mismo nombre que el adaptador. Así, si el nombre es “other”, se guardará dentro de some/path/other.

En algunas circunstancias, decidir qué valores añadir al archivo de checkpoint puede volverse un poco más complicado. Por ejemplo, en PEFT, DoRA se implementa como un caso especial de LoRA. Si deseas convertir un modelo DoRA a PEFT, debes crear un checkpoint de LoRA con entradas adicionales para DoRA. Puedes verlo en el __init__ del código anterior de LoraLayer:

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

Esto indica que hay un parámetro extra opcional por capa para DoRA.

adapter_config

Toda la demás información necesaria para cargar un modelo PEFT está contenida en el archivo adapter_config.json. Veamos este archivo para un modelo LoRA aplicado a 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
}

Esta configuración contiene muchas entradas y, a primera vista, puede resultar abrumador determinar todos los valores correctos que deben incluirse. Sin embargo, la mayoría de las entradas no son necesarias para cargar el modelo. Esto se debe a que utilizan valores predeterminados y no necesitan ser añadidos, o a que solo afectan la inicialización de los pesos de LoRA, lo cual es irrelevante a la hora de cargar el modelo. Si no sabes qué hace un parámetro específico, por ejemplo, "use_rslora",, no lo añadas y deberías estar bien. Ten en cuenta también que, a medida que se añadan más opciones, este archivo tendrá más entradas en el futuro, pero debería ser compatible hacia atrás.

Como mínimo, debe incluir las siguientes entradas:

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

Sin embargo, se recomienda añadir tantas entradas como sea posible, como el rango r o el base_model_name_or_path (si es un modelo de Transformers). Esta información puede ayudar a otros a comprender mejor el modelo y a compartirlo con mayor facilidad. Para comprobar qué claves y valores se esperan, consulte el archivo config.py (por ejemplo, este es el archivo de configuración para LoRA) en el código fuente de PEFT.

Almacenamiento de modelos

En algunas circunstancias, es posible que desee almacenar el modelo PEFT completo, incluidos los pesos base. Esto puede ser necesario si, por ejemplo, el modelo base no está disponible para los usuarios que intentan cargar el modelo PEFT. Puede fusionar los pesos primero o convertirlo en un modelo Transformer.

Fusionar los pesos

La forma más directa de almacenar el modelo PEFT completo es fusionar los pesos del adaptador en los pesos base:

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

Sin embargo, este enfoque tiene algunas desventajas:

• Una vez que se llama a merge_and_unload(), se obtiene un modelo básico sin ninguna funcionalidad específica de PEFT. Esto significa que ya no se pueden utilizar los métodos específicos de PEFT.
• No se pueden desfusionar los pesos, cargar varios adaptadores a la vez, deshabilitar el adaptador, etc.
• No todos los métodos de PEFT admiten la fusión de pesos.
• Algunos métodos de PEFT pueden permitir la fusión de forma general, pero no con configuraciones específicas (por ejemplo, al usar ciertas técnicas de cuantización).
• El modelo completo será mucho más grande que el modelo PEFT, ya que también contendrá todos los pesos base.

Pero la inferencia con un modelo fusionado debería ser un poco más rápida.

Convertir a un modelo de Transformers

Otra forma de guardar el modelo completo, asumiendo que el modelo base es un modelo de Transformers, es usar este enfoque poco ortodoxo para insertar directamente los pesos de PEFT en el modelo base y guardarlo, lo cual solo funciona si “engañas” a Transformers para que crea que el modelo PEFT no es un modelo PEFT. Esto solo funciona con LoRA porque otros adaptadores no están implementados en 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>)