Ajuste Fino de Modelos de Lenguaje Grandes (LLMs) con LoRA: Guía Práctica

La era de los Large Language Models (LLMs) ha llegado para quedarse, transformando la forma en que interactuamos con la inteligencia artificial. Sin embargo, entrenar o ajustar modelos tan masivos puede ser un desafío monumental debido a los enormes requisitos computacionales y de memoria. Aquí es donde técnicas como Low-Rank Adaptation (LoRA) entran en juego, permitiéndonos adaptar estos potentes modelos a tareas específicas con una fracción de los recursos.

En este tutorial, exploraremos qué es LoRA, por qué es tan efectivo y cómo puedes implementarlo paso a paso para ajustar tus propios LLMs. Prepárate para potenciar tus proyectos de NLP con esta técnica innovadora.

🎯 ¿Qué es LoRA y por qué es importante?

LoRA, o Low-Rank Adaptation, es una técnica de ajuste fino eficiente de parámetros que permite adaptar grandes modelos pre-entrenados a nuevas tareas con un costo computacional y de almacenamiento significativamente reducido. En lugar de ajustar todos los millones o miles de millones de parámetros del modelo base, LoRA introduce un pequeño conjunto de matrices de bajo rango que se superponen a las matrices de peso originales del modelo.

💡 El problema del Ajuste Fino Tradicional

Tradicionalmente, el ajuste fino de un LLM implicaría actualizar todos los pesos del modelo. Esto tiene varias desventajas:

• Alto costo computacional: Se requiere una gran cantidad de VRAM y tiempo de entrenamiento.
• Gran almacenamiento: Es necesario guardar una copia completa del modelo ajustado para cada tarea.
• Catastrophic Forgetting: Existe el riesgo de que el modelo olvide conocimientos previamente aprendidos sobre la tarea original al especializarse demasiado en la nueva tarea.

✅ La Solución LoRA

LoRA aborda estos problemas de la siguiente manera:

