Implementar LoRA desde cero: eficiencia en el fine-tuning de LLM

(lightning.ai)

1 puntos por GN⁺ 2024-01-23 | 1 comentarios | Compartir por WhatsApp

LoRA (Low-Rank Adaptation) es una técnica que reduce el costo del fine-tuning al actualizar solo pequeñas matrices de bajo rango, sin volver a entrenar todo el LLM, y este Studio implementa directamente las capas LoRA para comprobar cómo funcionan
La idea clave es aproximar el cambio de pesos ΔW del fine-tuning tradicional como el producto de dos matrices pequeñas A y B; mientras más pequeño sea el rango r, menor será tanto el número de parámetros entrenables como la capacidad de representación
Una matriz de pesos de 5,000×10,000 tiene 50 millones de parámetros, pero con LoRA r=8 solo se agregan B de 5,000×8 y A de 8×10,000, para un total de 120 mil parámetros, unas 400 veces menos
En clasificación de sentimiento sobre IMDb con DistilBERT, la configuración base de LoRA logró 89.44% de Test acc, por encima del 86.22% de entrenar solo las dos últimas capas, aunque por debajo del 92.31% del fine-tuning completo
Tras explorar hiperparámetros, LoRA logró Val acc de 92.96% y Test acc de 92.39% con cerca de 500 mil parámetros entrenables, superando ligeramente al fine-tuning completo que entrenó 66,955,010 parámetros

Cómo LoRA reduce el costo del fine-tuning

LoRA significa Low-Rank Adaptation y es una técnica para hacer más eficiente el fine-tuning de LLM
En el fine-tuning convencional se ajustan todos los parámetros del modelo de deep learning, pero LoRA actualiza solo un pequeño conjunto de matrices de bajo rango
Un LLM preentrenado puede usarse para muchas tareas, pero el fine-tuning sigue siendo útil para adaptarlo a un dataset o tarea específica
A medida que el modelo crece, actualizar todas las capas implica un costo computacional cada vez mayor

Aproximar ΔW como el producto de matrices pequeñas

En el fine-tuning tradicional, la actualización de la matriz de pesos W se calcula como ΔW
LoRA aproxima ΔW como el producto de dos matrices pequeñas A y B
- Si estás familiarizado con PCA o SVD, puedes verlo como una descomposición de ΔW en A y B
El rango r es un hiperparámetro de LoRA
- Un r más pequeño reduce la cantidad de parámetros entrenables, acelera el entrenamiento y baja los requisitos de cómputo
- Pero al mismo tiempo también reduce la capacidad de las matrices de bajo rango para capturar información específica de la tarea
Ejemplo con una matriz de pesos de 5,000×10,000:
- Actualización normal ΔW: 50 millones de parámetros en total
- LoRA con r=8: B de 5,000×8, A de 8×10,000
- Parámetros adicionales: 80,000 + 40,000 = 120,000
- Es 400 veces más pequeño que el fine-tuning tradicional
En uso real, conviene probar varios valores de r para encontrar el equilibrio entre rendimiento y costo

Implementar una capa LoRA con PyTorch

La LoRALayer básica recibe la dimensión de entrada, la dimensión de salida, el rango y el factor de escalado alpha

class LoRALayer(torch.nn.Module):
    def __init__(self, in_dim, out_dim, rank, alpha):
        super().__init__()
        std_dev = 1 / torch.sqrt(torch.tensor(rank).float())
        self.A = torch.nn.Parameter(torch.randn(in_dim, rank) * std_dev)
        self.B = torch.nn.Parameter(torch.zeros(rank, out_dim))
        self.alpha = alpha

    def forward(self, x):
        x = self.alpha * (x @ self.A @ self.B)
        return x

in_dim es la dimensión de entrada de la capa a la que se aplicará LoRA, y out_dim es la dimensión de salida
rank controla la complejidad de las matrices A y B, así como la cantidad de parámetros que agrega LoRA
alpha determina la magnitud del cambio que LoRA introduce sobre los pesos del modelo original
- Si alpha es alto, el comportamiento del modelo se ajusta más
- Si alpha es bajo, el cambio es más sutil
A se inicializa con pequeños valores aleatorios y la desviación estándar se define con la raíz cuadrada del rango
- Esto ayuda a evitar que los valores iniciales de A sean demasiado grandes
B se inicializa en 0
- Antes de empezar el entrenamiento, como B=0, entonces AB=0
- Hasta que A y B se actualicen por backpropagation, LoRALayer no afecta los pesos originales

