Modelos personalizados

Algumas técnicas de ajuste fino, como prompt tuning, são específicas para modelos de linguagem. Isso significa que, no 🤗 PEFT, assume-se que um modelo do 🤗 Transformers está sendo utilizado. No entanto, outras técnicas de ajuste fino - como LoRA - não são restritas a tipos específicos de modelos.

Neste guia, veremos como o LoRA pode ser aplicado a um perceptron multicamada, a um modelo de visão computacional da biblioteca timm ou a uma nova arquitetura do 🤗 Transformers.

Perceptron multicamada

Vamos supor que queremos fazer o ajuste fino de um perceptron multicamadas com LoRA. Aqui está a definição:

from torch import nn


class MLP(nn.Module):
    def __init__(self, num_units_hidden=2000):
        super().__init__()
        self.seq = nn.Sequential(
            nn.Linear(20, num_units_hidden),
            nn.ReLU(),
            nn.Linear(num_units_hidden, num_units_hidden),
            nn.ReLU(),
            nn.Linear(num_units_hidden, 2),
            nn.LogSoftmax(dim=-1),
        )

    def forward(self, X):
        return self.seq(X)

Este é um perceptron multicamada simples, com uma camada de entrada, uma camada oculta e uma camada de saída.

Para este exemplo ilustrativo, escolhemos um número excessivamente grande de unidades ocultas para destacar os ganhos de eficiência do PEFT, mas esses ganhos estão em linha com exemplos mais realistas.

Existem algumas camadas lineares neste modelo que podem ser ajustadas com LoRA. Ao trabalhar com modelos comuns do 🤗 Transformers, o PEFT saberá quais camadas aplicar o LoRA, mas, neste caso, cabe a nós, como usuários, escolher as camadas. Para determinar os nomes das camadas a serem ajustadas:

print([(n, type(m)) for n, m in MLP().named_modules()])

Isso deve imprimir:

[('', __main__.MLP),
 ('seq', torch.nn.modules.container.Sequential),
 ('seq.0', torch.nn.modules.linear.Linear),
 ('seq.1', torch.nn.modules.activation.ReLU),
 ('seq.2', torch.nn.modules.linear.Linear),
 ('seq.3', torch.nn.modules.activation.ReLU),
 ('seq.4', torch.nn.modules.linear.Linear),
 ('seq.5', torch.nn.modules.activation.LogSoftmax)]

Suponha que queiramos aplicar LoRA à camada de entrada e à camada oculta, que são 'seq.0' e 'seq.2'. Além disso, suponha que queiramos atualizar a camada de saída sem LoRA, que seria 'seq.4'. A configuração correspondente seria:

from peft import LoraConfig

config = LoraConfig(
    target_modules=["seq.0", "seq.2"],
    modules_to_save=["seq.4"],
)

Com isso, podemos criar nosso modelo PEFT e verificar a fração de parâmetros treinados:

from peft import get_peft_model

model = MLP()
peft_model = get_peft_model(model, config)
peft_model.print_healthcheck()

Por fim, podemos usar qualquer framework de treinamento que preferirmos, ou escrever nosso próprio loop de ajuste, para treinar o peft_model.

Para um exemplo completo, confira este notebook.

modelos timm

A biblioteca timm contém um grande número de modelos de visão computacional pré-treinados. Eles também podem ser ajustados finamente com PEFT. Vamos verificar como isso funciona na prática.

Para começar, verifique se o timm está instalado no ambiente Python:

python -m pip install -U timm

A seguir, carregamos um modelo timm para uma tarefa de classificação de imagens:

import timm

num_classes = ...
model_id = "timm/poolformer_m36.sail_in1k"
model = timm.create_model(model_id, pretrained=True, num_classes=num_classes)

Novamente, precisamos decidir em quais camadas aplicar o LoRA. Como o LoRA suporta camadas de convolução 2D e essas são um bloco de construção fundamental deste modelo, devemos aplicar o LoRA às camadas de convolução 2D. Para identificar os nomes dessas camadas, vamos examinar todos os nomes das camadas:

