peft por dentro

LoRA. Episodio seis: peft por dentro. LoRA es low-rank adaptation, adaptación de rango bajo. En el episodio cinco la escribimos a mano, en una clase de veinte líneas. En la práctica casi nadie hace eso: se usa peft, la librería de Hugging Face. Hoy la abrimos para ver qué hace de verdad, y comprobamos que calcula exactamente lo mismo que nuestra clase.

Tres llamadas

• LoraConfig: la receta (r, alpha, qué capas)
• get_peft_model: reemplaza esas capas
• Entrenás como siempre: solo A y B tienen gradiente
• save_pretrained: guarda solo el adapter

peft viene de parameter-efficient fine-tuning: ajustar un modelo entrenando pocos parámetros. Con LoRA, el uso entero son tres llamadas. LoraConfig es la receta: el rango r, el alpha y la lista de capas a modificar. get_peft_model recorre el modelo y reemplaza esas capas por otras con A y B adentro. Después entrenás con tu loop de siempre: solo A y B reciben gradiente. Y save_pretrained guarda solo el adapter, es decir, las matrices chicas, sin el modelo base.

Parte 1 · La receta y el reemplazo

Un modelo de juguete con los mismos nombres que Gemma, para poder leer todo lo que pasa.

LoraConfig y get_peft_model

El juguete es un bloque con cuatro capas lineales: q proj y v proj, de ocho por ocho, y up y down, que agrandan a treinta y dos y vuelven a ocho. La receta: rango dos, lora alpha cuatro, sin dropout. target modules dice qué capas tocar, por nombre. get_peft_model aplica la receta. Contamos los parámetros. Y miramos de qué clase quedó cada capa. Entrenables, sesenta y cuatro de setecientos cuarenta y cuatro: ocho coma seis por ciento. q proj y v proj siguen llamándose Linear, pero ahora son la clase de peft; up sigue siendo la de torch. Y solo cuatro tensores entrenan: lora A y lora B de cada capa reemplazada.

cfg = LoraConfig(r=2, lora_alpha=4, lora_dropout=0.0,
                 target_modules=["q_proj", "v_proj"])
modelo = get_peft_model(Bloque(), cfg)
e, t = modelo.get_nb_trainable_parameters()
print(f"entrenables: {e} de {t} ({e / t:.1%})")

for nombre in ["q_proj", "v_proj", "up"]:
    mod = modelo.get_submodule("base_model.model." + nombre)
    origen = type(mod).__module__.split(".")[0]   # peft o torch
    print(f"{nombre:6} → {type(mod).__name__} de {origen}")

for n, p in modelo.named_parameters():
    if p.requires_grad:
        print("entrena:", n.removeprefix("base_model.model."))

entrenables: 64 de 744 (8.6%)
q_proj → Linear de peft
v_proj → Linear de peft
up     → Linear de torch
entrena: q_proj.lora_A.default.weight
entrena: q_proj.lora_B.default.weight
entrena: v_proj.lora_A.default.weight
entrena: v_proj.lora_B.default.weight

Qué hace get_peft_model

Lo que pasa adentro, en tres pasos. peft recorre todos los submódulos del modelo, con sus nombres completos, como base model punto model punto q proj. Si el final del nombre coincide con algo de target modules, saca esa capa y pone en su lugar un lora punto Linear, que guarda la original adentro, intacta. Después pone requires grad en falso para todos los parámetros del modelo, salvo los de A y B. Ojo: si un nombre no existe, peft avisa; si existe pero no es la capa que querías, no.

modelo base → named_modules() → ¿nombre en target? → lora.Linear → congela la base → A y B entrenables

Contar los 64 entrenables

por cada capa reemplazada

dos capas · rango 2 · d = k = 8

el resto queda congelado

El conteo, a mano. Llamemos d a las salidas de la capa y k a las entradas; r es el rango. B tiene d filas y r columnas; A, r filas y k columnas. En el dibujo, B por A da una matriz de d por k, del tamaño de W, pero se guardan solo r por d más k números. Acá d y k valen ocho, r vale dos, y son dos capas: dos por dos por dieciséis, sesenta y cuatro. En el dibujo, dieciséis por tensor. El resto del bloque: sesenta y cuatro en cada W, doscientos ochenta y ocho en up y doscientos sesenta y cuatro en down, contando los sesgos. Seiscientos ochenta congelados, más sesenta y cuatro: setecientos cuarenta y cuatro.

Parte 2 · Adentro de un lora.Linear

La capa que puso peft: qué guarda, cómo arranca y qué cuenta hace en el forward.

Lo que guarda la capa

