Cuello de botella común desde F16 hasta F23 — Coste de ajuste fino

Desde DNABERT en F16 hasta UCE en F23, todos los modelos presentados en esta serie tienen una escala de millones a miles de millones de parámetros. Para realizar el ajuste fino (fine-tuning) de estos modelos en un laboratorio para una tarea específica (por ejemplo, la clasificación de tipos celulares de un cáncer determinado), en principio, se deben cargar en la memoria de la GPU los gradientes y los estados del optimizador para todos los parámetros del modelo. Esto es difícil de gestionar con los recursos de GPU que poseen la mayoría de los laboratorios (una sola T4, A10, etc.). En F24, abordaremos la técnica estándar para eludir este cuello de botella: LoRA (Low-Rank Adaptation).

Principio — Añadir solo una corrección de bajo rango sin cambiar todo

Coste de memoria del ajuste fino completo

Supongamos que realizamos un ajuste fino completo (full fine-tuning) de un modelo con $\theta$ parámetros utilizando el optimizador Adam convencional. Si sumamos los pesos del modelo, los gradientes y los momentos de 1.er y 2.º orden de Adam, el límite inferior de almacenamiento es aproximadamente

$\text{Límite inferior de almacenamiento básico} \approx \theta \times (\underbrace{2}_{\text{pesos(fp16)}} + \underbrace{2}_{\text{gradientes}} + \underbrace{8}_{\text{momentos de Adam, fp32}}) \text{ bytes} = 12\theta \text{ bytes}$

Para un modelo de 1.000 millones de parámetros, este límite inferior es aproximadamente $12 \times 10^9$ bytes $\approx 12$ GB. En el entrenamiento real, se añaden las activaciones, los buffers temporales y, dependiendo de la implementación, los pesos maestros en fp32, por lo que es difícil de gestionar con los 16 GB de VRAM de una T4.

LoRA — Aproximar el cambio de peso mediante bajo rango

En lugar de actualizar por completo la matriz de pesos original $W \in \mathbb{R}^{d \times k}$ durante el proceso de ajuste fino, LoRA aproxima ese cambio $\Delta W$ como el producto de dos matrices mucho más pequeñas.

$W' = W + \Delta W \approx W + BA, \quad B \in \mathbb{R}^{d \times r},\ A \in \mathbb{R}^{r \times k},\ r \ll \min(d, k)$

Aquí, $r$ es el rango (rank), que suele ser un valor pequeño entre 4 y 64. Los pesos originales $W$ permanecen congelados (freeze), y solo se entrenan $A$ y $B$.

Ejemplo de cálculo manual: Ratio de reducción de parámetros

Supongamos que ajustamos con LoRA una matriz de pesos cuadrada con $d = k = 4096$. El número de parámetros originales sujetos al ajuste fino completo es

$d \times k = 4096 \times 4096 = 16{,}777{,}216$

Si aplicamos un LoRA con rango $r=8$, el número de parámetros sujetos al entrenamiento es

$d \times r + r \times k = 4096 \times 8 + 8 \times 4096 = 65{,}536$

La proporción de ahorro es

$\frac{16{,}777{,}216}{65{,}536} = 256\text{ veces}$

Es decir, si consideramos solo esta capa, los parámetros que deben aprenderse se reducen a aproximadamente el 0,4 % del original. Aunque se extienda al modelo completo, esta proporción se mantiene aproximadamente; este es el fundamento por el cual LoRA reduce drásticamente la memoria de GPU y el tiempo de entrenamiento.

¿Por qué funciona la aproximación de bajo rango?

Los modelos grandes preentrenados ya han aprendido representaciones generales a partir de datos masivos. La evidencia empírica de LoRA indica que la "corrección" necesaria para adaptar un modelo a una tarea específica de downstream suele ocurrir dentro de un subespacio de dimensión mucho menor (dimensión intrínseca), sin necesidad de explorar nuevamente todo el espacio de pesos original.

$\text{intrinsic dimension of task adaptation} \ll \text{full parameter space}$

Práctica: Ajuste fino LoRA de modelos tipo Nucleotide Transformer con HuggingFace PEFT

# Escala ejecutable en el nivel gratuito de Colab T4 (para un modelo pequeño de variantes).

from transformers import AutoModel, AutoTokenizer
from peft import LoraConfig, get_peft_model, TaskType