Reemplazar una capa Linear por LinearWithLoRA

LoRA suele aplicarse a las capas Linear/feedforward de una red neuronal
Si el forward original llama dos capas Linear en secuencia, al aplicar LoRA se suma la salida de LoRA a la salida de cada Linear

def forward(self, x):
    x = self.linear_1(x) + self.lora_1(x)
    x = F.relu(x)
    x = self.linear_2(x) + self.lora_2(x)
    return logits

Al modificar un modelo existente en PyTorch, una forma simple es reemplazar cada capa Linear por LinearWithLoRA

class LinearWithLoRA(torch.nn.Module):
    def __init__(self, linear, rank, alpha):
        super().__init__()
        self.linear = linear
        self.lora = LoRALayer(
            linear.in_features, linear.out_features, rank, alpha
        )

    def forward(self, x):
        return self.linear(x) + self.lora(x)

LinearWithLoRA conserva tanto la capa Linear original como la nueva LoRALayer
Si reemplazas las capas Linear de un modelo preentrenado por LinearWithLoRA, puedes montarle LoRA y luego hacer fine-tuning

Experimento de clasificación IMDb con DistilBERT

El ejemplo práctico usa clasificación de texto, ya que evaluar exactitud es más sencillo que en generación de texto
El modelo usa DistilBERT preentrenado de Hugging Face transformers

from transformers import AutoModelForSequenceClassification

model = AutoModelForSequenceClassification.from_pretrained(
    "distilbert-base-uncased", num_labels=2)

Para entrenar solo los nuevos pesos LoRA, se establece requires_grad=False en todos los parámetros del modelo

for param in model.parameters():
    param.requires_grad = False

DistilBERT tiene 6 capas Transformer, y dentro de cada una hay capas Linear
- En attention están q_lin, k_lin, v_lin, out_lin
- En la FFN están lin1, lin2
- En la salida están las dos capas Linear pre_classifier y classifier

Configuración para aplicar LoRA de forma selectiva

La configuración base de LoRA lo aplica solo a las matrices de pesos de query y value en attention

lora_r = 8
lora_alpha = 16
lora_dropout = 0.05
lora_query = True
lora_key = False
lora_value = True
lora_projection = False
lora_mlp = False
lora_head = False

Recorriendo en un loop cada capa Transformer de DistilBERT, se reemplazan las capas Linear seleccionadas por LinearWithLoRA
- Si lora_query=True, se reemplaza q_lin
- Si lora_key=True, se reemplaza k_lin
- Si lora_value=True, se reemplaza v_lin
- Si lora_projection=True, se reemplaza out_lin
- Si lora_mlp=True, se reemplazan ffn.lin1 y ffn.lin2
- Si lora_head=True, se reemplazan pre_classifier y classifier
Después del reemplazo, en la salida del modelo se puede verificar que q_lin, v_lin, etc. pasaron a ser LinearWithLoRA

Comparación entre LoRA base y fine-tuning tradicional

Resultado al entrenar la clasificación de IMDb Movie Reviews con la configuración base de LoRA:
- Train acc: 92.15%
- Val acc: 89.98%
- Test acc: 89.44%
Resultado al hacer fine-tuning solo de las dos últimas capas de salida:
- Train acc: 86.68%
- Val acc: 87.26%
- Test acc: 86.22%
- Parámetros entrenables: 592,130
LoRA base obtuvo mejor Test acc que entrenar solo las dos últimas capas, y además usó menos parámetros entrenables: 147,456
Resultado del fine-tuning tradicional de todas las capas:
- Train acc: 96.41%
- Val acc: 92.80%
- Test acc: 92.31%
- Parámetros entrenables: 66,955,010
El fine-tuning completo logra un Test acc alrededor de 2% mayor que LoRA base, pero actualiza cerca de 450 veces más parámetros

Búsqueda de hiperparámetros para LoRA

