Implementando LoRA desde cero

Estrategias para implementar LoRA desde cero

Cómo implementar LoRA desde cero y algunos consejos prácticos

En esta publicación de blog, te mostraré cómo implementar LoRA desde cero.

LoRA, acrónimo de Adaptación de Baja Rango o Adaptadores de Baja Rango, ofrece un método eficiente y liviano para ajustar modelos de lenguaje preexistentes. Esto incluye modelos de lenguaje enmascarados como BERT y RoBERTa, así como modelos causales (o de chatbot) como GPT, Llama y Mistral.

Una de las principales ventajas de los adaptadores de bajo rango es su eficiencia. Al utilizar menos parámetros, los LoRA reducen significativamente la complejidad computacional y el uso de memoria. Esto nos permite entrenar modelos grandes en GPUs de consumo y distribuir sin esfuerzo nuestros compactos LoRA (en términos de megabytes) a otros.

Además, los LoRA pueden mejorar el rendimiento de generalización. Al limitar la complejidad del modelo, ayudan a prevenir el sobreajuste, especialmente en escenarios con datos de entrenamiento limitados. Esto resulta en modelos más resilientes que sobresalen con datos nuevos o no vistos, o al menos, mantienen el conocimiento de sus tareas de entrenamiento iniciales.

• La Tierra no es plana, y tus diagramas de Voronoi tampoco deberían serlo
• Segmentar cualquier cosa en 3D para nubes de puntos Guía completa (SAM 3D)
• Crea una interfaz de usuario web para interactuar con LLMs utilizando Amazon SageMaker JumpStart

Además, los adaptadores de bajo rango se pueden integrar fácilmente en arquitecturas existentes de redes neuronales. Esta integración permite el ajuste fino y la adaptación de modelos preentrenados con un costo adicional de entrenamiento mínimo, lo que los hace muy adecuados para aplicaciones de transferencia de aprendizaje.

Comenzaremos profundizando en cómo funciona LoRA, luego demostraré cómo desarrollarlo desde cero para un modelo RoBERTa, seguido de una prueba de nuestro desarrollo utilizando los criterios de evaluación GLUE y SQuAD, junto con una discusión sobre consejos generales y mejoras.

Cómo funciona LoRA

La idea básica de LoRA es mantener las matrices preentrenadas (es decir, los parámetros del modelo original) congeladas (es decir, en un estado fijo) y solo agregar un pequeño delta a la matriz original, que tiene menos parámetros que la matriz original.

Por ejemplo, considera la matriz W, que podría ser los parámetros de una capa completamente conectada o una de las matrices del mecanismo de autoatención de un transformador:

Obviamente, si W-orig tiene dimensiones n×m y simplemente inicializamos una nueva matriz delta con las mismas dimensiones para ajustarlo fino, no habríamos ganado nada; al contrario, habríamos duplicado los parámetros.

El truco es hacer que ΔW sea menos “dimensional” que la matriz original, construyéndola mediante multiplicación de matrices a partir de matrices de menor dimensión B y A.

Donde primero definimos un rango r, considerablemente menor que las dimensiones base de la matriz r≪n y r≪m. Luego, la matriz B es n×r y la matriz A es r×m. Al multiplicarlos, obtenemos una matriz con las mismas dimensiones de W, pero construida a partir de un conteo de parámetros mucho menor.

Obviamente, queremos que nuestro delta sea cero al inicio del entrenamiento, de modo que el ajuste fino comience como el modelo original. Por lo tanto, B a menudo se inicializa como ceros y A se inicializa como valores aleatorios (generalmente distribuidos normalmente).

Por ejemplo, esto podría verse así:

Imagina una situación en la que nuestra dimensionalidad base es 1024 y elegimos un rango LoRA de 4 entonces:

• W tiene 1024 * 1024 ≈ 1 millón de parámetros
• A & B tienen r * 1024 = 4 * 1024 ≈ 4k parámetros cada uno, lo que da un total de 8k
• Por lo tanto, solo tenemos que entrenar el 0,8% de los parámetros para actualizar nuestra matriz con LoRA

Un pequeño aparte, en el artículo de LoRA pesan la matriz delta con un parámetro alpha:

Si solo estableces tu α al primer r con el que experimentas y ajustas la tasa de aprendizaje, generalmente puedes cambiar el parámetro r más tarde sin tener que ajustar la tasa de aprendizaje nuevamente (al menos aproximadamente). Si bien podemos pasar por alto este detalle en nuestra implementación, es una característica común en muchas otras bibliotecas de LoRA, como PEFT de Hugging Face.

Implementando LoRA

Para nuestra implementación, queremos ceñirnos de cerca al artículo original de LoRA. Allí probaron qué matrices de un transformador realmente debes reemplazar. Descubrieron que, al comparar diferentes estrategias en una tarea de ajuste fino de GPT-3, era suficiente adaptar solo los vectores de consulta y valor del mecanismo de autoatención.

Ten en cuenta que hoy en día muchas personas ignoran esta evaluación y permiten que se ajuste cada matriz, sin importar la tarea o el modelo (ver artículo de QLoRA).

Nuestra implementación aquí se realizará en PyTorch, pero debería ser fácilmente adaptable a diferentes frameworks.

Para esta publicación de blog, simplifiqué un poco el código, de modo que debería ser más fácil de leer, al tiempo que muestra los elementos esenciales. El código completo y algunos pesos de LoRA entrenados se pueden encontrar aquí: https://github.com/Montinger/Transformer-Workbench.

Volver a implementar el modelo de autoatención

El modelo que deseamos adaptar es el modelo RoBERTa de Huggingface. La forma más directa es simplemente envolver nuevamente el mecanismo de autoatención original RobertaSelfAttention. La nueva clase LoraRobertaSelfAttention inicializará las matrices de LoRA. Todas las matrices B se inicializarán con ceros y todas las matrices A con números aleatorios de una distribución normal.

class LoraRobertaSelfAttention(RobertaSelfAttention):    """    Extiende RobertaSelfAttention con matrices de LoRA (Adaptación de bajo rango).    LoRA mejora la eficiencia al actualizar solo las matrices de consulta y valor.    Esta clase agrega matrices de LoRA y aplica la lógica de LoRA en el método forward.    Parámetros:    - r (int): Rango para matrices de LoRA.    - config: Configuración del modelo Roberta.    """    def __init__(self, r=8, *args, **kwargs):        super().__init__(*args, **kwargs)        d = self.all_head_size        # Inicializar matrices de LoRA para query y value        self.lora_query_matrix_B = nn.Parameter(torch.zeros(d, r))        self.lora_query_matrix_A = nn.Parameter(torch.randn(r, d))        self.lora_value_matrix_B = nn.Parameter(torch.zeros(d, r))        self.lora_value_matrix_A = nn.Parameter(torch.randn(r, d))

Dadas estas matrices, ahora definimos nuevos métodos de clase lora_query y lora_value. Estos calculan la matriz ΔW, es decir, BA, y la agregan a la matriz original, a la cual hacemos referencia desde los métodos originales query y value.

class LoraRobertaSelfAttention(RobertaSelfAttention):    # ...    def lora_query(self, x):        """        Aplica LoRA a la parte de consulta. Calcula una salida de consulta modificada agregando         la adaptación de LoRA a la salida de consulta estándar. Requiere que la capa lineal regular         esté congelada antes del entrenamiento.        """        lora_query_weights = torch.matmul(self.lora_query_matrix_B, self.lora_query_matrix_A)        return self.query(x) + F.linear(x, lora_query_weights)    def lora_value(self, x):        """        Aplica LoRA a la parte de valor. Calcula una salida de valor modificada agregando         la adaptación de LoRA a la salida de valor estándar. Requiere que la capa lineal regular         esté congelada antes del entrenamiento.        """        lora_value_weights = torch.matmul(self.lora_value_matrix_B, self.lora_value_matrix_A)        return self.value(x) + F.linear(x, lora_value_weights)

Ahora la parte fea: para usar los métodos, debemos sobrescribir la función forward original de RobertaSelfAttention. Aunque esto está un poco codificado (ver la discusión sobre mejoras más adelante), es bastante simple. Primero, copiamos el código forward original de https://github.com/huggingface/transformers/blob/main/src/transformers/models/roberta/modeling_roberta.py. Segundo, reemplazamos cada llamada a query por lora_query y cada llamada a value por lora_value. La función se verá así:

class LoraRobertaSelfAttention(RobertaSelfAttention):    # ...    def forward(self, hidden_states, *args, **kwargs):        """Copiado desde https://github.com/huggingface/transformers/blob/main/src/transformers/models/roberta/modeling_roberta.py         pero reemplazando las llamadas a query y value con llamadas a         las funciones lora_query y lora_value.         Aquí solo daremos una idea de cómo ajustar esto.          Cambiar todas las llamadas a self.value y self.query en la versión actual.        """        # código original para query:        ## mixed_query_layer = self.query(hidden_states)        # query actualizada para LoRA:        mixed_query_layer = self.lora_query(hidden_states)        # La clave no tiene LoRA, por lo que dejamos estas llamadas sin cambios        key_layer = self.transpose_for_scores(self.key(hidden_states))        # código original para value:        ## value_layer = self.transpose_for_scores(self.value(hidden_states))        # value actualizada para LoRA:        value_layer = self.transpose_for_scores(self.lora_value(hidden_states))                 # ... (resto del código forward, sin cambios)

¡Tada, ahí lo tenemos! Nuestra implementación de nuestra autoatención LoRA. Ahora, la única tarea que queda es reemplazar los módulos de atención en el modelo RoBERTa original.

Reemplazando los módulos

Genial, hemos reemplazado la autoatención con nuestra propia implementación; pero, ¿cómo conseguimos que esta nueva clase esté en el viejo modelo RoBERTa? Básicamente, tenemos que recorrer cada componente con nombre del modelo RoBERTa, comprobar si es de la clase RobertaSelfAttention y, si es así, reemplazarlo por LoraRobertaSelfAttention, asegurándonos de que se conserven las matrices de peso originales.

Para lograr esto, escribiremos una nueva función de envoltura que pueda hacer esta sustitución. Además, también queremos agregar la funcionalidad para afinar el modelo RoBERTa en algunas tareas reales más adelante.

class LoraWrapperRoberta(nn.Module):    def __init__(self, task_type, num_classes=None, dropout_rate=0.1, model_id="roberta-large",                 lora_rank=8, train_biases=True, train_embedding=False, train_layer_norms=True):        """        Un envoltorio para RoBERTa con Adaptación de baja clasificación (LoRA) para varias tareas de Procesamiento del Lenguaje Natural (PLN).         - task_type: Tipo de tarea de PLN ('glue', 'squad_v1', 'squad_v2').         - num_classes: Número de clases para clasificación (varía según la tarea).         - dropout_rate: Tasa de abandono en el modelo.         - model_id: ID del modelo RoBERTa preentrenado.         - lora_rank: Rango para la adaptación LoRA.         - train_biases, train_embedding, train_layer_norms:            Banderas para mantener ciertos parámetros entrenables            después de inicializar LoRA.                Ejemplo:            model = LoraWrapperRoberta(task_type='glue')        """        super().__init__()        # 1. Inicializar el modelo base con los parámetros        self.model_id = model_id        self.tokenizer = RobertaTokenizer.from_pretrained(model_id)        self.model = RobertaModel.from_pretrained(model_id)        self.model_config = self.model.config        # 2. Agregar la capa para las tareas de referencia        d_model = self.model_config.hidden_size        self.finetune_head_norm = nn.LayerNorm(d_model)        self.finetune_head_dropout = nn.Dropout(dropout_rate)        self.finetune_head_classifier = nn.Linear(d_model, num_classes)        # 3. Configurar el modelo LoRA para el entrenamiento        self.replace_multihead_attention()        self.freeze_parameters_except_lora_and_bias()

Como puedes ver, llamamos a dos métodos auxiliares en la inicialización:

1. self.replace_multihead_attention: Esto reemplaza la atención de todas las partes de la red neuronal por nuestra LoraRobertaSelfAttention previamente escrita
2. self.freeze_parameters_except_lora_and_bias: Esto congela todos los parámetros principales para el entrenamiento, de manera que los gradientes y los pasos del optimizador solo se apliquen a los parámetros LoRA y a los otros parámetros de sesgo y de normalización de capa que queremos mantener entrenables.

