Formato de checkpoint do PEFT

Este documento descreve como os arquivos de checkpoint do PEFT são estruturados e como converter entre o formato PEFT e outros formatos.

Arquivos PEFT

Os métodos de PEFT (fine-tuning de parâmetros eficientes) atualizam apenas um subconjunto pequeno dos parâmetros de um modelo, em vez de todos eles. Isso é vantajoso porque os arquivos de checkpoint geralmente podem ser muito menores do que os arquivos do modelo original e são mais fáceis de armazenar e compartilhar. No entanto, isso também significa que, para carregar um modelo PEFT, você precisa ter o modelo original disponível também.

Quando você chama save_pretrained() em um modelo PEFT, o modelo PEFT salva três arquivos, descritos abaixo:

1. adapter_model.safetensors ou adapter_model.bin

Por padrão, o modelo é salvo no formato safetensors, uma alternativa segura ao formato bin, que é conhecido por ser suscetível a vulnerabilidades de segurança porque utiliza a utilidade pickle por baixo dos panos. Ambos os formatos armazenam o mesmo state_dict, no entanto, e são intercambiáveis.

O state_dict contém apenas os parâmetros do módulo de adaptador, não do modelo base. Para ilustrar a diferença de tamanho, um modelo BERT normal requer ~420MB de espaço em disco, enquanto um adaptador IA³ sobre esse modelo BERT requer apenas ~260KB.

2. adapter_config.json

O arquivo adapter_config.json contém a configuração do módulo de adaptador, que é necessária para carregar o modelo. Abaixo está um exemplo de um adapter_config.json para um adaptador IA³ com configurações padrão aplicadas a um 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
}

O arquivo de configuração contém:

• o tipo de módulo de adaptador armazenado, "peft_type": "IA3"
• informações sobre o modelo base, como "base_model_name_or_path": "bert-base-uncased"
• a revisão do modelo (se houver), "revision": null

Se o modelo base não for um modelo pré-treinado do Transformers, as duas últimas entradas serão null. Além disso, as configurações estão todas relacionadas ao adaptador IA³ específico que foi usado para o ajuste fino do modelo.

3. README.md

O README.md gerado é o cartão do modelo de um modelo PEFT e contém algumas entradas pré-preenchidas. A intenção disso é facilitar o compartilhamento do modelo com outras pessoas e fornecer algumas informações básicas sobre o modelo. Este arquivo não é necessário para carregar o modelo.

Converter para o formato PEFT

Ao converter de outro formato para o formato PEFT, exigimos tanto o arquivo adapter_model.safetensors (ou adapter_model.bin) quanto o arquivo adapter_config.json.

adapter_model

Para os pesos do modelo, é importante usar o mapeamento correto de nome de parâmetro para valor para que o PEFT carregue o arquivo. Obter esse mapeamento correto é um exercício de verificação dos detalhes da implementação, pois não há um formato geralmente acordado para adaptadores PEFT.

Felizmente, descobrir esse mapeamento não é excessivamente complicado para os casos base comuns. Vamos examinar um exemplo concreto, o 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

No código __init__ usado por todas as classes LoraLayer no PEFT, há uma série de parâmetros usados para inicializar o modelo, mas apenas alguns são relevantes para o arquivo de checkpoint: lora_A, lora_B, lora_embedding_A e lora_embedding_B. Esses parâmetros estão listados no atributo de classe adapter_layer_names e contêm os parâmetros aprendíveis, portanto, devem ser incluídos no arquivo de checkpoint. Todos os outros parâmetros, como o rank r, são derivados do adapter_config.json e devem ser incluídos nele (a menos que o valor padrão seja usado).

Vamos verificar o state_dict de um modelo PEFT LoRA aplicado ao BERT. Ao imprimir as cinco primeiras chaves usando as configurações padrão do LoRA (as chaves restantes são as mesmas, apenas com números de camada diferentes), obtemos:

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

Vamos analisar isso:

• Por padrão, para modelos BERT, o LoRA é aplicado às camadas query e value do módulo de atenção. É por isso que você vê attention.self.query e attention.self.value nos nomes das chaves de cada camada.
• LoRA decompõe os pesos em duas matrizes de baixo posto, lora_A e lora_B. É daí que lora_A e lora_B vêm nos nomes das chaves.
• Essas matrizes LoRA são implementadas como camadas nn.Linear, portanto, os parâmetros são armazenados no atributo .weight (lora_A.weight, lora_B.weight).
• Por padrão, o LoRA não é aplicado à camada de incorporação do BERT, portanto não há entradas para lora_A_embedding e lora_B_embedding.
• As chaves do state_dict sempre começam com "base_model.model.". O motivo é que, no PEFT, envolvemos o modelo base dentro de um modelo específico do tuner (LoraModel neste caso), que por sua vez é envolvido em um modelo PEFT geral (PeftModel). Por esse motivo, esses dois prefixos são adicionados às chaves. Ao converter para o formato PEFT, é necessário adicionar esses prefixos.