El rendimiento de LoRA puede variar según lora_r, lora_alpha y la configuración de capas a las que se aplica
03_finetune-lora.py recibe los hiperparámetros como argumentos por línea de comandos

python 03_finetune-lora.py --lora_alpha 32 --lora_r 16

También se pueden activar otros objetivos de aplicación de LoRA

python 03_finetune-lora.py \
--lora_alpha 32 \
--lora_r 16 \
--lora_query True \
--lora_key True \
--lora_value True \
--lora_projection True \
--lora_mlp True \
--lora_head True

03_gridsearch.py ejecuta la siguiente grilla en todas las GPU disponibles
- alpha_values = [1, 4, 8, 16, 32, 64]
- rank_values = [1, 2, 4, 8, 16, 32]
- lora_query = ["True"]
- lora_key = ["False", "True"]
- lora_value = ["True"]
- lora_projection = ["False", "True"]
- lora_mlp = ["False", "True"]
- lora_head = ["False", "True"]
El script puede ejecutarse desde Visual Studio Code, la terminal de línea de comandos o como Job, y el Job se apaga automáticamente al terminar
Los resultados se guardan en results.txt

Mejor configuración encontrada en la búsqueda en grilla

Según results.txt, la mejor combinación de hiperparámetros fue la siguiente

lora_r: 8
lora_alpha: 1
lora_query: True
lora_key: False
lora_value: True
lora_projection: False
lora_mlp: True
lora_head: False

Resultado de esta configuración:
- Val acc: 92.96%
- Test acc: 92.39%
Esta configuración LoRA usa cerca de 500k parámetros entrenables, muchísimo menos que los 66M del fine-tuning completo
Su exactitud fue ligeramente superior al fine-tuning completo, que obtuvo Val acc de 92.80% y Test acc de 92.31%

Entorno de ejecución y materiales adicionales

Si haces clic en Run en la parte superior del Studio, puedes clonar el entorno con el código incluido
Después de clonar el Studio, puedes ejecutar los archivos de código sin pasos extra de instalación, descarga o configuración
Notebooks y scripts relacionados:
- 00_lora-layer.ipynb: implementación de la capa LoRA
- 01_finetune-last-layers.ipynb: fine-tuning de las últimas capas
- 02_finetune-with-lora.ipynb: fine-tuning con LoRA
- 03_finetune-lora.py: ejecución con hiperparámetros de LoRA por argumentos
- 03_gridsearch.py: búsqueda en grilla de hiperparámetros de LoRA
- 04_finetune-all-layers.ipynb: fine-tuning de todas las capas
Material adicional:
- Parameter-Efficient LLM Finetuning With Low-Rank Adaptation (LoRA): presenta con más detalle los conceptos y el ejemplo práctico de LoRA
- Practical Tips for Finetuning LLMs Using LoRA (Low-Rank Adaptation): ofrece preguntas frecuentes y consejos prácticos sobre fine-tuning con LoRA

1 comentarios

GN⁺ 2024-01-23

Opiniones de Hacker News

El flujo de la técnica sigue el LLMs 101 de Maxime Labonne: https://github.com/mlabonne/llm-course#4-supervised-fine-tun...
LoRA != LoRa, así que me sigo confundiendo. No me gusta que hayan reutilizado una sigla que ya existía
- A mí me pasa igual. Mi trabajo principal es machine learning y aun así —o quizá precisamente por eso— cada vez que veo esta sigla en un lugar con casi nada de contexto, me detengo un momento
  Especialmente en lugares como la portada de HN, donde ambos significados son naturales
- ¿Qué significa además de “Low-Rank Adaptation”? Es difícil hasta buscar la diferencia
- Esto pasa cuando la gente se especializa demasiado y no le importa qué ocurre fuera de su burbuja
- No me gusta la tendencia de la gente de software a tomar nombres relacionados con hardware
- Es una lástima que dos tecnologías sin relación entre sí hayan terminado usando la misma sigla
Todavía me parece raro que en ciencias de la computación se digan cosas como “no sabemos exactamente cómo estos números, es decir, los hiperparámetros, afectan el resultado, así que probemos varios valores y usemos el que mejor funcione”
- ¿“Probar varios valores y usar el que mejor funcione” no es parecido a usar una simulación de Monte Carlo para encontrar valores?
  A veces puedes quedarte atrapado en un máximo local en vez del óptimo/la respuesta correcta, pero aun así funciona. Como no se puede resolver con una fórmula en forma cerrada, se toman muestras aleatorias miles de millones de veces hasta encontrar el valor deseado; no digo que los LLM sean lo mismo, pero este enfoque se usa bastante seguido