print([(n, type(m)) for n, m in model.named_modules()])

Isso exibirá uma lista muito longa; mostraremos apenas as primeiras:

[('', timm.models.metaformer.MetaFormer),
 ('stem', timm.models.metaformer.Stem),
 ('stem.conv', torch.nn.modules.conv.Conv2d),
 ('stem.norm', torch.nn.modules.linear.Identity),
 ('stages', torch.nn.modules.container.Sequential),
 ('stages.0', timm.models.metaformer.MetaFormerStage),
 ('stages.0.downsample', torch.nn.modules.linear.Identity),
 ('stages.0.blocks', torch.nn.modules.container.Sequential),
 ('stages.0.blocks.0', timm.models.metaformer.MetaFormerBlock),
 ('stages.0.blocks.0.norm1', timm.layers.norm.GroupNorm1),
 ('stages.0.blocks.0.token_mixer', timm.models.metaformer.Pooling),
 ('stages.0.blocks.0.token_mixer.pool', torch.nn.modules.pooling.AvgPool2d),
 ('stages.0.blocks.0.drop_path1', torch.nn.modules.linear.Identity),
 ('stages.0.blocks.0.layer_scale1', timm.models.metaformer.Scale),
 ('stages.0.blocks.0.res_scale1', torch.nn.modules.linear.Identity),
 ('stages.0.blocks.0.norm2', timm.layers.norm.GroupNorm1),
 ('stages.0.blocks.0.mlp', timm.layers.mlp.Mlp),
 ('stages.0.blocks.0.mlp.fc1', torch.nn.modules.conv.Conv2d),
 ('stages.0.blocks.0.mlp.act', torch.nn.modules.activation.GELU),
 ('stages.0.blocks.0.mlp.drop1', torch.nn.modules.dropout.Dropout),
 ('stages.0.blocks.0.mlp.norm', torch.nn.modules.linear.Identity),
 ('stages.0.blocks.0.mlp.fc2', torch.nn.modules.conv.Conv2d),
 ('stages.0.blocks.0.mlp.drop2', torch.nn.modules.dropout.Dropout),
 ('stages.0.blocks.0.drop_path2', torch.nn.modules.linear.Identity),
 ('stages.0.blocks.0.layer_scale2', timm.models.metaformer.Scale),
 ('stages.0.blocks.0.res_scale2', torch.nn.modules.linear.Identity),
 ('stages.0.blocks.1', timm.models.metaformer.MetaFormerBlock),
 ('stages.0.blocks.1.norm1', timm.layers.norm.GroupNorm1),
 ('stages.0.blocks.1.token_mixer', timm.models.metaformer.Pooling),
 ('stages.0.blocks.1.token_mixer.pool', torch.nn.modules.pooling.AvgPool2d),
 ...
 ('head.global_pool.flatten', torch.nn.modules.linear.Identity),
 ('head.norm', timm.layers.norm.LayerNorm2d),
 ('head.flatten', torch.nn.modules.flatten.Flatten),
 ('head.drop', torch.nn.modules.linear.Identity),
 ('head.fc', torch.nn.modules.linear.Linear)]
 ]

Ao examinar mais de perto, vemos que as camadas de convolução 2D têm nomes como "stages.0.blocks.0.mlp.fc1" e "stages.0.blocks.0.mlp.fc2". Como podemos corresponder especificamente a esses nomes de camada? Você pode escrever uma expressão regular para corresponder aos nomes das camadas. No nosso caso, a regex r".*\.mlp\.fc\d" deve fazer o trabalho.

Além disso, como no primeiro exemplo, devemos garantir que a camada de saída, neste caso a cabeça de classificação, também seja atualizada. Olhando para o final da lista impressa acima, podemos ver que ela é nomeada 'head.fc'. Com isso em mente, aqui está nossa configuração LoRA:

config = LoraConfig(target_modules=r".*\.mlp\.fc\d", modules_to_save=["head.fc"])

Então, só precisamos criar o modelo PEFT passando nosso modelo base e a configuração para get_peft_model:

peft_model = get_peft_model(model, config)
peft_model.print_healthcheck()

Isso nos mostra que precisamos treinar menos de 2% de todos os parâmetros, o que representa um enorme ganho de eficiência.

Para um exemplo completo, confira este notebook.

Novas arquiteturas de transformers

Quando novas arquiteturas populares de transformers são lançadas, fazemos o possível para adicioná-las rapidamente ao PEFT. Se você encontrar um modelo de transformers que não é suportado nativamente, não se preocupe, ele provavelmente ainda funcionará se a configuração estiver definida corretamente. Especificamente, você precisa identificar as camadas que devem ser adaptadas e defini-las corretamente ao inicializar a classe de configuração correspondente, por exemplo, LoraConfig. Aqui estão algumas dicas para ajudar com isso.

Como primeiro passo, é uma boa ideia verificar os modelos existentes para se inspirar. Você pode encontrá-los dentro de constants.py no repositório PEFT. Frequentemente, você encontrará uma arquitetura semelhante que usa os mesmos nomes. Por exemplo, se a nova arquitetura do modelo for uma variação do modelo “mistral” e você quiser aplicar LoRA, pode ver que a entrada para “mistral” em TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING contém ["q_proj", "v_proj"]. Isso indica que, para os modelos “mistral”, os target_modules para LoRA devem ser ["q_proj", "v_proj"]:

from peft import LoraConfig, get_peft_model

my_mistral_model = ...
config = LoraConfig(
    target_modules=["q_proj", "v_proj"],
    ...,  # other LoRA arguments
)
peft_model = get_peft_model(my_mistral_model, config)

Se isso não ajudar, verifique os módulos existentes na arquitetura do seu modelo com o método named_modules e tente identificar as camadas de atenção, especialmente as camadas de chave, consulta e valor. Essas geralmente terão nomes como c_attn, query, q_proj, etc. A camada de chave nem sempre é adaptada e, idealmente, você deve verificar se incluí-la resulta em melhor desempenho.

Além disso, as camadas lineares são alvos comuns para adaptação (por exemplo, no artigo QLoRA, os autores sugerem adaptá-las também). Seus nomes frequentemente contêm as strings fc ou dense.

Se você deseja adicionar um novo modelo ao PEFT, crie uma entrada em constants.py e abra um pull request no repositório. Não se esqueça de atualizar o README também.

Verificar parâmetros e camadas

Você pode verificar se aplicou corretamente um método PEFT ao seu modelo de algumas maneiras.

• Execute print_healthcheck() para inspecionar a fração de parâmetros treináveis e o estado do adaptador. Se a fração de parâmetros for menor ou maior do que o esperado, verifique o repr do modelo imprimindo o modelo. Isso exibe os nomes de todos os tipos de camadas no modelo. Certifique-se de que apenas as camadas-alvo pretendidas sejam substituídas pelas camadas do adaptador. Por exemplo, se o LoRA for aplicado às camadas nn.Linear, você deve ver apenas camadas lora.Linear sendo usadas.

peft_model.print_healthcheck()

• Outra forma de visualizar as camadas adaptadas é usar o atributo targeted_module_names para listar o nome de cada módulo que foi adaptado.

print(peft_model.targeted_module_names)

Tipos de módulo não suportados

Métodos como LoRA só funcionam se os módulos-alvo forem suportados pelo PEFT. Por exemplo, é possível aplicar LoRA às camadas nn.Linear e nn.Conv2d, mas não, por exemplo, à nn.LSTM. Se você encontrar uma classe de camada à qual deseja aplicar o PEFT e ela não for suportada, você pode:

• define um mapeamento personalizado para despachar dinamicamente módulos personalizados em LoRA
• abra uma issue e solicite o recurso, onde os mantenedores o implementarão ou orientarão você sobre como implementá-lo você mesmo, caso a demanda por esse tipo de módulo seja suficientemente alta

Suporte experimental para despacho dinâmico de módulos personalizados em LoRA

Este recurso é experimental e sujeito a alterações, dependendo da aceitação pela comunidade. Introduziremos uma API pública e estável caso haja demanda significativa por ela.

O PEFT oferece uma API experimental para tipos de módulos personalizados para LoRA. Suponha que você tenha uma implementação de LoRA para LSTMs. Normalmente, você não conseguiria informar ao PEFT para usá-la, mesmo que, teoricamente, ela funcionasse com o PEFT. No entanto, isso é possível com o despacho dinâmico de camadas personalizadas.

A API experimental atualmente é assim:

class MyLoraLSTMLayer:
    ...

base_model = ...  # load the base model that uses LSTMs

# add the LSTM layer names to target_modules
config = LoraConfig(..., target_modules=["lstm"])
# define a mapping from base layer type to LoRA layer type
custom_module_mapping = {nn.LSTM: MyLoraLSTMLayer}
# register the new mapping
config._register_custom_module(custom_module_mapping)
# after registration, create the PEFT model
peft_model = get_peft_model(base_model, config)
# do training

Quando você chama get_peft_model(), verá um aviso porque o PEFT não reconhece o tipo de módulo alvo. Nesse caso, você pode ignorar esse aviso.

Ao fornecer um mapeamento personalizado, o PEFT primeiro verifica as camadas do modelo base contra o mapeamento personalizado e direciona para o tipo de camada LoRA personalizada, caso haja uma correspondência. Se não houver correspondência, o PEFT verifica os tipos de camada LoRA embutidos em busca de uma correspondência.

Portanto, esse recurso também pode ser usado para substituir a lógica de despacho existente, por exemplo, se você quiser usar sua própria camada LoRA para nn.Linear em vez de usar a fornecida pelo PEFT.

Ao criar seu módulo LoRA personalizado, siga as mesmas regras dos módulos LoRA existentes. Algumas restrições importantes a considerar:

• O módulo personalizado deve herdar de nn.Module e peft.tuners.lora.layer.LoraLayer.
• O método __init__ do módulo personalizado deve ter os argumentos posicionais base_layer e adapter_name. Após isso, há argumentos adicionais **kwargs que você pode usar ou ignorar à sua discrição.
• Os parâmetros aprendíveis devem ser armazenados em um nn.ModuleDict ou nn.ParameterDict, onde a chave corresponde ao nome do adaptador específico (lembre-se de que um modelo pode ter mais de um adaptador ao mesmo tempo).
• O nome desses atributos de parâmetros aprendíveis deve começar com "lora_", por exemplo, self.lora_new_param = ....
• Alguns métodos são opcionais, por exemplo, você só precisa implementar merge e unmerge se quiser suportar a mesclagem de pesos.

Atualmente, as informações sobre o módulo personalizado não persistem ao salvar o modelo. Ao carregar o modelo, é necessário registrar os módulos personalizados novamente.

# saving works as always and includes the parameters of the custom modules
peft_model.save_pretrained(<model-path>)

# loading the model later:
base_model = ...
# load the LoRA config that you saved earlier
config = LoraConfig.from_pretrained(<model-path>)
# register the custom module again, the same way as the first time
custom_module_mapping = {nn.LSTM: MyLoraLSTMLayer}
config._register_custom_module(custom_module_mapping)
# pass the config instance to from_pretrained:
peft_model = PeftModel.from_pretrained(model, tmp_path / "lora-custom-module", config=config)

Se você usar esse recurso e achar útil, ou se encontrar problemas, nos avise criando uma issue ou uma discussão no GitHub. Isso nos permite estimar a demanda por esse recurso e adicionar uma API pública, caso ela seja suficientemente alta.