Esse último ponto não se aplica a técnicas de prefix tuning, como o prompt tuning. Nelas, os embeddings extras são armazenados diretamente no state_dict sem que nenhum prefixo seja adicionado às chaves.

Ao inspecionar os nomes dos parâmetros no modelo carregado, você pode se surpreender ao perceber que eles parecem um pouco diferentes, por exemplo, base_model.model.encoder.layer.0.attention.self.query.lora_A.default.weight. A diferença está na parte .default no penúltimo segmento. Essa parte existe porque o PEFT geralmente permite a adição de vários adaptadores de uma vez (usando um nn.ModuleDict ou nn.ParameterDict para armazená-los). Por exemplo, se você adicionar outro adaptador chamado “other”, a chave para esse adaptador seria base_model.model.encoder.layer.0.attention.self.query.lora_A.other.weight.

Quando você chama save_pretrained(), o nome do adaptador é removido das chaves. A razão é que o nome do adaptador não é uma parte importante da arquitetura do modelo; é apenas um nome arbitrário. Ao carregar o adaptador, você poderia escolher um nome totalmente diferente, e o modelo ainda funcionaria da mesma maneira. É por isso que o nome do adaptador não é armazenado no arquivo de checkpoint.

Se você chamar save_pretrained("some/path") e o nome do adaptador não for "default", o adaptador será armazenado em um subdiretório com o mesmo nome do adaptador. Portanto, se o nome for “other”, ele será armazenado dentro de some/path/other.

Em algumas circunstâncias, decidir quais valores adicionar ao arquivo de checkpoint pode se tornar um pouco mais complicado. Por exemplo, no PEFT, o DoRA é implementado como um caso especial do LoRA. Se você quiser converter um modelo DoRA para PEFT, deve criar um checkpoint LoRA com entradas extras para o DoRA. Você pode ver isso no __init__ do código anterior de LoraLayer:

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

Isso indica que há um parâmetro extra opcional por camada para o DoRA.

adapter_config

Todas as outras informações necessárias para carregar um modelo PEFT estão contidas no arquivo adapter_config.json. Vamos verificar esse arquivo para um modelo LoRA aplicado ao 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
}

Isso contém muitas entradas e, à primeira vista, pode parecer esmagador descobrir todos os valores corretos para inserir. No entanto, a maioria das entradas não é necessária para carregar o modelo. Isso ocorre porque elas usam os valores padrão e não precisam ser adicionadas, ou porque afetam apenas a inicialização dos pesos do LoRA, o que é irrelevante quando se trata de carregar o modelo. Se você descobrir que não sabe o que um parâmetro específico faz, por exemplo, "use_rslora", não o adicione e tudo deve funcionar sem problemas. Observe também que, à medida que mais opções forem adicionadas, este arquivo terá mais entradas no futuro, mas deve manter compatibilidade reversa.

No mínimo, você deve incluir as seguintes entradas:

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

No entanto, é recomendável adicionar o máximo de entradas possível, como o rank r ou o base_model_name_or_path (se for um modelo do Transformers). Essas informações podem ajudar outras pessoas a entender melhor o modelo e compartilhá-lo com mais facilidade. Para verificar quais chaves e valores são esperados, consulte o arquivo config.py (por exemplo, este é o arquivo de configuração para LoRA) no código-fonte do PEFT.

Armazenamento de modelos

Em algumas circunstâncias, você pode querer armazenar o modelo PEFT inteiro, incluindo os pesos da base. Isso pode ser necessário se, por exemplo, o modelo base não estiver disponível para os usuários que tentam carregar o modelo PEFT. Você pode mesclar os pesos primeiro ou convertê-lo em um modelo Transformer.

Mergear os pesos

A maneira mais direta de armazenar o modelo PEFT completo é mesclar os pesos do adaptador nos pesos da base:

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

No entanto, essa abordagem apresenta algumas desvantagens:

• Quando merge_and_unload() é chamado, você obtém um modelo básico sem nenhuma funcionalidade específica do PEFT. Isso significa que você não pode mais usar nenhum dos métodos específicos do PEFT.
• Você não pode desmesclar os pesos, carregar vários adaptadores de uma vez, desabilitar o adaptador, etc.
• Nem todos os métodos de PEFT suportam a mesclagem de pesos.
• Alguns métodos de PEFT podem, em geral, permitir a fusão, mas não com configurações específicas (por exemplo, ao usar certas técnicas de quantização).
• O modelo completo será muito maior que o modelo PEFT, pois conterá todos os pesos da base também.

Mas a inferência com um modelo mesclado deve ser um pouco mais rápida.

Converter para um modelo do Transformers

Outra forma de salvar o modelo inteiro, assumindo que o modelo base seja um modelo do Transformers, é usar essa abordagem “hacky” para inserir diretamente os pesos do PEFT no modelo base e salvá-lo, o que só funciona se você “enganar” o Transformers para que ele acredite que o modelo PEFT não é um modelo PEFT. Isso só funciona com LoRA, pois outros adaptadores não estão implementados no 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>)