- Se siente como la diferencia entre algo diseñado por ingeniería y algo descubierto
  Hasta ahora, la mayor parte de la industria había sido diseñada por ingeniería, y los LLM se parecen más a algo descubierto
- Ese tipo de parcheo de abajo hacia arriba también se parece a la forma en que, como observó el propio Dijkstra, empezó la informática en Estados Unidos: https://www.cs.utexas.edu/users/EWD/transcriptions/EWD06xx/E...
  Idealmente se necesita una base teórica, pero para extraer suficientes datos como para crear o verificar una teoría, a veces hace falta exploración aleatoria
- También pesó mucho que Minsky y otros desacreditaran los perceptrones diciendo que no podían modelar funciones no lineales. Los LLM difícilmente habrían surgido sin las CPU y GPU modernas, pero eso no significa que no pudiéramos haber tenido antes una mejor base teórica
  Estamos varios años por detrás de donde deberíamos estar. Cuando trabajaba en la industria de videojuegos en los 90, era “sentido común” que las redes neuronales eran, en el mejor de los casos, un callejón sin salida y, en el peor, una estafa. Es realmente lamentable que hayamos perdido tanto tiempo porque algunas figuras de autoridad desalentaron a todos, y esta vez tenemos que evitar que eso se repita
- Investigar configuraciones de Stable Diffusion se siente parecido. Lo que uno descubre rápido es que hay muchas conjeturas mezcladas
Todavía no está claro cuándo conviene hacer fine-tuning y cuándo usar RAG
Antes pensaba que el fine-tuning se usaba principalmente para cambiar el comportamiento del modelo, pero últimamente parece que algunas empresas también lo usan para agregar conocimiento. Me pregunto cuáles son los usos principales del fine-tuning
- Creo que el uso principal sigue siendo el cambio de comportamiento. Por ejemplo, fine-tuning de instrucciones, fine-tuning para clasificación, etc.
  Agregar conocimiento a los pesos se hace mejor con preentrenamiento. O, si tienes una base de datos externa o documentos para consultar durante la generación, puedes usar RAG como en preguntas y respuestas. Como referencia, en el NeurIPS 2023 LLM Efficiency Challenge, todos los ganadores que hicieron fine-tuning del “mejor” LLM en 24 horas con una sola GPU usaron LoRA o QLoRA (LoRA cuantizado)
- Si los datos adicionales no son concisos o requieren contexto, el fine-tuning es mejor que RAG
  Si hay demasiado contexto o el foco es difuso, la capacidad de seguir el prompt puede diluirse, y RAG no permite que el modelo aprenda asociaciones de tokens de mayor dimensión. Por eso hay que tener suerte y extraer el contenido necesario del material complementario, lo que no lo hace mucho mejor que un motor de búsqueda avanzado. Esto es especialmente problemático al trabajar con corpus especializados que tienen microdialectos propios y que no aparecen bien en datasets públicos, como documentos internos de gobiernos o grandes empresas
- Según entiendo, el fine-tuning es anormalmente efectivo [0]. Esto se debe a que el aprendizaje en contexto depende mucho de qué tan potente sea el modelo base y de cómo esté construido RAG, es decir, del procesamiento de consultas, la búsqueda por embeddings, el ranking de resultados, etc. [1]
  Según los papers que leí, el fine-tuning puede agregar nuevo conocimiento de dominio o reforzar conocimiento específico, mientras que RAG se limita al refuerzo. Dicho eso, las dos técnicas, con distintos trade-offs, a veces muestran capacidades de nivel similar [2]
  
  [0] Fast.ai: Can Models learn from one sample, https://www.fast.ai/posts/2023-09-04-learning-jumps/ / https://archive.is/eJMPR
  
  [1] LlamaIndex: Advanced RAG, https://blog.llamaindex.ai/a-cheat-sheet-and-some-recipes-fo... / https://archive.is/qtBXX
  
  [2] Microsoft: RAG vs Fine-tuning: Pipelines, Tradeoffs, and a Case Study, https://arxiv.org/html/2401.08406v2#S6 / https://archive.is/UQ8Sa#S6