1. Matrices de bajo rango: Se añade un par de matrices (A y B) de rango bajo a las capas de atención del modelo. Estas matrices son mucho más pequeñas que la matriz de pesos original.
2. Entrenamiento solo de LoRA: Durante el ajuste fino, solo se entrenan los pesos de estas matrices A y B. Los pesos originales del modelo base se congelan y no se modifican.
3. Representación eficiente: El cambio total en los pesos de una capa se representa como la multiplicación de estas matrices de bajo rango ($W' = W + BA$).

🛠️ Requisitos Previos

Antes de sumergirnos en el código, asegúrate de tener lo siguiente:

• Python 3.8+
• Pip para la gestión de paquetes
• Acceso a una GPU (recomendado para un rendimiento decente, aunque LoRA reduce la necesidad de VRAM)

Instalación de Librerías

Necesitaremos transformers para los modelos, peft para la implementación de LoRA, accelerate para la optimización del entrenamiento, y datasets para manejar los datos.

pip install transformers datasets peft accelerate evaluate bitsandbytes scipy trl

📖 Entendiendo la Cuantización (4-bit/8-bit)

Para trabajar con LLMs aún más grandes, a menudo es necesario cargarlos en formatos de menor precisión (cuantización). Esto reduce drásticamente el uso de memoria de la GPU, a costa de una ligera pérdida de precisión que a menudo es aceptable.

• 16-bit (bfloat16/float16): Precisión estándar para muchos LLMs. Requiere VRAM considerable.
• 8-bit (int8): Reduce la memoria a la mitad. Útil para modelos grandes.
• 4-bit (nf4): Reduce la memoria a la cuarta parte. Ideal para GPU con VRAM limitada, como las de consumo.

LoRA es especialmente potente cuando se combina con la cuantización, ya que permite ajustar modelos gigantes que de otra manera no cabrían en una única GPU.

📝 Preparación del Conjunto de Datos

Para este tutorial, utilizaremos un conjunto de datos pequeño para una tarea de clasificación o generación simple. Elegiremos el conjunto de datos imdb de Hugging Face datasets y lo adaptaremos para una tarea de sentimiento.

from datasets import load_dataset

# Cargar el dataset IMDB (para demostración, usaremos solo una fracción)
dataset = load_dataset("imdb")

# Tomaremos una muestra más pequeña para agilizar el proceso en el tutorial
def format_data(example):
    # Formatearemos los datos para que el LLM los procese como una conversación
    template = "### Instrucción:\nClasifica el siguiente sentimiento de la reseña como positivo o negativo.\n\n### Reseña:\n{text}\n\n### Sentimiento:\n{label}"
    # Mapear 0 a 'negativo' y 1 a 'positivo'
    label_map = {0: 'negativo', 1: 'positivo'}
    example['text'] = template.format(text=example['text'], label=label_map[example['label']])
    return example

# Aplicar la función de formato solo a una pequeña parte para un ejemplo rápido
train_dataset = dataset['train'].shuffle(seed=42).select(range(1000)).map(format_data)
val_dataset = dataset['test'].shuffle(seed=42).select(range(200)).map(format_data)

print("Ejemplo de dato entrenado:\n", train_dataset[0]['text'])

# El label original ya no es necesario si el LLM generará la respuesta
# Si usaras un modelo de clasificación, mantenerías el label original

El formato que hemos creado es un estilo de prompt instructivo. El LLM aprenderá a completar la parte ### Sentimiento: basándose en la ### Reseña:. Esto es típico para el ajuste fino de LLMs.

⚙️ Configuración del Modelo Base y el Tokenizador

Seleccionaremos un LLM pequeño y eficiente de Hugging Face. Para este tutorial, usaremos NousResearch/Llama-2-7b-chat-hf (una versión de Llama 2). Puedes elegir otros como mistralai/Mistral-7B-v0.1.

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

model_name = "NousResearch/Llama-2-7b-chat-hf" # O "mistralai/Mistral-7B-v0.1"

# Configuración de cuantización (4-bit)
nfq4_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4", # normal_float 4-bit
    bnb_4bit_use_double_quant=True, # doble cuantización para más ahorro de memoria
    bnb_4bit_compute_dtype=torch.bfloat16 # usar bfloat16 para cómputos si la GPU lo soporta
)

# Cargar el modelo base y el tokenizador
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    quantization_config=nfq4_config,
    torch_dtype=torch.bfloat16, # Asegúrate de que el modelo se cargue con el dtype correcto
    device_map="auto" # Distribuye el modelo automáticamente en las GPUs disponibles
)
model.config.use_cache = False # Deshabilitar cache para fine-tuning
model.config.pretraining_tp = 1 # Recomendado para Llama 2

tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
tokenizer.pad_token = tokenizer.eos_token # Usa EOS como padding token
tokenizer.padding_side = "right" # Importante para modelos generativos

print("Modelo cargado y cuantizado con éxito.")

🚀 Configuración de LoRA con PEFT

Ahora, configuraremos los parámetros de LoRA usando la librería PEFT (Parameter-Efficient Fine-tuning).

from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training

# 1. Preparar el modelo para entrenamiento k-bit (cuantizado)
model = prepare_model_for_kbit_training(model)

# 2. Configuración de LoRA
lora_config = LoraConfig(
    lora_alpha=16, # Escala el efecto de los nuevos pesos. Un valor común es 16 o 32.
    lora_dropout=0.1, # Dropout aplicado a las matrices LoRA para regularización.
    r=8, # El "rango" de las matrices de bajo rango. Controla la expresividad de LoRA.
    bias="none", # Tipo de ajuste de bias. 'none' es común.
    task_type="CAUSAL_LM", # Tipo de tarea: Causal Language Modeling (generación de texto).
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
    # Estos son los módulos (capas) del modelo base a los que se aplicará LoRA. 
    # Normalmente se dirigen a las capas de atención (q, k, v, o) y MLP (gate, up, down) para Llama.
)

