LoRA desde cero

LoRA, low-rank adaptation. Episodio cinco: LoRA desde cero. Hasta acá, fórmulas y cuentas. Hoy escribimos la capa entera en PyTorch, en diez líneas, la metemos en la red de dígitos del episodio dos, la entrenamos, verificamos a mano sus gradientes, y al final la fusionamos con los pesos y comprobamos que la salida no cambia.

Parte 1 · La clase

Una capa que envuelve a una nn.Linear existente, la congela, y le agrega A y B al costado.

LoRALinear

La clase recibe una capa lineal ya entrenada, el rango r y alfa. Guarda la capa original como base, y le apaga el gradiente: W sub cero y el sesgo b quedan congelados. A es un nn punto Parameter, un tensor que PyTorch entrena: r por k, sacado de una campana de Gauss con varianza uno sobre k. La N caligráfica es la distribución normal. B, d por r, en cero. El forward: lo que da la base, más alfa sobre r por x por A transpuesta por B transpuesta. Va con transpuestas porque las filas de x son ejemplos. La probamos sobre una capa de sesenta y cuatro a doscientos cincuenta y seis. A y B entrenan; base weight y base bias, no. Y la salida es exactamente igual a la de la base.

Cada línea, su término:

• self.base = base.requires_grad_(False) → $\nabla_{W_0} = 0,\;\; \nabla_b = 0$ — la capa original, congelada
• self.A = nn.Parameter(torch.randn(r, k) → $A \sim \mathcal{N}\!\left(0,\, \tfrac{1}{k}\right)$ — $r \times k$, al azar
• self.B = nn.Parameter(torch.zeros(d, r)) → $B_0 = 0$ — $d \times r$: arranca igual
• return self.base(x) + self.escala → $h = x W_0^\top + b + \tfrac{\alpha}{r}\, x A^\top B^\top$ — filas de x: ejemplos

class LoRALinear(nn.Module):
    def __init__(self, base, r, alpha):
        super().__init__()
        self.base = base.requires_grad_(False)          # W0 y b: congelados
        k, d = base.in_features, base.out_features
        self.A = nn.Parameter(torch.randn(r, k) / k**0.5)  # r × k, al azar
        self.B = nn.Parameter(torch.zeros(d, r))        # d × r, en cero
        self.escala = alpha / r

    def forward(self, x):
        return self.base(x) + self.escala * (x @ self.A.T @ self.B.T)

capa = LoRALinear(nn.Linear(64, 256), r=4, alpha=8)
for nombre, p in capa.named_parameters():
    print(f"{nombre:12} {str(tuple(p.shape)):10} entrena={p.requires_grad}")
x = torch.randn(5, 64)
print("igual a la base:", torch.equal(capa(x), capa.base(x)))

A            (4, 64)    entrena=True
B            (256, 4)   entrena=True
base.weight  (256, 64)  entrena=False
base.bias    (256,)     entrena=False
igual a la base: True

Antes y después de envolver

Lado a lado. La capa original: W y b entrenan, dieciséis mil seiscientos cuarenta números. Envuelta con rango cuatro y alfa ocho: W sub cero y b quietos, y entrenan A, de cuatro por sesenta y cuatro, y B, de doscientos cincuenta y seis por cuatro. Mil doscientos ochenta números, casi trece veces menos. La escala alfa sobre r es ocho sobre cuatro: dos.

Meterla en una red entera

Ahora en la red de dígitos del episodio dos: tres capas lineales, preentrenada con dígitos normales. Primero congelamos toda la red. Recorremos sus capas, y cada lineal se reemplaza por una LoRALinear que la envuelve. Esto es, en chico, lo que hace peft con un modelo de verdad: episodio seis. La aplicamos con rango cuatro y alfa ocho. Y contamos. La capa dos ahora es una LoRALinear, con A de cuatro por doscientos cincuenta y seis y B de doscientos cincuenta y seis por cuatro. Entrenables: cuatro mil trescientos noventa y dos de ochenta y nueve mil, el cuatro coma nueve por ciento. Y antes de entrenar, en espejo acierta el cuarenta y tres coma tres: igual que la red base, como tenía que ser.

def agregar_lora(red, r, alpha):
    red.requires_grad_(False)                         # congelar todo
    for i, capa in enumerate(red):
        if isinstance(capa, nn.Linear):               # cada capa lineal
            red[i] = LoRALinear(capa, r, alpha)       # se envuelve
    return red

red = agregar_lora(cargar_base(), r=4, alpha=8)       # la red del ep. 2
entrena = sum(p.numel() for p in red.parameters() if p.requires_grad)
total = sum(p.numel() for p in red.parameters())
print("capa 2:", type(red[2]).__name__,
      "A", tuple(red[2].A.shape), "B", tuple(red[2].B.shape))
print(f"entrenables: {entrena:,} de {total:,} ({entrena / total:.1%})")
print(f"en espejo, sin entrenar: {precision(red, espejo(Xte), yte):.1%}")

capa 2: LoRALinear A (4, 256) B (256, 4)
entrenables: 4,392 de 89,394 (4.9%)
en espejo, sin entrenar: 43.3%

Qué entrena en cada capa

Los entrenables, capa por capa. Cada una aporta r por d más k. La primera: cuatro por sesenta y cuatro más doscientos cincuenta y seis, mil doscientos ochenta. La del medio: cuatro por quinientos doce, dos mil cuarenta y ocho. La última: cuatro por doscientos sesenta y seis, mil sesenta y cuatro. Contra esto, los ochenta y cinco mil dos parámetros de la red base, congelados: una barra cuarenta veces más alta que la más alta de acá.

• capa 0: 64→256: 1.280
• capa 2: 256→256: 2.048
• capa 4: 256→10: 1.064

Parte 2 · Entrenarla y verificarla

El mismo training loop de siempre. Lo único distinto es qué parámetros recibe el optimizador.

El training loop, solo con A y B

Guardamos una copia de W sub cero de la primera capa, para comprobar al final que nadie la tocó. Los parámetros que entrenan son los que tienen requires grad: los A y los B. Adam, adaptive moment estimation, recibe solo esos, con learning rate cero coma cero uno. Trescientos pasos con batches de dígitos en espejo. Al final, W sub cero contra su copia. Y cuántos tensores tienen estado en Adam. La pérdida baja de siete a cero coma cero cero cinco. En espejo acierta el noventa y seis coma dos por ciento. W sub cero está intacta. Y Adam guarda estado para seis tensores: A y B de tres capas. Los pesos congelados no pagan memoria de optimizador, lo del episodio uno.

W0 = red[0].base.weight.clone()                       # para comprobar después
loras = [p for p in red.parameters() if p.requires_grad]
opt = torch.optim.Adam(loras, lr=1e-2)                # solo A y B
g, historia = torch.Generator().manual_seed(0), []
for paso in range(301):
    i = torch.randint(0, len(Xtr), (64,), generator=g)
    perdida = F.cross_entropy(red(espejo(Xtr[i])), ytr[i])
    opt.zero_grad(); perdida.backward(); opt.step()
    historia.append(perdida.item())
    if paso % 100 == 0:
        print(f"paso {paso:3}: pérdida {perdida.item():.3f}")

print("en espejo:", f"{precision(red, espejo(Xte), yte):.1%}")
print("W0 intacta:", torch.equal(W0, red[0].base.weight))
print("estados de Adam:", len(opt.state), "tensores (A y B de 3 capas)")

paso   0: pérdida 7.048
paso 100: pérdida 0.064
paso 200: pérdida 0.055
paso 300: pérdida 0.005
en espejo: 96.2%
W0 intacta: True
estados de Adam: 6 tensores (A y B de 3 capas)

La pérdida durante el entrenamiento

La pérdida, paso a paso, promediada de a diez pasos porque cada batch es distinto. Arranca en siete: la red preentrenada está muy segura de respuestas equivocadas cuando el dígito está en espejo. En unos cincuenta pasos ya casi resolvió la tarea, y después afina. Los saltitos son batches más difíciles que otros. Todo esto moviendo cuatro mil números de ochenta y nueve mil, con la red base exactamente igual que antes: si mañana sacamos A y B, vuelve a reconocer los dígitos normales como el primer día.

Los gradientes, en forma de batch

$X$: $n \times k$, un ejemplo por fila · $H$: $n \times d$

$\Delta = \partial L / \partial H$: $n \times d$, un $\delta$ por fila

$r \times k$, como $A$

En el episodio tres vimos los gradientes para un solo ejemplo. Con un batch, se apilan. X tiene un ejemplo por fila, n por k; H, la salida, n por d. En el dibujo, X por A transpuesta: lo que lee A de cada ejemplo. Delta mayúscula es el gradiente de la pérdida respecto de H: un delta minúsculo por fila. El gradiente de B es s por delta mayúscula transpuesta por X A transpuesta. En el dibujo, el producto suma sobre los ejemplos y queda de d por r, la forma de B. El de A es s por delta mayúscula por B, transpuesto, por X: r por k, la forma de A.

Gradientes a mano contra autograd

Lo verificamos. Una capa chica, y le ponemos B al azar, como después de entrenar: con B en cero, el gradiente de A daría cero, y la prueba no probaría nada. retain grad le pide a PyTorch que guarde el gradiente de h, que es delta mayúscula. El gradiente de B con la fórmula. Y el de A. Las dos fórmulas coinciden con autograd, el motor de derivadas de PyTorch. Y el gradiente de W sub cero es None: ni siquiera se calcula.

Cada línea, su término:

• capa.B.normal_() → $B \ne 0$ — si no, el gradiente de A daría cero
• h.retain_grad() → $\Delta = \partial L / \partial H$ — PyTorch guarda el gradiente de h
• gB = s * d.T @ (x @ capa.A.T) → $s\, \Delta^\top (X A^\top)$ — la fórmula de B
• gA = s * (d @ capa.B).T @ x → $s\, (\Delta B)^\top X$ — la fórmula de A

capa = LoRALinear(nn.Linear(6, 5), r=2, alpha=4)
with torch.no_grad():
    capa.B.normal_()                                  # como después de entrenar
x = torch.randn(3, 6)                                 # 3 ejemplos
h = capa(x)
h.retain_grad()                                       # guardar δ = ∂L/∂h
perdida = (h ** 2).sum()
perdida.backward()

s, d = capa.escala, h.grad                            # d: 3 × 5
gB = s * d.T @ (x @ capa.A.T)                         # s · δ (Ax)ᵀ
gA = s * (d @ capa.B).T @ x                           # s · (Bᵀδ) xᵀ
print("∇B a mano = autograd:", torch.allclose(gB, capa.B.grad))
print("∇A a mano = autograd:", torch.allclose(gA, capa.A.grad))
print("∇W0:", capa.base.weight.grad)

∇B a mano = autograd: True
∇A a mano = autograd: True
∇W0: None

Parte 3 · Fusionar

Al terminar, B por A se puede sumar a W sub cero una sola vez. La red vuelve a ser una red común.

Por qué fusionar no cambia la salida

se arma una vez: $256 \times 64$ en la capa 0

sacamos x factor común: (BA)ᵀ = AᵀBᵀ

mismo costo que la red original

Delta W es alfa sobre r por B por A. Durante el entrenamiento nunca la armamos; ahora sí, una vez. En el dibujo, la capa cero: doscientos cincuenta y seis por cuatro, por cuatro por sesenta y cuatro. El paso clave: la transpuesta de un producto es el producto de las transpuestas al revés, B por A transpuesta es A transpuesta por B transpuesta. Así, x por W sub cero transpuesta más s por x por A transpuesta por B transpuesta es x por la transpuesta de W sub cero más s B A: sacamos x factor común. En el dibujo, las dos matrices se suman en una. Entonces la capa fusionada es una lineal común. En el dibujo, la precisión que vamos a medir: noventa y seis coma dos antes y después.

Fusionar y comparar

Una función que convierte una LoRALinear en una nn punto Linear común. Arma delta W, una sola vez. Los pesos nuevos son W sub cero más delta W; el sesgo, el de la base. Armamos la red fusionada, capa por capa. Y comparamos las salidas de las dos redes con los mismos dígitos. Además, el rango del cambio de la primera capa. La diferencia máxima entre los logits es dos coma tres por diez a la menos cinco: redondeo de float treinta y dos, porque las sumas se hacen en otro orden. La precisión, noventa y seis coma dos en las dos. El cambio tiene rango cuatro, el r que elegimos. Y la red fusionada tiene ochenta y cinco mil dos parámetros: los mismos que la base.

def fusionar(capa):                                   # LoRALinear → nn.Linear
    base = capa.base
    nueva = nn.Linear(base.in_features, base.out_features)
    with torch.no_grad():
        dW = capa.escala * capa.B @ capa.A            # d × k, se arma una vez
        nueva.weight.copy_(capa.base.weight + dW)     # W = W0 + (α/r)·BA
        nueva.bias.copy_(capa.base.bias)
    return nueva, dW

fusionada = nn.Sequential(*[fusionar(c)[0] if isinstance(c, LoRALinear)
                            else c for c in red])
X = espejo(Xte)
with torch.no_grad():
    dif = (red(X) - fusionada(X)).abs().max().item()
print(f"diferencia máxima de logits: {dif:.1e}")
p1, p2 = precision(red, X, yte), precision(fusionada, X, yte)
print(f"precisión: {p1:.1%} → {p2:.1%}")
dW0 = fusionar(red[0])[1]
print("rango de ΔW, capa 0:", torch.linalg.matrix_rank(dW0).item())
print("parámetros:", sum(p.numel() for p in fusionada.parameters()))

diferencia máxima de logits: 2.3e-05
precisión: 96.2% → 96.2%
rango de ΔW, capa 0: 4
parámetros: 85002

Sin fusionar contra fusionada

Cuándo conviene cada una. Sin fusionar, W sub cero queda intacta y el adapter, A y B, vive aparte: lo podés guardar en un archivo chico, sacarlo o cambiarlo por otro. A cambio, cada forward hace el camino extra. Fusionada, es una red común: mismo tamaño y misma velocidad que la original. Volver atrás exige restar, con un poco de error de redondeo. Servir varios adapters sobre un mismo modelo, y fusionar y desfusionar, es el episodio siete.

Hu et al., 2021 — LoRA: Low-Rank Adaptation of Large Language Models · arXiv:2106.09685

sin latencia extra · «When deployed in production, we can explicitly compute and store W = W0 + BA and perform inference as usual.»

El paper lo plantea como una ventaja frente a otros adapters de la época, que agregaban capas en serie y hacían más lento cada forward. En producción se calcula y se guarda W igual a W sub cero más B A, y se hace inferencia como siempre. Por construcción, LoRA no agrega latencia, el tiempo de respuesta, comparado con un modelo ajustado completo.

Lo que verificamos hoy

Repasemos las pruebas. Con B en cero, la capa da exactamente lo mismo que la base. Los gradientes de A y B, con la fórmula, coinciden con autograd. W sub cero, después de trescientos pasos, es bit a bit la misma. La red aprende la tarea: de cuarenta y tres a noventa y seis por ciento. Y fusionada, la salida cambia solo por redondeo. Cinco pruebas que cualquier implementación de LoRA tendría que pasar; en el episodio seis se las hacemos a peft.

Parte 4 · Lo que falta

Nuestra LoRALinear funciona. Una librería agrega lo que hace falta para usarla en modelos de verdad.

Lo que agrega peft

• Elegir las capas por nombre: target_modules
• lora_dropout: dropout en la entrada del camino LoRA
• Varios adapters con nombre en el mismo modelo
• Guardar y cargar solo A y B

Lo que agrega peft, la librería de Hugging Face, que desarmamos en el episodio seis. En vez de recorrer un Sequential, elige las capas por nombre, como q proj y v proj. Un dropout opcional: apaga al azar parte de la entrada del camino LoRA mientras entrena. Varios adapters con nombre sobre el mismo modelo base. Y guarda solo A y B, el archivo chico del episodio uno. Por dentro, la cuenta es la misma que la nuestra.

Resumiendo

• LoRALinear: base congelada + $A$ al azar + $B$ en cero, 10 líneas
• Inyectar: congelar todo y envolver cada nn.Linear
• Juguete: 4.392 entrenables (4,9 %), espejo 43,3 % → 96,2 %
• $\nabla_B = s\,\Delta^\top XA^\top$, $\nabla_A = s\,(\Delta B)^\top X$: igual que autograd
• Fusionar: $W = W_0 + \frac{\alpha}{r}BA$, misma salida y mismo costo
• Próximo: peft por dentro

Papers del episodio

• Hu et al., 2021 — LoRA: Low-Rank Adaptation of Large Language Models