- Estos son modelos autorregresivos. Si hay un nuevo tipo de secuencia en el que se puede predecir lo que viene después a partir de lo anterior, y esa forma difiere de lo que el modelo vio antes, el fine-tuning parece tener sentido
  Como criterio para decidir qué hacer en una situación de datos específica es bastante ambiguo, pero puede bastar como heurística general. Si agregar conocimiento entra en esto quizá sea cuestión de preferencia sin experimentos
Buen artículo. No soy de este campo, pero cuando leí el paper original entendí que LoRA se aplicaba solo a la última capa densa, y no de forma independiente a todas las capas. Puede que lo haya leído mal.
Estuve investigando un poco por qué la implementación del enlace es así, y parece que QLoRA usó este enfoque y que tiene efectos interesantes. Sería bueno agregar una nota sobre la decisión de QLoRA. Aun así, no tengo muy claro por qué funciona; desde el punto de vista de un principiante, aplicar LoRA a la última capa tiene sentido, pero no logro intuir la justificación para aplicarlo repetidamente a cada capa lineal. ¿Podrías explicar la intuición?
- Como muchas cosas en machine learning, qué capas usar depende más de evidencia empírica que de teoría. En un pipeline típico de entrenamiento con LoRA, se congela el modelo base y solo se ajustan las capas LoRA.
  Mientras más capas reemplaces por capas LoRA, mayor es la libertad de optimización. Algunos enfoques de fine-tuning recomiendan ajustar solo la última capa, porque se supone que contiene la representación de “más alto nivel” de la entrada. Otros enfoques ajustan todas las capas. En gran medida depende de los datos y del problema, y LoRA simplemente refleja esa práctica.
Prefiero el enfoque de Axolotl, que no empieza “desde cero”, sino desde la configuración. Axolotl soporta fine-tuning de Mistral y Llama 2, y también muchas técnicas modernas como sample packing, FlashAttention y xFormers.
En vez de aprender LoRA desde cero, me concentro en reunir y curar datos de fine-tuning para hacer fine-tuning centrado en datos.
- Nuestra librería lit-gpt (de Lightning AI) hace lo mismo.
  https://github.com/Lightning-AI/lit-gpt
Poner nombres sí que es difícil. Al principio pensé que esto hablaba de LoRa de “long range”, o de LoRaWAN para comunicación de sensores IoT.
¿Cuáles son las librerías más usadas para fine-tuning? Me refiero a enfoques que no sean “desde cero”.
- https://github.com/OpenAccess-AI-Collective/axolotl
Wow, yo también al principio di por hecho que hablaban de LoRa.
¿Cuál es el costo de rendimiento de LoRA?
- Durante el entrenamiento es más eficiente que el fine-tuning completo, porque con backpropagation solo se actualiza una parte de los parámetros.
  Durante la inferencia hay dos posibilidades. Si se suman dinámicamente los valores de LoRA durante el forward pass, en teoría podría ser un poco más lento, pero también puede ser una ventaja si quieres mantener conjuntos pequeños de pesos separados para cada cliente. Esto se debe a que puedes ejecutar un solo modelo base grande y aplicar al vuelo los pesos LoRA de cada cliente. Si vuelves a fusionar los pesos LoRA en el modelo base, puedes obtener exactamente el mismo rendimiento que con el modelo base.

Implementar LoRA desde cero: eficiencia en el fine-tuning de LLM

Cómo LoRA reduce el costo del fine-tuning

Aproximar ΔW como el producto de matrices pequeñas

Implementar una capa LoRA con PyTorch

Reemplazar una capa Linear por LinearWithLoRA

Experimento de clasificación IMDb con DistilBERT

Configuración para aplicar LoRA de forma selectiva

Comparación entre LoRA base y fine-tuning tradicional

Búsqueda de hiperparámetros para LoRA

Mejor configuración encontrada en la búsqueda en grilla

Entorno de ejecución y materiales adicionales

Lecturas relacionadas

1 comentarios

Opiniones de Hacker News