class LoraWrapperRoberta(nn.Module):    # ...    def replace_multihead_attention_recursion(self, model):        """        Reemplaza RobertaSelfAttention con LoraRobertaSelfAttention en el modelo.        Este método aplica el reemplazo de forma recursiva a todos los subcomponentes.        Parámetros        ----------        model : nn.Module            El módulo o modelo PyTorch que se modificará.        """        for name, module in model.named_children():            if isinstance(module, RobertaSelfAttention):                # Reemplazar RobertaSelfAttention con LoraRobertaSelfAttention                nueva_capa = LoraRobertaSelfAttention(r=self.lora_rank, config=self.model_config)                nueva_capa.load_state_dict(module.state_dict(), strict=False)                setattr(model, name, nueva_capa)            else:                # Llamada recursiva para módulos hijos                self.replace_multihead_attention_recursion(module)

Debemos recorrer de forma recursiva todas las partes del modelo, ya que en PyTorch las partes de la red pueden (y de hecho lo están en RoBERTa) empaquetarse en un módulo PyTorch separado.

Ahora debemos congelar todos los parámetros que ya no queremos entrenar:

class LoraWrapperRoberta(nn.Module):    # ...    def freeze_parameters_except_lora_and_bias(self):        """        Congela todos los parámetros del modelo excepto ciertas capas y tipos basados en la configuración.        Los parámetros en capas LoRA, la cabeza fine-tune, los parámetros de sesgo, las incrustaciones y las capas de normalización         se pueden configurar como entrenables según las configuraciones de la clase.        """        for name, param in self.model.named_parameters():            es_entrenable = (                "lora_" in name or                "finetune_head_" in name or                (self.train_biases and "bias" in name) or                (self.train_embeddings and "embeddings" in name) or                (self.train_layer_norms and "LayerNorm" in name)            )            param.requires_grad = es_entrenable

Además, debemos implementar los métodos forward para tener en cuenta las tareas en las que vamos a realizar el ajuste fino, así como dos métodos para guardar y cargar los pesos LoRA, de manera que podamos cargar los adaptadores de un modelo previamente entrenado.

Cliffhanger: Existe una forma que habría hecho que el código fuera mucho más limpio y fácil de generalizar a otras arquitecturas de red (ya que el nuestro está bastante específico al modelo RoBERTa). ¿Puedes adivinar cuál podría ser? Tienes tiempo para reflexionar sobre esta pregunta hasta que lo discutamos en la sección de Mejoras Posibles más adelante. Pero hasta entonces: Probemos en algunos puntos de referencia si nuestra implementación realmente funciona.

Probando los resultados con GLUE y SQuAD

Nuestra implementación está lista para ser evaluada utilizando los puntos de referencia de GLUE (Evaluación General del Entendimiento del Lenguaje) y SQuAD (Conjunto de Datos de Preguntas y Respuestas de Stanford).

El punto de referencia GLUE, una suite de ocho tareas de procesamiento del lenguaje natural diversas, evalúa las capacidades de comprensión integral de un modelo de lenguaje. Incluye desafíos como el análisis de sentimiento, la implicación textual y la similitud de frases, ofreciendo una medida sólida de la capacidad de adaptación y competencia lingüística de un modelo.

SQuAD, por otro lado, se enfoca en evaluar modelos de preguntas y respuestas. Involucra la extracción de respuestas de pasajes de Wikipedia, donde el modelo identifica el fragmento de texto relevante. SQuAD v2, una versión más avanzada, introduce preguntas sin respuesta, añadiendo complejidad y reflejando situaciones de la vida real donde los modelos deben reconocer cuando el texto no tiene una respuesta.

Es importante destacar que para el siguiente punto de referencia, no ajusté ningún hiperparámetro, no realicé múltiples ejecuciones (especialmente en los conjuntos de datos más pequeños de GLUE, que tienden a tener ruido estocástico), no hice paradas tempranas y no comencé con una puesta a punto en una tarea de GLUE anterior (como se hace a menudo para disminuir la variabilidad del ruido del conjunto de datos pequeño y evitar el sobreajuste).

Todas las ejecuciones:

• Partieron de una inyección fresca de LoRA con un rango de 8 en el modelo RoBERTa-base
• El entrenamiento se realizó exactamente durante 6 épocas para cada tarea, sin ninguna detención temprana.
• Durante las primeras 2 épocas, la tasa de aprendizaje se aumentó linealmente hasta el valor máximo, y luego disminuyó linealmente hasta cero durante las 4 épocas restantes.
• La tasa de aprendizaje máxima para todas las tareas fue de 5e-4.
• El tamaño del lote (batch size) para todas las tareas fue de 16.

El modelo RoBERTa-base tiene 124.6 millones de parámetros. Con los parámetros de LoRA, los sesgos y las reglas de capa, solo tenemos 420 mil parámetros descongelados para entrenar. Esto significa que básicamente entrenamos solo con el 0.34% de los parámetros originales.

El número de parámetros introducidos por LoRA para estas tareas específicas es notablemente mínimo, que equivale a solo 1.7 MB de tamaño real en disco. Puedes encontrar los LoRAs entrenados en el repositorio de Git en la carpeta de Salida.

Después del entrenamiento, volvimos a cargar los parámetros de LoRA, los aplicamos nuevamente y probamos el rendimiento en el conjunto de validación de cada tarea. A continuación se muestran los resultados:

Es probable que estos resultados puedan mejorar significativamente con una afinación de hiperparámetros. Sin embargo, esto claramente demuestra que nuestra implementación de LoRA está funcionando y nuestras matrices de rango bajo están aprendiendo.

Posibles Mejoras

Reflexionando sobre nuestra implementación, uno podría preguntarse: “¿Podría haber existido un enfoque más eficiente y generalizable (es decir, aplicable a otras arquitecturas de red) que volver a codificar la clase de autoatención y realizar reemplazos complejos?”

De hecho, podríamos haber simplemente implementado un envoltorio alrededor de la función nn.Linear de pytorch y ser más específicos sobre qué capas queremos reemplazar con ella, revisando sus nombres. Del mismo modo, podrías escribir envoltorios alrededor de la mayoría de las capas pytorch base y ser capaz de adaptar rápidamente LoRA a nuevas arquitecturas de red. Para dar un boceto rápido de cómo se podría hacer esto:

class LoraLinear(nn.Linear):    """    Extiende una capa lineal de PyTorch con Adaptación de Bajo Rango (LoRA).    LoRA agrega dos matrices a la capa, permitiendo un entrenamiento eficiente de modelos grandes.    """    def __init__(self, in_features, out_features, r=8, *args, **kwargs):        super().__init__(in_features, out_features, *args, **kwargs)        # Inicializa las matrices LoRA        self.lora_matrix_B = nn.Parameter(torch.zeros(out_features, r))        self.lora_matrix_A = nn.Parameter(torch.randn(r, in_features))                # Congela la matriz de pesos original        self.weight.requires_grad = False    def forward(self, x: Tensor) -> Tensor:        # Calcula el ajuste de peso LoRA        lora_weights = torch.matmul(self.lora_matrix_B, self.lora_matrix_A)        # Aplica las transformaciones lineales originales y ajustadas por LoRA        return super().forward(x) + F.linear(x, lora_weights)

De hecho, esta es la forma en que la biblioteca PEFT (Parameter-Efficient Fine-Tuning) de huggingface implementa LoRA. Para cualquier aplicación práctica, donde no estés tratando de aprender, te recomiendo encarecidamente que lo uses en lugar de codificar el tuyo propio.

También se ha vuelto una práctica bastante común inyectar LoRA en todas las capas lineales también (es decir, todas las matrices de autoatención y las dos capas lineales para la red directa completamente conectada). Por lo general, es una buena idea mantener los sesgos y las reglas de capa entrenables, además de los parámetros de LoRA. Como ya son pequeños, no necesitarás inyectar rango bajo en ellos.

Es recomendable cuantificar los pesos de la matriz original para conservar la VRAM de la GPU, lo que facilita el entrenamiento de modelos más grandes en una GPU dada. Esto se puede hacer eficientemente utilizando la biblioteca bits-and-bytes, ahora totalmente integrada con Hugging Face (consultar referencias).

En resumen, aquí están los Cinco Mandamientos de la Adaptación de Bajo Rango en un entorno serio:

Si encuentras difícil leer la tabla de piedra inscrita, aquí están de nuevo en texto plano:

Los Cinco Mandamientos de la Adaptación de Bajo Rango

1. Utiliza LoRA para afinar eficientemente el modelo, enfocándote en mantener el tamaño de los parámetros mínimo.2. Emplea la biblioteca PEFT para la implementación de LoRA, evitando la necesidad de programación compleja.3. Extiende las adaptaciones de LoRA a todas las capas lineales, mejorando las capacidades generales del modelo.4. Mantén los sesgos y las capas normales entrenables, ya que son críticos para la adaptabilidad del modelo y no requieren adaptaciones de bajo rango.5. Aplica Quantized-LoRA – QLoRA – para preservar la memoria VRAM de la GPU y entrenar tu modelo, permitiendo el entrenamiento de modelos más grandes.

Recuerda, el entrenamiento con QLoRA puede ser un poco más lento que LoRA, ya que implica descuantificar matrices durante cada multiplicación. Por ejemplo, al afinar algo masivo como Llama-7B, QLoRA requiere aproximadamente un 75% menos de memoria VRAM pero es aproximadamente un 40% más lento en comparación con LoRA estándar. Para obtener más información, consulta los artículos de blog que he vinculado en las referencias.

Guía Paso a Paso para la Implementación de PEFT

Veamos cómo realmente obedecer nuestros mandamientos e implementar una mejor versión a través de PEFT.

Primero, carguemos nuestro modelo de manera cuantizada. Gracias a la integración de bitsandbytes con la biblioteca Huggingface transformers (introducida en mayo de 2023), esto es muy sencillo.

Tenemos que especificar un archivo de configuración y luego cargar el modelo directamente desde huggingface con esta cuantización. Por lo general, es mejor usar los objetos AutoModel de transformers. Es difícil cargar un modelo cuantizado como un submódulo de un objeto nn.module más grande y recién definido. Debes trabajar generalmente con los modelos crudos de huggingface e importar directamente un AutoModelForSequenceClassification para las tareas de GLUE y un AutoModelForQuestionAnswering para los benchmarks de SQuAD. En la configuración también podemos especificar qué parámetros no cuantizar: aquí debemos registrar las cabezas de clasificación o de salida de qa, ya que queremos entrenar estas por completo, es decir, sin LoRA, ya que fueron inicializadas nuevamente para el ajuste fino y nunca fueron parte del modelo base pre-entrenado.

import bitsandbytes as bnbfrom transformers import AutoModel, AutoModelForSequenceClassification, BitsAndBytesConfig# Configuración para cargar un modelo cuantizadobnb_config = BitsAndBytesConfig(    load_in_4bit=True,  # Habilitar la carga de 4 bits    bnb_4bit_quant_type="nf4",    bnb_4bit_compute_dtype=torch.bfloat16,    llm_int8_skip_modules=['classifier', 'qa_outputs'],  # Omitir estos para la cuantización)# Carga el modelo de Huggingface con cuantizaciónmodel = AutoModelForSequenceClassification.from_pretrained('roberta-base',          torch_dtype="auto", quantization_config=bnb_config)

Puedes verificar la carga de 4 bits inspeccionando los módulos del modelo y los tipos de datos de los parámetros:

# Verificar la carga de 4 bitsprint("Verificando elementos de 4 bits (Linear4bit) en la capa de atención:")print(model.roberta.encoder.layer[4].attention)print("Comprobando el tipo de datos uint8:")print(model.roberta.encoder.layer[4].attention.self.query.weight.dtype)

Ahora, inyectemos los parámetros de LoRA con PEFT. Ten en cuenta que la biblioteca PEFT es mucho más flexible, también cuando se trabaja con modelos personalizados u otras estructuras complicadas, siempre y cuando solo estés haciendo LoRA en lugar de QLoRA (la cuantización suele ser la parte difícil).

La biblioteca PEFT apunta a los módulos a reemplazar a través de sus nombres; por lo tanto, debemos echar un vistazo a los parámetros nombrados del modelo model.named_parameters(). Así es como se ve esto para el modelo roberta-base no cuantizado.

