LoRA para adaptación eficiente

LoRA (Low-Rank Adaptation) es una técnica ideada en 2021 por un equipo de investigadores de Microsoft que reduce drásticamente el número de parámetros entrenables en modelos de lenguaje de gran tamaño.

En el ajuste fino tradicional, el modelo ajusta una matriz de pesos para obtener la matriz adaptada. LoRA congela el modelo base y calcula una matriz de actualización de bajo rango.

La fórmula de predicción se expresa matemáticamente de la siguiente manera:

• x: Vector de entrada de la capa.
• : Matriz de pesos original del modelo preentrenado (que se congela y no se actualiza).
• : Matriz de actualización de los pesos.
• y : Matrices de bajo rango entrenables.

La matriz A se inicializa con valores aleatorios, la matriz B se inicializa en 0 y el número de parámetros entrenables sería r(d+k) en lugar de dk.

Vamos a ver un ejemplo con opt-125m, un modelo decoder con 125 millones de parámetros, el más pequeño de la familia Open Pre-trained Transformer (OPT). Hemos importado los módulos de transformers y peft y hemos definido una función para describir el modelo y otra para imprimir los parámetros entrenables.

Primero, cargamos el modelo:

# 1. Cargar modelo base pequeño (OPT-125M)
model_id = "facebook/opt-125m"

print(f"Loading base model and tokenizer: {model_id}...")
tokenizer = AutoTokenizer.from_pretrained(model_id)
base_model = AutoModelForCausalLM.from_pretrained(
        model_id,
        torch_dtype=torch.float32,
        device_map="auto"
    )

Ahora, mostramos los componentes del modelo, configuramos LoRA, lo aplicamos al modelo base, imprimimos los parámetros entrenables y simulamos un forward y backward. En la configuración de LoRA indicamos que es un modelo generativo, ponemos un rango r de 8, un factor de escala 2r, que se aplique a las matrices de consulta (query) y valor (value) y que haya dropout del 0.05.

# Mostrar componentes usando el recorrido recursivo corregido
describe_model_components(base_model)

# 2. Configurar PEFT / LoRA
# r=8 reduce drásticamente el tamaño frente al espacio original de 768x768
peft_config = LoraConfig(
        task_type=TaskType.CAUSAL_LM,
        r=8,
        lora_alpha=16,
        target_modules=["q_proj", "v_proj"],
        lora_dropout=0.05,
        bias="none"
    )

# 3. Aplicar PEFT al modelo base
print("Applying LoRA configuration to base model...")
peft_model = get_peft_model(base_model, peft_config)

# 4. Verificar parámetros congelados vs actualizados
print_trainable_parameters(peft_model)

# 5. Simulación de pasada forward y backward para verificar gradientes
print("Simulating a dummy forward and backward pass...")
inputs = tokenizer("Fine-tuning models with PEFT is", return_tensors="pt").to(peft_model.device)
labels = inputs["input_ids"]

outputs = peft_model(**inputs, labels=labels)
loss = outputs.loss

loss.backward()
print(f"Loss calculated: {loss.item():.4f}")
print("Gradients successfully generated for updated parameters.")

Vemos los componentes del modelo. La capa de embedding, de positional embedding y las capas de 12 decoders. Después vemos el detalle de cada decoder, con las proyecciones del mecanismo de atención (query, key, value, output) y las dos capas lineales. Vemos que las matrices de atención tienen dimensión 768×768.

Por último, vemos los parámetros entrenables del modelo, que son las matrices adaptadas de LoRA A, de dimensión 8×768 y B, de dimensión 768×8 respectivamente, tanto para el query como el value de cada capa de decoder. También vemos un resumen de los parámetros entrenables, que serían un 0.2349% de los parámetros totales.

UPDATED (Trainable): base_model.model.model.decoder.layers.0.self_attn.v_proj.lora_A.default.weight | Shape: [8, 768]
UPDATED (Trainable): base_model.model.model.decoder.layers.0.self_attn.v_proj.lora_B.default.weight | Shape: [768, 8]
UPDATED (Trainable): base_model.model.model.decoder.layers.0.self_attn.q_proj.lora_A.default.weight | Shape: [8, 768]
UPDATED (Trainable): base_model.model.model.decoder.layers.0.self_attn.q_proj.lora_B.default.weight | Shape: [768, 8]
UPDATED (Trainable): base_model.model.model.decoder.layers.1.self_attn.v_proj.lora_A.default.weight | Shape: [8, 768]
UPDATED (Trainable): base_model.model.model.decoder.layers.1.self_attn.v_proj.lora_B.default.weight | Shape: [768, 8]
UPDATED (Trainable): base_model.model.model.decoder.layers.1.self_attn.q_proj.lora_A.default.weight | Shape: [8, 768]
UPDATED (Trainable): base_model.model.model.decoder.layers.1.self_attn.q_proj.lora_B.default.weight | Shape: [768, 8]