# 3. Obtener el modelo PEFT (Parameter-Efficient Fine-Tuning)
model = get_peft_model(model, lora_config)

print("\nModelo con LoRA configurado. Parámetros entrenables:")
model.print_trainable_parameters()

Verás una salida similar a:

trainable params: 4,194,304 || all params: 3,504,186,368 || trainable%: 0.11969411986423011

Esto significa que estamos entrenando solo unos 4 millones de parámetros (¡menos del 0.2% del total!), en lugar de los miles de millones del modelo base. ¡Una eficiencia increíble!

🏃‍♂️ Entrenamiento del Modelo con Trainer

Ahora utilizaremos la clase Trainer de transformers para manejar el bucle de entrenamiento. trl proporciona SFTTrainer que simplifica el fine-tuning supervisado para LLMs.

from transformers import TrainingArguments
from trl import SFTTrainer

# Configuración de los argumentos de entrenamiento
training_arguments = TrainingArguments(
    output_dir="./results", # Directorio para guardar los checkpoints y logs
    num_train_epochs=1, # Para este tutorial, 1 época es suficiente
    per_device_train_batch_size=4, # Tamaño del batch por GPU
    gradient_accumulation_steps=2, # Acumulación de gradientes para simular un batch más grande
    optim="paged_adamw_8bit", # Optimizador optimizado para entrenamiento 8-bit
    save_steps=100, # Guardar checkpoint cada 100 pasos
    logging_steps=10, # Loguear métricas cada 10 pasos
    learning_rate=2e-4, # Tasa de aprendizaje
    weight_decay=0.001, # Regularización
    fp16=False, # Ya estamos en bfloat16, no necesitamos fp16
    bf16=True, # Usar bfloat16 para entrenamiento
    max_grad_norm=0.3, # Clipping de gradientes
    max_steps=-1, # -1 para entrenar por num_train_epochs
    warmup_ratio=0.03, # Warmup inicial de la tasa de aprendizaje
    group_by_length=True, # Agrupa ejemplos de longitud similar para eficiencia
    lr_scheduler_type="constant", # Scheduler de tasa de aprendizaje
    report_to="none" # No reportar a ninguna plataforma de logging por ahora
)

# Inicializar SFTTrainer
trainer = SFTTrainer(
    model=model,
    train_dataset=train_dataset,
    eval_dataset=val_dataset,
    peft_config=lora_config,
    dataset_text_field="text", # Campo del dataset que contiene el texto de entrenamiento
    tokenizer=tokenizer,
    args=training_arguments,
    max_seq_length=512, # Longitud máxima de secuencia para el tokenizador
    packing=False, # No empaquetar múltiples ejemplos en una sola secuencia
)

# Entrenar el modelo
trainer.train()

print("\nEntrenamiento completado!")

Este proceso puede tardar un tiempo dependiendo de tu GPU y el tamaño del dataset. Con el pequeño dataset de ejemplo, debería ser relativamente rápido.

💾 Guardar y Cargar el Modelo Ajustado

Una de las ventajas de LoRA es que solo necesitas guardar los pesos de LoRA, no el modelo base completo. Esto ahorra mucho espacio.

# Guardar los pesos de LoRA
output_dir = "./llama2-7b-chat-sentiment-lora"
trainer.model.save_pretrained(output_dir)
tokenizer.save_pretrained(output_dir)

print(f"\nModelo LoRA guardado en {output_dir}")

Para cargar el modelo ajustado y usarlo para inferencia, primero cargas el modelo base (cuantizado) y luego le aplicas los pesos de LoRA.

from peft import PeftModel, PeftConfig

# Cargar la configuración de LoRA guardada
config = PeftConfig.from_pretrained(output_dir)

# Cargar el modelo base original (cuantizado)
base_model = AutoModelForCausalLM.from_pretrained(
    config.base_model_name_or_path,
    quantization_config=nfq4_config,
    torch_dtype=torch.bfloat16,
    device_map="auto"
)