# Cargar el modelo preentrenado (ejemplo de ruta tratado en F17)
# base_model = AutoModel.from_pretrained("InstaDeepAI/nucleotide-transformer-v2-50m-multi-species")
# tokenizer = AutoTokenizer.from_pretrained("InstaDeepAI/nucleotide-transformer-v2-50m-multi-species")

lora_config = LoraConfig(
    task_type=TaskType.FEATURE_EXTRACTION,
    r=8,                       # Rango: ajustar entre 4 y 64 según la dificultad de la tarea
    lora_alpha=16,             # Factor de escala, normalmente alrededor de 2 veces r
    lora_dropout=0.1,
    # Los nombres reales de los módulos varían entre modelos. Comprobarlos con named_modules() antes de especificarlos.
    target_modules=["query", "value"],
)

# peft_model = get_peft_model(base_model, lora_config)
# peft_model.print_trainable_parameters()
# Ejemplo de salida: "trainable params: 294,912 || all params: 50,219,904 || trainable%: 0.587%"

print("Ejemplo de configuración de LoRA:", lora_config)

Limitar las matrices de proyección Query y Value de target_modules como atención es una concesión común, pero los nombres (query, value) varían según la arquitectura del modelo. Primero, verifique el nombre real de la capa lineal con base_model.named_modules() y compruebe la inserción del adaptador y el paso forward con un lote pequeño. Aplicar LoRA a todas las capas lineales podría mejorar ligeramente el rendimiento, pero reduciría el margen de ahorro de parámetros; por lo tanto, ajustar el alcance de aplicación en función de la memoria GPU disponible y el rendimiento objetivo es una práctica profesional.

Mapeo CS

• Factorización de matrices de bajo rango (low-rank matrix factorization): El núcleo de LoRA es la aproximación de bajo rango del álgebra lineal. Aproximar una matriz grande mediante el producto de dos matrices más pequeñas utiliza las mismas herramientas matemáticas que la factorización de matrices en sistemas de recomendación o la proyección de baja dimensión en PCA.
• Aprendizaje eficiente en parámetros (Parameter-Efficient Fine-Tuning, PEFT): LoRA es un ejemplo destacado dentro del amplio grupo de técnicas PEFT. Otras técnicas PEFT, como la inserción de capas adaptadoras (adapters) o el ajuste por prompts, comparten la misma filosofía: "mantener fijo el modelo original y aprender solo módulos adicionales pequeños".
• Despliegue modular (modular deployment): Los pesos de LoRA aprendidos pueden almacenarse y desplegarse por separado del modelo original, lo que permite intercambiar adaptadores específicos para diferentes tareas sobre un único modelo base, similar a una arquitectura de plugins en ingeniería de software.

Defectos comunes

• Establecer el rango $r$ arbitrariamente pequeño: Si se fija $r$ demasiado bajo, puede faltar la capacidad de representación necesaria para la tarea. Es más seguro comparar el rendimiento de validación con varios valores de $r$.
• Reducir arbitrariamente target_modules y atribuir erróneamente la caída de rendimiento a las limitaciones de LoRA: Si se restringe demasiado el alcance de aplicación, no se obtendrá la adaptabilidad necesaria. Si el rendimiento no cumple las expectativas, se debe revisar conjuntamente el rango y el alcance de los módulos aplicados.
• Confundir cuantización (quantization) con LoRA: Aunque existen técnicas que combinan cuantización con LoRA, como QLoRA, LoRA en sí mismo no es una técnica para cambiar la precisión de los pesos, sino para reducir la cantidad de parámetros aprendidos. Es importante distinguirlos claramente.

Para profundizar más

El texto ha sido reescrito directamente por el equipo de investigación BPD. Profundice consultando el artículo original y los materiales oficiales.

• Artículo original de LoRA: Hu et al. (2021), LoRA: Low-Rank Adaptation of Large Language Models, preimpresión en arXiv (publicado en ICLR 2022).
• Artículo original de QLoRA: Dettmers et al. (2023), QLoRA: Efficient Finetuning of Quantized LLMs, NeurIPS.
• Documentación oficial de PEFT de HuggingFace: Métodos de uso de la biblioteca y lista de técnicas compatibles.

En el siguiente capítulo (F25), cerraremos la capa F1 de modelos fundacionales de esta serie con un recorrido práctico donde exploramos, seleccionamos y comparamos los modelos tratados en toda la serie (F01～F24) en el Hub de modelos Bio de HuggingFace.