Injeção de adaptadores

Com o PEFT, você pode injetar adaptadores treináveis em qualquer módulo torch, o que permite usar métodos de adaptador sem depender das classes de modelagem no PEFT. Isso funciona para todos os adaptadores, exceto para aqueles baseados em aprendizado de prompt (por exemplo, prefix tuning ou p-tuning) e adaptadores que mantêm estado compartilhado entre várias camadas de destino.

Consulte a tabela abaixo para ver quando você deve injetar adaptadores.

inject_adapter_in_model não suporta métodos PEFT que mantêm o estado do adaptador compartilhado entre várias camadas de destino. Isso inclui atualmente TinyLoRA, UniLoRA, VeRA, PVeRA, VBLoRA e FRoD. Use get_peft_model() para esses métodos, em vez disso.

Criando um novo modelo PEFT

Para realizar a injeção do adaptador, use o método inject_adapter_in_model(). Este método aceita 3 argumentos: a configuração PEFT, o modelo e um nome de adaptador opcional. Você também pode anexar múltiplos adaptadores ao modelo se chamar inject_adapter_in_model() várias vezes com diferentes nomes de adaptador.

Por exemplo, para injetar adaptadores LoRA no submódulo linear do 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 o modelo para verificar se os adaptadores foram injetados corretamente.

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

Injeção com base em um state_dict

Às vezes, pode acontecer de haver um checkpoint de adaptador PEFT, mas a configuração PEFT correspondente não ser conhecida por qualquer motivo. Para injetar as camadas PEFT para esse checkpoint, normalmente seria necessário recriar a configuração PEFT correspondente, em especial o argumento target_modules, com base no state_dict do checkpoint. Isso pode ser trabalhoso e propenso a erros. Para evitar isso, também é possível chamar inject_adapter_in_model() e passar o state_dict carregado 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)

Nesse caso, o PEFT usará o state_dict como referência para quais camadas direcionar, em vez de usar a configuração do PEFT. Como usuário, você não precisa definir os target_modules exatos da configuração do PEFT para que isso funcione. No entanto, você ainda deve passar uma configuração do PEFT do tipo correto, neste exemplo LoraConfig, e pode deixar o target_modules como None.

Observe que isso ainda apenas cria as camadas PEFT não inicializadas; os valores de state_dict não são usados para preencher os pesos do modelo. Para preencher os pesos, prossiga chamando set_peft_model_state_dict(), conforme descrito abaixo.

⚠️ Observe que, se houver uma incompatibilidade entre o que está configurado no config do PEFT e o que é encontrado no state_dict, o PEFT emitirá um aviso sobre isso. Você pode ignorar o aviso se souber que o config do PEFT não está especificado corretamente.

Se os adaptadores PEFT originais estavam usando target_parameters em vez de target_modules, a injeção a partir de um state_dict não funcionará corretamente. Nesse caso, é obrigatório usar a configuração PEFT correta para a injeção.

Salvando o modelo

Para salvar apenas o adaptador, use a função 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)

Caso contrário, model.state_dict() retorna o state dict completo do modelo.

Carregando o modelo

Após carregar o state_dict salvo, ele pode ser aplicado usando a função 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)

Se a injeção do adaptador for lenta ou você precisar carregar um grande número de adaptadores, pode usar uma otimização que permite criar um adaptador “vazio” no meta device e preencher os pesos apenas com pesos reais quando o set_peft_model_state_dict() é chamado. Para isso, passe low_cpu_mem_usage=True tanto para inject_adapter_in_model() quanto para 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 carregar pesos do hub, existe a função de baixo nível load_peft_weights():

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

Definindo e carregando os pesos base

As funções acima lidam com o state dict do adapter. Existem também situações em que os pesos do base model precisam ser lidos ou gravados através do wrapper PEFT. Isso não é totalmente trivial, pois o PEFT renomeia os parâmetros dos módulos-alvo (por exemplo, q_proj.weight se torna q_proj.base_layer.weight) e adiciona parâmetros de adapter que não existem no modelo base. Use get_base_model_state_dict() e set_base_model_state_dict() para traduzir entre as duas 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)

O principal caso de uso para isso é carregar os pesos base após o modelo já ter sido envolvido pelo PEFT. Por exemplo, configurações de treinamento com FSDP, como o TorchTitan, inicializam o modelo no dispositivo meta, aplicam o PEFT e fazem o sharding do resultado. A memória real só existe após o sharding, portanto, o checkpoint, cujas chaves usam os nomes originais, precisa ser carregado no modelo já envolvido. Observe que get_base_model_state_dict() retorna os tensores vivos do modelo, não cópias.