# Cargar los pesos de LoRA en el modelo base
model_for_inference = PeftModel.from_pretrained(base_model, output_dir)
model_for_inference.eval() # Poner el modelo en modo evaluación

tokenizer_for_inference = AutoTokenizer.from_pretrained(output_dir)

print("Modelo ajustado cargado para inferencia.")

🧪 Realizar Inferencias y Evaluar

Vamos a probar nuestro modelo con algunas reseñas de ejemplo.

from transformers import GenerationConfig

def generate_sentiment(review_text):
    prompt = f"### Instrucción:\nClasifica el siguiente sentimiento de la reseña como positivo o negativo.\n\n### Reseña:\n{review_text}\n\n### Sentimiento:\n"
    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")

    # Configuración de generación, importante para LLMs
    generation_config = GenerationConfig(
        max_new_tokens=10, # Generar hasta 10 tokens nuevos (suficiente para 'positivo' o 'negativo')
        do_sample=True, # Habilitar muestreo
        temperature=0.7, # Controlar la aleatoriedad
        top_p=0.0, # Deshabilitar top_p sampling (para resultados más deterministas)
        pad_token_id=tokenizer.pad_token_id, # Asegúrate de que el pad token esté configurado
        eos_token_id=tokenizer.eos_token_id # Asegúrate de que el EOS token esté configurado
    )

    with torch.no_grad():
        outputs = model_for_inference.generate(**inputs, generation_config=generation_config)
    
    # Decodificar y extraer solo la parte generada
    generated_text = tokenizer.decode(outputs[0][len(inputs.input_ids[0]):], skip_special_tokens=True)
    return generated_text.strip().split('\n')[0] # Extraer solo la primera línea de la respuesta

# Ejemplos de prueba
review1 = "¡Esta película es increíble, me encantó cada segundo!"
review2 = "Qué pérdida de tiempo, no la recomiendo en absoluto."
review3 = "Estuvo bien, pero esperaba más de la trama."

print(f"Reseña 1: '{review1}' -> Sentimiento: {generate_sentiment(review1)}")
print(f"Reseña 2: '{review2}' -> Sentimiento: {generate_sentiment(review2)}")
print(f"Reseña 3: '{review3}' -> Sentimiento: {generate_sentiment(review3)}")

Los resultados pueden variar ligeramente debido a la aleatoriedad del muestreo y el tamaño limitado del dataset de entrenamiento. Con un dataset más grande y más épocas, el rendimiento mejoraría sustancialmente.

📈 Ventajas y Desventajas de LoRA

LoRA no es una bala de plata, pero ofrece beneficios significativos en la mayoría de los escenarios de ajuste fino de LLMs.

📚 Más Allá de LoRA: Otras Técnicas PEFT

LoRA es solo una de las muchas técnicas de Parameter-Efficient Fine-Tuning (PEFT). Otras incluyen:

• Prefix-Tuning: Añade una secuencia de tokens especiales (prefijo) al inicio de la entrada que se entrena para la tarea específica.
• P-Tuning: Similar a Prefix-Tuning, pero entrena prompts continuos en lugar de tokens discretos.
• Prompt-Tuning: Entrena un conjunto de tokens virtuales que se concatena con el prompt de entrada.

Estas técnicas tienen sus propias ventajas y desventajas, y la elección depende de la tarea y los recursos disponibles. LoRA se ha demostrado como una de las más robustas y eficaces en muchas aplicaciones.

Conclusion

El ajuste fino de LLMs con LoRA es una habilidad fundamental para cualquier ingeniero de Machine Learning o científico de datos que trabaje con modelos de lenguaje. Te permite aprovechar el poder de los modelos más grandes sin necesidad de un clúster de GPUs de vanguardia. Al dominar LoRA, no solo ahorras recursos, sino que también puedes adaptar rápidamente los LLMs a una multitud de tareas empresariales y de investigación, abriendo un mundo de posibilidades.

Esperamos que esta guía te haya proporcionado una comprensión sólida y las herramientas prácticas para comenzar tu propio viaje con LoRA. ¡Feliz ajuste fino!