Abrimos la q proj reemplazada. base layer es la lineal original, de ocho por ocho. lora A y lora B son lineales sin sesgo, guardadas en un diccionario bajo el nombre default: ese nombre es el del adapter, y en el episodio siete vamos a tener varios. scaling es lora alpha sobre r. B arranca en cero. Y A, al azar. A es de dos por ocho, B de ocho por dos, scaling dos. Máximo de B: cero, así que al principio el modelo da exactamente lo mismo que la base. A cae entre menos cero coma treinta y dos y cero coma treinta y cinco: peft la inicializa uniforme, como cualquier nn.Linear, con cota uno sobre raíz de k.

q = modelo.base_model.model.q_proj           # el lora.Linear
print("base_layer:", tuple(q.base_layer.weight.shape))
A = q.lora_A["default"].weight               # r × k
B = q.lora_B["default"].weight               # d × r
print("lora_A:", tuple(A.shape), " lora_B:", tuple(B.shape))
print("scaling:", q.scaling["default"])      # lora_alpha / r

print("B al empezar, máx |B|:", B.abs().max().item())
k = A.shape[1]
print(f"A en [{A.min():.2f}, {A.max():.2f}],",
      f"cota 1/√k = {1 / math.sqrt(k):.2f}")

base_layer: (8, 8)
lora_A: (2, 8)  lora_B: (8, 2)
scaling: 2.0
B al empezar, máx |B|: 0.0
A en [-0.32, 0.35], cota 1/√k = 0.35

El forward de lora.Linear

El forward, como grafo de cómputo. Arriba, el camino de siempre: x por la W original. Abajo, el desvío: x pasa por A, después por B, se multiplica por scaling, y las dos ramas se suman. Las formas, con cinco tokens de ocho números: A comprime cada token a dos números, el rango, y B los devuelve a ocho. En el backward, el gradiente llega a A y a B, pero no a W: está congelada.

La cuenta del forward

$x$: un token, vector de $k = 8$ números

primero se comprime: 8 → 2

después se expande: 2 → 8

lora_alpha / r

La misma cuenta, en fórmula, para un token. x es un vector columna de ocho números. W sub cero es la matriz original, congelada; en el dibujo, W sub cero por x da ocho números, una fila por salida. El desvío nunca arma la matriz B por A: primero calcula A por x, que da r números, dos. En el dibujo, cada fila de A produce uno. Después B por ese resultado vuelve a ocho números. Así cuesta r por k más d por r multiplicaciones, en vez de d por k. Y scaling es alfa sobre r: cuatro sobre dos, dos. En el código, x viene en filas, uno por token, por eso peft hace x por A transpuesta.

Parte 3 · ¿Calcula lo mismo que la nuestra?

Copiamos A y B de peft a la LoRALinear del episodio cinco y comparamos las salidas número a número.

peft contra nuestra LoRALinear

Primero, B al azar, para que el desvío no sea cero. Creamos nuestra LoRALinear sobre la misma base layer, y le copiamos A y B. Pasamos cinco tokens por las dos y buscamos la mayor diferencia entre las cuarenta salidas. get delta weight arma la matriz del cambio, delta W: alfa sobre r por B por A. Y medimos su rango. La diferencia es cero, exacto: las mismas operaciones en el mismo orden dan los mismos bits. Delta W es de ocho por ocho con rango dos, y coincide con escala por B por A de nuestra clase. Es nuestra fórmula.

Cada línea, su término:

• mia.A.copy_( → $A_{\text{mía}} \leftarrow A_{\text{peft}},\;\; B_{\text{mía}} \leftarrow B_{\text{peft}}$ — los mismos números en las dos
• (y_peft - y_mia).abs().max() → $\max_{i}\, \lvert h_i^{\text{peft}} - h_i^{\text{mía}} \rvert$ — la mayor diferencia, en 40 salidas
• q.get_delta_weight("default") → $\Delta W = \tfrac{\alpha}{r}\, B A$ — la matriz del cambio, $8 \times 8$
• matrix_rank(dW) → $\operatorname{rango}(\Delta W) \le r = 2$ — una matriz de 8 × 8 con rango 2

with torch.no_grad():                         # B ≠ 0, como tras entrenar
    q.lora_B["default"].weight.normal_()

mia = LoRALinear(q.base_layer, r=2, alpha=4)
with torch.no_grad():                         # copiamos A y B de peft
    mia.A.copy_(q.lora_A["default"].weight)
    mia.B.copy_(q.lora_B["default"].weight)

y_peft, y_mia = q(x), mia(x)
print("salida:", tuple(y_peft.shape))
print("máx |peft − mía|:", (y_peft - y_mia).abs().max().item())
dW = q.get_delta_weight("default")            # (α/r)·B·A
print("ΔW:", tuple(dW.shape), "rango", torch.linalg.matrix_rank(dW).item())
print("igual a escala·B·A:", torch.allclose(dW, mia.escala * mia.B @ mia.A))

salida: (5, 8)
máx |peft − mía|: 0.0
ΔW: (8, 8) rango 2
igual a escala·B·A: True

Entrenar: lo congelado sigue congelado

Ahora entrenamos, con datos sintéticos: sesenta y cuatro entradas y salidas al azar. Le pasamos al optimizador solo los parámetros con requires grad. AdamW, que es Adam, adaptive moment estimation, con weight decay desacoplado, doscientos pasos, learning rate uno por diez a la menos dos. Al final miramos el gradiente de la W original. Y si W cambió. La pérdida baja de cero coma nueve nueve siete a cero coma ocho seis nueve: poco, porque los objetivos son ruido y rango dos no alcanza para copiarlos, pero baja. El gradiente de W es None, nunca se calculó. W está intacta, bit a bit. Y B ya no es cero: su mayor valor es cero coma cincuenta y seis.

q = modelo.base_model.model.q_proj
W_antes = q.base_layer.weight.clone()
entrenables = [p for p in modelo.parameters() if p.requires_grad]
opt = torch.optim.AdamW(entrenables, lr=1e-2)

for paso in range(200):
    perdida = ((modelo(X) - Y) ** 2).mean()
    if paso == 0:
        print(f"pérdida inicial: {perdida.item():.3f}")
    opt.zero_grad(); perdida.backward(); opt.step()

print(f"pérdida final:   {perdida.item():.3f}")
print("grad de W:", q.base_layer.weight.grad)
print("W intacta:", torch.equal(W_antes, q.base_layer.weight))
print(f"máx |B| ahora: {q.lora_B['default'].weight.abs().max():.2f}")

pérdida inicial: 0.997
pérdida final:   0.869
grad de W: None
W intacta: True
máx |B| ahora: 0.56

Nuestra clase y la de peft

Lado a lado. Nuestra clase tiene lo esencial: A, B, la escala y la base congelada, para una sola tarea. peft hace la misma cuenta y le suma infraestructura: dropout opcional sobre la entrada del desvío, A y B guardadas por nombre de adapter, para tener varios en la misma capa, la opción de fusionarlos con W, y variantes como DoRA, weight-decomposed LoRA, que vemos en el episodio nueve. Y el reemplazo automático, en modelos de cientos de capas.

Parte 4 · Gemma y el archivo

El mismo reemplazo sobre Gemma 2 2B, sin cargar sus pesos. Y qué queda en disco cuando guardás.

LoRA sobre Gemma 2 2B, sin pesos

Para contar no hacen falta los pesos. Con el device meta, PyTorch crea los tensores con su forma pero sin memoria: el modelo entero se arma en un segundo. Dos recetas con rango ocho: solo q y v, como en el paper original, y all-linear, todas las lineales de cada capa. Contamos las capas que quedaron con lora A. Y los entrenables. Con q y v: cincuenta y dos capas, veintiséis por dos, y un millón quinientos noventa y siete mil cuatrocientos cuarenta entrenables, el mismo número del episodio uno: cero coma cero seis por ciento. Con all-linear: ciento ochenta y dos capas, siete por bloque, y diez coma cuatro millones. En q proj, A es de ocho por dos mil trescientos cuatro, y B de dos mil cuarenta y ocho por ocho.

for nombre, objetivo in [("q y v", ["q_proj", "v_proj"]),
                         ("all-linear", "all-linear")]:
    with torch.device("meta"):                # sin pesos: solo las formas
        base = AutoModelForCausalLM.from_config(cfg_gemma)
    cfg = LoraConfig(r=8, lora_alpha=16, target_modules=objetivo)
    m = get_peft_model(base, cfg)
    capas = [n for n, x in m.named_modules() if hasattr(x, "lora_A")]
    entr, total = m.get_nb_trainable_parameters()
    print(f"{nombre + ':':11} {len(capas)} capas,",
          f"{entr:,} ({entr / total:.3%})")

q = m.base_model.model.model.layers[12].self_attn.q_proj
print("capa 12, q_proj: A", tuple(q.lora_A["default"].weight.shape),
      "B", tuple(q.lora_B["default"].weight.shape))
print(f"total: {total:,}")

q y v:      52 capas, 1,597,440 (0.061%)
all-linear: 182 capas, 10,383,360 (0.396%)
capa 12, q_proj: A (8, 2304) B (2048, 8)
total: 2,624,725,248

Cuánto entrena cada receta (r = 8)

En una tabla. q y v: un millón seiscientos mil números; a cuatro bytes cada uno, el adapter pesa seis coma cuatro megabytes. all-linear: seis veces y media más, cuarenta y un megabytes. Y Gemma entera, dos mil seiscientos catorce millones de parámetros: diez coma cinco gigabytes. Los dos adapters son menos del medio por ciento del modelo.

Qué guarda save_pretrained

Volvemos al juguete. save pretrained escribe una carpeta. Abrimos el archivo de pesos, un safetensors. Leemos la configuración. Y cargamos el adapter sobre una base nueva, creada con la misma semilla. Tres archivos: un README, la configuración en JSON, y seiscientos setenta y dos bytes de pesos. Adentro, solo cuatro tensores: lora A y lora B de q y de v. Fijate que en las claves desapareció la palabra default: el nombre del adapter se elige al cargar. La configuración guarda rango dos, alpha cuatro y los dos nombres. Y la salida cargada es igual a la original.

modelo.save_pretrained(carpeta)
for f in sorted(os.listdir(carpeta)):
    print(f, os.path.getsize(os.path.join(carpeta, f)), "bytes")

pesos = load_file(os.path.join(carpeta, "adapter_model.safetensors"))
for k, v in pesos.items():
    print(" ", k.removeprefix("base_model.model."), tuple(v.shape))
c = json.load(open(os.path.join(carpeta, "adapter_config.json")))
print("config:", c["r"], c["lora_alpha"], sorted(c["target_modules"]))

otro = PeftModel.from_pretrained(base(), carpeta)   # base + adapter
print("misma salida:", torch.allclose(modelo(x), otro(x)))

README.md 5074 bytes
adapter_config.json 1139 bytes
adapter_model.safetensors 672 bytes
  q_proj.lora_A.weight (2, 8)
  q_proj.lora_B.weight (8, 2)
  v_proj.lora_A.weight (2, 8)
  v_proj.lora_B.weight (8, 2)
config: 2 4 ['q_proj', 'v_proj']
misma salida: True

El archivo pesa lo que pesan A y B

float32: 4 bytes por número

Gemma, q y v, r = 8

El archivo de pesos son cuatro bytes por número, porque guardamos en float32. En el dibujo: dieciséis números por tensor, sesenta y cuatro bytes cada uno, doscientos cincuenta y seis en total. Los seiscientos setenta y dos del archivo incluyen un encabezado con los nombres y las formas. En Gemma, con q y v: un millón quinientos noventa y siete mil cuatrocientos cuarenta por cuatro, seis coma cuatro megabytes. En el dibujo, contra los diez mil cuatrocientos cincuenta y siete del modelo base, la barra del adapter casi no se ve. Por eso se comparten adapters y no modelos.

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

350 GB → 35 MB · Con GPT-3 175B, r = 4 y solo q y v: «the checkpoint size is reduced by roughly 10,000× (from 350GB to 35MB)».

El paper original ya lo decía. Con GPT-3, de ciento setenta y cinco mil millones de parámetros, rango cuatro y solo las proyecciones de query y value, el checkpoint pasa de trescientos cincuenta gigabytes a treinta y cinco megabytes: unas diez mil veces menos. Y agregan la cuenta de servir cien tareas: un modelo base más cien adapters, en vez de cien modelos enteros.

Lo que conviene saber al usar peft

• target_modules busca por nombre: mirá named_modules()
• El adapter solo sirve sobre la misma base
• scaling = lora_alpha / r: cambiar r cambia la escala
• get_delta_weight te da ΔW para inspeccionar

Cuatro cosas prácticas. target modules busca por nombre: antes de configurar, imprimí los nombres del modelo. El adapter es una diferencia sobre una base concreta: cargado sobre otra base, la salida no tiene sentido, y nada te avisa si las formas coinciden. scaling es lora alpha sobre r: si subís r sin tocar alpha, el desvío se achica. En el episodio nueve vemos una variante que cambia esa escala. Y get delta weight te da delta W como matriz, para medirla: rango, norma, lo que quieras.

Resumiendo

• LoraConfig + get_peft_model: reemplazo por nombre
• lora.Linear: base congelada + $\tfrac{\alpha}{r} B(Ax)$, B en cero
• Misma salida que la LoRALinear del ep. 5: diferencia 0
• Solo A y B tienen gradiente; W queda intacta
• Gemma, q y v, r = 8: 1,6 M entrenables, 6,4 MB
• Próximo: fusionar, servir y cambiar de adapter

Papers del episodio

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