Módulo                                                        Parámetros----------------------------------------------------------  ------------roberta.embeddings.word_embeddings.weight                     38_603_520roberta.embeddings.position_embeddings.weight                    394_752roberta.embeddings.token_type_embeddings.weight                      768roberta.embeddings.LayerNorm.weight                                  768roberta.embeddings.LayerNorm.bias                                    768roberta.encoder.layer.0.attention.self.query.weight              589_824roberta.encoder.layer.0.attention.self.query.bias                    768roberta.encoder.layer.0.attention.self.key.weight                589_824roberta.encoder.layer.0.attention.self.key.bias                      768roberta.encoder.layer.0.attention.self.value.weight              589_824roberta.encoder.layer.0.attention.self.value.bias                    768roberta.encoder.layer.0.attention.output.dense.weight            589_824roberta.encoder.layer.0.attention.output.dense.bias                  768roberta.encoder.layer.0.attention.output.LayerNorm.weight            768roberta.encoder.layer.0.attention.output.LayerNorm.bias              768roberta.encoder.layer.0.intermediate.dense.weight              2_359_296roberta.encoder.layer.0.intermediate.dense.bias                    3_072roberta.encoder.layer.0.output.dense.weight                    2_359_296roberta.encoder.layer.0.output.dense.bias                            768roberta.encoder.layer.0.output.LayerNorm.weight                      768roberta.encoder.layer.0.output.LayerNorm.bias                        768roberta.encoder.layer.1.attention.self.query.weight              589_824...roberta.encoder.layer.11.output.LayerNorm.bias                       768classifier.dense.weight                                          589_824classifier.dense.bias                                                768classifier.out_proj.weight                                         1_536classifier.out_proj.bias                                               2----------------------------------------------------------  ------------TOTAL                                                        124_647_170

A continuación, podemos especificar los objetivos de LoRA para seleccionar estas cadenas. La verificación se realiza si contiene la subcadena especificada en su nombre completo. Por lo tanto, escribir query y value es equivalente a nuestra implementación desde cero anterior. Para las capas densas, debemos tener un poco más de cuidado, ya que el clasificador también tiene una salida densa. Si deseamos ajustar los demás capas densas, debemos ser más específicos a través de intermediate.dense y output.dense.

Todos los parámetros que no se inyectaron con parámetros LoRA se congelan automáticamente, es decir, no recibirán actualizaciones de gradiente. Si hay alguna capa que deseamos entrenar en su forma original, podemos especificarlas pasando una lista al parámetro modules_to_save de la configuración de Lora-Config. En nuestro caso, queremos agregar LayerNorm aquí y afinar las cabezas para GLUE y SQuAD. Tenga en cuenta que no todos los elementos de las listas tienen que coincidir con algo. Simplemente podemos agregar el classifier y qa_outputs a esta lista y luego tener un solo archivo de configuración que funcionará correctamente para ambas tareas.

Para los parámetros de sesgo, puede utilizar el parámetro de configuración conveniente bias. Puede especificar ya sea todos para volver a entrenar todos los sesgos de todos los módulos, lora_only para entrenar solo los inyectados o ninguno para mantener todos los sesgos constantes durante el entrenamiento.

El siguiente ejemplo inyecta una LoRA con rango 2. Especificamos los parámetros alpha con el valor 8 anterior, ya que este fue el rango que probamos primero y debería permitirnos mantener la tasa de aprendizaje original de nuestro ejemplo desde cero.

import peft# Configuración para la inyección de LoRA a través de PEFTpeft_config = peft.LoraConfig(    r=2, # dimensión del rango de las matrices inyectadas de LoRA    lora_alpha=8, # parámetro de escala, use 8 aquí para que sea comparable con nuestra propia implementación    target_modules=['query', 'key', 'value', 'intermediate.dense', 'output.dense'], # sé preciso sobre dense porque el clasificador también tiene dense    modules_to_save=["LayerNorm", "classifier", "qa_outputs"], # vuelva a entrenar la capa de normalización; classifier es la cabeza de afinación; qa_outputs es para SQuAD    lora_dropout=0.1, # probabilidad de deserción para las capas    bias="all", # ninguno, todos o lora_only)model = peft.get_peft_model(model, peft_config)

