Inyección de adaptadores

Con PEFT, puedes inyectar adaptadores entrenables en cualquier módulo de torch, lo que te permite utilizar métodos de adaptadores sin depender de las clases de modelado de PEFT. Esto funciona para todos los adaptadores, excepto para aquellos basados en el aprendizaje de prompts (por ejemplo, prefix tuning o p-tuning) y para los adaptadores que mantienen un estado compartido entre múltiples capas objetivo.

Consulte la tabla siguiente para ver cuándo debe inyectar adaptadores.

inject_adapter_in_model no admite los métodos de PEFT que mantienen el estado del adaptador compartido entre varias capas objetivo. Esto incluye actualmente TinyLoRA, UniLoRA, VeRA, PVeRA, VBLoRA y FRoD. Utilice get_peft_model() para estos métodos en su lugar.

Creación de un nuevo modelo PEFT

Para realizar la inyección del adaptador, use el método inject_adapter_in_model(). Este método toma 3 argumentos: la configuración de PEFT, el modelo y un nombre de adaptador opcional. También puede adjuntar múltiples adaptadores al modelo si llama a inject_adapter_in_model() varias veces con diferentes nombres de adaptador.

Por ejemplo, para inyectar adaptadores LoRA en el submódulo linear del módulo DummyModel:

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)

Imprima el modelo para ver que los adaptadores se han inyectado correctamente.

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)
)

Inyección basada en un state_dict

A veces, es posible que exista un checkpoint de adaptador PEFT pero la configuración PEFT correspondiente no sea conocida por cualquier motivo. Para inyectar las capas PEFT de este checkpoint, normalmente habría que ingeniar al revés la configuración PEFT correspondiente, en particular el argumento target_modules, basándose en el state_dict del checkpoint. Esto puede ser tedioso y propenso a errores. Para evitarlo, también es posible llamar a inject_adapter_in_model() y pasar el state_dict cargado como argumento:

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)

En este caso, PEFT usará el state_dict como referencia para qué capas apuntar, en lugar de usar la configuración de PEFT. Como usuario, no tiene que establecer los target_modules exactos de la configuración de PEFT para que esto funcione. Sin embargo, aún debe pasar una configuración de PEFT del tipo correcto, en este ejemplo LoraConfig, y puede dejar los target_modules como None.

Tenga en cuenta que esto solo crea las capas PEFT sin inicializar; los valores de state_dict no se utilizan para poblar los pesos del modelo. Para poblar los pesos, continúe llamando a set_peft_model_state_dict() como se describe a continuación.

⚠️ Ten en cuenta que si hay una discrepancia entre lo configurado en la configuración de PEFT y lo que se encuentra en state_dict, PEFT te emitirá una advertencia al respecto. Puedes ignorar la advertencia si sabes que la configuración de PEFT no está especificada correctamente.

Si los adaptadores PEFT originales estaban usando target_parameters en lugar de target_modules, la inyección desde un state_dict no funcionará correctamente. En este caso, es obligatorio usar la configuración PEFT correcta para la inyección.

Guardado del modelo

Para guardar solo el adaptador, use la función 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)

En caso contrario, model.state_dict() devuelve el estado completo del modelo.

Carga del modelo

Después de cargar el state_dict guardado, se puede aplicar mediante la función 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)

Si la inyección del adaptador es lenta o necesitas cargar un gran número de adaptadores, puedes usar una optimización que permite crear un adaptador “vacío” en el dispositivo meta y solo rellena los pesos con pesos reales cuando se llama a set_peft_model_state_dict(). Para hacer esto, pasa low_cpu_mem_usage=True tanto a inject_adapter_in_model() como a set_peft_model_state_dict().

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

Para cargar pesos desde el hub existe la función de bajo nivel load_peft_weights():

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

Configuración y carga de pesos base

Las funciones anteriores tratan con el state dict del adapter. También hay situaciones en las que es necesario leer o escribir los pesos del base model a través del envoltorio de PEFT. Esto no es del todo trivial, ya que PEFT cambia el nombre de los parámetros de los módulos objetivo (por ejemplo, q_proj.weight pasa a ser q_proj.base_layer.weight) y añade parámetros de adapter que no existen en el modelo base. Utilice get_base_model_state_dict() y set_base_model_state_dict() para traducir entre las dos nomenclaturas:

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)

El caso de uso principal de esto es cargar los pesos base después de que el modelo ya haya sido envuelto por PEFT. Por ejemplo, las configuraciones de entrenamiento con FSDP, como TorchTitan, inicializan el modelo en el dispositivo meta, aplican PEFT y fragmentan el resultado. La memoria real solo existe después de la fragmentación, por lo que el punto de control, cuyas claves usan los nombres originales, debe cargarse en el modelo ya envuelto. Tenga en cuenta que get_base_model_state_dict() devuelve los tensores en vivo del modelo, no copias.