Recuerda que especificar más módulos para inyecciones de LoRA podría aumentar los requisitos de VRAM. Si encuentras limitaciones de VRAM, considera reducir el número de módulos objetivo o el rango de LoRA.

Para el entrenamiento, especialmente con QLoRA, elige un optimizador que sea compatible con matrices cuantizadas. Reemplaza tu optimizador estándar de torch por una variante de bitsandbytes de la siguiente manera:

import torchimport bitsandbytes as bnb# reemplaza estooptimizer = torch.optim.AdamW(argumentos aquí)# con estooptimizer = bnb.optim.AdamW8bit(argumentos iguales aquí)

Luego puedes entrenar este modelo como antes, sin tener que preocuparte explícitamente por QLoRA durante el entrenamiento.

Una vez que se haya completado el entrenamiento, el proceso para guardar y volver a cargar tu modelo es sencillo. Utiliza model.save_pretrained para guardar tu modelo, especificando el nombre de archivo deseado. La biblioteca PEFT creará automáticamente un directorio en esta ubicación, donde almacenará los pesos del modelo y un archivo de configuración. Este archivo incluye detalles esenciales como el modelo base y los parámetros de configuración de LoRA.

Para volver a cargar el modelo, utiliza peft.AutoPeftModel.from_pretrained, pasando la ruta del directorio como argumento. Un punto crucial para recordar es que la configuración de LoRA actualmente no conserva la cantidad de clases para las que se inicializó AutoModelForSequenceClassification. Al usar from_pretrained, debes ingresar manualmente este número de clase como un parámetro adicional. No hacerlo resultará en un error.

El modelo recargado comprenderá el modelo base original con los adaptadores LoRA aplicados. Si decides integrar permanentemente los adaptadores LoRA en las matrices del modelo base, simplemente ejecuta model.merge_and_unload().

Para obtener una comprensión más práctica y obtener instrucciones detalladas, echa un vistazo al repositorio de GitHub. Allí encontrarás dos notebooks titulados Train-QLoRA-with-PEFT.ipynb y Load-LoRA-Weights-PEFT.ipynb, que proporcionan un ejemplo paso a paso para entrenar y cargar modelos con PEFT.

Conclusión

“No dejaremos de explorar, y el final de toda nuestra exploración será llegar donde comenzamos y conocer el lugar por primera vez”.

— de “Little Gidding” por T.S. Eliot

Este viaje nos ha llevado desde una implementación sencilla, aunque con código duro, de LoRA a una comprensión más profunda de los adaptadores de rango bajo, su implementación práctica y las pruebas de referencia.

Exploramos una estrategia alternativa y más eficiente de implementación y nos adentramos en la elegancia de bibliotecas existentes como PEFT para la integración de LoRA.

Nuestra aventura concluye con pautas prácticas para emplear LoRA, encapsuladas en los ‘Cinco Mandamientos’, asegurando un uso eficiente y efectivo de esta técnica en aplicaciones del mundo real y una guía paso a paso sobre cómo implementarlos en la práctica.

Referencias

Todas las imágenes, a menos que se indique lo contrario, son del autor.

• Artículo original de LoRA: https://arxiv.org/pdf/2106.09685.pdf
• Artículo de QLoRA: https://arxiv.org/abs/2305.14314
• Guía de Sentdex sobre afinamiento de QLoRA: https://www.youtube.com/watch?v=J_3hDqSvpmg
• Entrada de blog sobre el afinamiento de LoRA en Llama: https://www.anyscale.com/blog/fine-tuning-llms-lora-or-full-parameter-an-in-depth-analysis-with-llama-2
• Integración de bitsandbytes con Hugging Face: https://huggingface.co/blog/4bit-transformers-bitsandbytes
• Información sobre el entrenamiento de LoRA: https://lightning.ai/pages/community/lora-insights/
• Ahorros esperados de VRAM LoRA vs QLoRA al afinar un modelo de Llama: https://cloud.google.com/vertex-ai/docs/model-garden/lora-qlora
• La fuente que utilicé para el texto en la losa de piedra, en caso de que quieras crear la tuya propia: https://www.fontspace.com/sharp-objects-nbp-font-f14469

We will continue to update Zepes; if you have any questions or suggestions, please contact us!