Por qué no ajustar todo

LoRA. Episodio uno: por qué no ajustar todo. LoRA quiere decir low-rank adaptation: adaptación de rango bajo. Es la forma más común de ajustar un modelo abierto a una tarea propia. En diez episodios la vamos a entender entera: la matemática, el código desde cero y lo que hace la librería. Hoy, el problema que resuelve, medido en bytes.

El mapa de la serie

El recorrido. Los tres primeros son la idea: por qué no conviene ajustar todos los pesos, la hipótesis de que el cambio necesario vive en pocas direcciones, y la fórmula. Después, implementarla: en qué matrices del modelo se pone, una capa LoRA escrita a mano en PyTorch, y qué hace por dentro peft, la librería de Hugging Face. Luego, usarla: fusionar el cambio con los pesos para servir el modelo, y QLoRA, quantized LoRA, sobre un modelo en cuatro bits. Y al final, las variantes, y qué aprende LoRA comparado con un ajuste completo.

Por qué no todo → Rango intrínseco → La fórmula → Dónde ponerlo → LoRA desde cero → peft por dentro → Fusionar y servir → QLoRA → Variantes → Qué aprende

Parte 1 · Contar antes de entrenar

Antes de hablar de memoria, hay que saber cuántos números tiene el modelo. Los contamos sin cargarlo.

Los parámetros de Gemma 2 2B, por grupo

Gemma 2 2B guarda sus pesos en archivos safetensors, con un índice que dice qué tensor está en qué archivo. Un tensor es una tabla de números de cualquier cantidad de dimensiones. Vamos a sumar por grupo. Recorremos los doscientos ochenta y ocho tensores, leyendo solo su forma, no los números. La cantidad de números es el producto de los lados: dos mil cuarenta y ocho por dos mil trescientos cuatro, por ejemplo. Cada nombre dice a qué parte pertenece. Y el total. Dos mil seiscientos catorce millones de parámetros. El embedding, la tabla con un vector por token, es el veintidós por ciento. La atención, el catorce. El MLP, multilayer perceptron, el sesenta y tres. Las normas casi no pesan.

grupos = {"embedding": 0, "atención": 0, "MLP": 0, "normas": 0}
for nombre in indice:                       # 288 tensores con nombre
    n = 1
    for lado in forma(nombre):              # p. ej. [2048, 2304]
        n *= lado
    if "embed" in nombre:      grupos["embedding"] += n
    elif "self_attn" in nombre: grupos["atención"] += n
    elif "mlp" in nombre:      grupos["MLP"] += n
    else:                      grupos["normas"] += n

total = sum(grupos.values())
for g, n in grupos.items():
    print(f"{g:10} {n:>13,}  {n / total:6.1%}")
print(f"{'total':10} {total:>13,}")

embedding    589,824,000   22.6%
atención     368,050,176   14.1%
MLP        1,656,225,792   63.4%
normas           241,920    0.0%
total      2,614,341,888

Dónde están los 2.614 millones

En barras, en millones. El embedding: casi seiscientos millones, por un vocabulario de doscientos cincuenta y seis mil tokens. La atención, en las veintiséis capas: trescientos sesenta y ocho millones. El MLP: la red densa que cada capa tiene después de la atención. Es la barra más larga, con mil seiscientos cincuenta y seis millones. Y las normas, que no se ven. En el episodio cuatro volvemos a estas barras para decidir dónde poner LoRA.

• embedding: 589,8 M
• atención: 368,1 M
• MLP: 1.656,2 M
• normas: 0,24 M

Parte 2 · Cuatro números por parámetro

Para usar el modelo alcanza con los pesos. Para entrenarlo, cada peso trae compañía.

Un paso de entrenamiento, y lo que deja en memoria

Un paso de entrenamiento, que viste en Deep learning a secas, episodio siete. El forward: la entrada pasa por el modelo. La pérdida: un número que mide cuánto se equivocó. El backward, o backpropagation: un gradiente por cada peso, del mismo tamaño que el modelo. El optimizador Adam, adaptive moment estimation, guarda dos promedios por peso: m y v. Y con todo eso mueve los pesos. Los gradientes se pueden tirar después de cada paso, pero hay que tenerlos a la vez que el resto; m y v viven todo el entrenamiento.

Qué guarda Adam, en PyTorch

Con una capa chica se ve todo. Una capa lineal de tres entradas y dos salidas: una matriz de pesos de dos por tres, más dos sesgos. Una pérdida cualquiera, y backward llena el gradiente. El primer paso del optimizador crea su estado. opt punto state guarda, para cada parámetro, lo que Adam necesita. Cuántos bytes ocupa cada número. Y cuánto se movió cada peso. Peso, gradiente, exp avg y exp avg sq: cuatro tablas de dos por tres. Cuatro bytes por número, porque es float treinta y dos. Y cada peso se movió exactamente cero coma uno, para un lado o para el otro. Ya vamos a ver por qué.

capa = torch.nn.Linear(3, 2)                # 6 pesos + 2 sesgos
opt = torch.optim.Adam(capa.parameters(), lr=0.1)

x = torch.tensor([[1.0, 2.0, 3.0]])
perdida = capa(x).pow(2).sum()
perdida.backward()                          # llena .grad
w0 = capa.weight.detach().clone()
opt.step()                                  # crea m y v

w = capa.weight
estado = opt.state[w]
print("peso      ", tuple(w.shape))
print("gradiente ", tuple(w.grad.shape))
print("m (exp_avg)   ", tuple(estado["exp_avg"].shape))
print("v (exp_avg_sq)", tuple(estado["exp_avg_sq"].shape))
print("bytes por número:", w.element_size())
print("cambio:", [round(c, 3) for c in (w.detach() - w0).flatten().tolist()])

peso       (2, 3)
gradiente  (2, 3)
m (exp_avg)    (2, 3)
v (exp_avg_sq) (2, 3)
bytes por número: 4
cambio: [0.1, 0.1, 0.1, -0.1, -0.1, -0.1]

Los dos estados de Adam

$g$: el gradiente · $\beta_1 = 0{,}9$

$\beta_2 = 0{,}999$ · el cuadrado, número por número

$\eta$: learning rate · paso 1: $\hat m / \sqrt{\hat v} = \pm 1$

peso, gradiente, m y v en float32

Adam, por dentro, lo desarma la serie de optimizadores en su episodio cinco. g es el gradiente de un peso. m es un promedio que se va actualizando: beta uno, cero coma nueve, de lo que había, más un décimo del gradiente nuevo. En el dibujo, dos columnas pesadas que se suman. v hace lo mismo con el gradiente al cuadrado, con beta dos, cero coma nueve nueve nueve. El paso: m sombrero sobre la raíz de v sombrero, por eta, el learning rate. Los sombreros corrigen que m y v arrancan en cero; épsilon evita dividir por cero. En el primer paso, m sombrero es g y v sombrero es g al cuadrado: el cociente da más o menos uno. En el dibujo, el cambio real: cero coma uno exacto. Y la cuenta que importa: dieciséis bytes por parámetro.

La memoria de un ajuste completo de Gemma

Ahora con Gemma. N es el total que contamos: dos mil seiscientos catorce millones. Los pesos en float treinta y dos: cuatro bytes cada uno. Un GB son mil millones de bytes: diez coma cinco GB. Los gradientes, m y v: tres copias más, del mismo tamaño. Y el total. Diez coma cinco cada una, cuarenta y uno coma ocho en total. Y falta lo que no contamos: las activaciones, los resultados intermedios del forward que el backward necesita, y crecen con el largo del texto y el batch. Entrenar en precisión mixta mueve estos números, pero no la conclusión.

Cada línea, su término:

• N = c["total"] → $N = 2.614.341.888$ — parámetros de Gemma 2 2B
• pesos = N * b → $4N \approx 10{,}5 \text{ GB}$ — float32: 4 bytes cada uno
• v = N * b → $4N + 4N + 4N + 4N$ — pesos, gradientes, m y v
• print(f"{'total' → $16N \approx 41{,}8 \text{ GB}$ — sin contar activaciones

N = c["total"]                              # 2.614.341.888 parámetros
b = 4                                       # float32: 4 bytes
pesos = N * b
grads = N * b                               # un gradiente por peso
m = N * b                                   # Adam: promedio del gradiente
v = N * b                                   # Adam: promedio del cuadrado
GB = 1e9
for nombre, x in [("pesos", pesos), ("gradientes", grads),
                  ("Adam m", m), ("Adam v", v)]:
    print(f"{nombre:11} {x / GB:5.1f} GB")
print(f"{'total':11} {(pesos + grads + m + v) / GB:5.1f} GB")
print("bytes por parámetro:", 4 * b)

pesos        10.5 GB
gradientes   10.5 GB
Adam m       10.5 GB
Adam v       10.5 GB
total        41.8 GB
bytes por parámetro: 16

Lo que se va apilando

Apilado. Solo los pesos, que es lo que hace falta para usar el modelo: diez coma cinco GB. Más los gradientes: veintiuno. Más m: treinta y uno. Más v: cuarenta y dos. Cuatro veces lo que pesa el modelo. Como referencia, una placa de video de consumo de las grandes trae veinticuatro GB. Con un modelo de los más chicos que hay, el ajuste completo ya no entra en una GPU, graphics processing unit, de esas.

• pesos: 10,5 GB
• + gradientes: 20,9 GB
• + Adam m: 31,4 GB
• + Adam v: 41,8 GB
• placa de 24 GB: 24 GB

Parte 3 · Congelar

Lo que no se entrena no necesita gradiente ni estado de Adam. Ese es el primer ahorro.

Un peso congelado no cuesta nada extra

Probémoslo. Una capa base de dos mil trescientos cuatro por dos mil trescientos cuatro, el ancho de Gemma. Hace de los pesos del modelo. Una parte entrenable mucho más chica, de ocho salidas. requires grad en falso congela la base: PyTorch no le calcula gradiente. El optimizador recibe solo los parámetros de la parte entrenable. Y miramos qué tiene cada uno. La base: gradiente None, nada; y sin estado en Adam. La parte extra: gradiente de ocho por dos mil trescientos cuatro, y estado. La base tiene cinco millones trescientos mil números; la extra, dieciocho mil cuatrocientos cuarenta.

base = torch.nn.Linear(2304, 2304)          # hace de W: se congela
extra = torch.nn.Linear(2304, 8)            # hace de la parte entrenable
base.requires_grad_(False)                  # sin gradiente

opt = torch.optim.Adam(extra.parameters(), lr=1e-3)
x = torch.randn(4, 2304)
(base(x).sum() + extra(x).sum()).backward()
opt.step()

for nombre, capa in [("base", base), ("extra", extra)]:
    w = capa.weight
    g = None if w.grad is None else tuple(w.grad.shape)
    n = sum(p.numel() for p in capa.parameters())
    print(f"{nombre:5} grad={g}  Adam={w in opt.state}  n={n:,}")

base  grad=None  Adam=False  n=5,310,720
extra grad=(8, 2304)  Adam=True  n=18,440

La cuenta con el modelo congelado

$M$: memoria en bytes · $N$: todos los parámetros

$N_e$: los entrenables · el resto, solo pesos

$r = 8$: 1.597.440

En fórmulas. M es la memoria en bytes y N la cantidad de parámetros. Ajustando todo, dieciséis bytes por cada uno. En el dibujo, las cuatro barras de diez coma cinco. Si congelamos casi todo y entrenamos solo N sub e parámetros, el modelo cuesta cuatro bytes por peso, y solo los entrenables pagan los dieciséis. En el dibujo, las tres barras de la derecha casi desaparecen. ¿Cuántos entrenables? Un adelanto del episodio tres: LoRA agrega dos matrices chicas, B y A, al lado de una matriz grande. En el dibujo, las formas de q proj, la proyección de queries de la atención. Con r igual a ocho, en q proj y v proj de las veintiséis capas, son un millón quinientos noventa y siete mil.

La cuenta, con LoRA

En código. Rango ocho, veintiséis capas. Para q proj, B y A suman ocho por dos mil cuarenta y ocho más dos mil trescientos cuatro. Para v proj, igual con mil veinticuatro. Los entrenables, en todas las capas. La memoria: el modelo a cuatro bytes, y los entrenables a dieciséis. Y lo que hay que guardar por cada tarea. Uno coma seis millones de entrenables: cero coma cero seis uno por ciento del modelo. La memoria baja de cuarenta y uno coma ocho a diez coma cinco GB: casi todo es el modelo congelado. Y el archivo por tarea pesa seis coma cuatro megabytes, contra diez coma cinco GB de una copia entera.

r, capas = 8, 26
q = r * (2048 + 2304)                       # B y A de q_proj
v = r * (1024 + 2304)                       # B y A de v_proj
Ne = capas * (q + v)                        # lo único que se entrena
print(f"entrenables: {Ne:,} ({Ne / N:.3%})")

GB = 1e9
completo = 16 * N                           # 4 copias de 4 bytes
lora = 4 * N + 16 * Ne                      # W solo pesos; A y B, las 4
print(f"ajuste completo: {completo / GB:5.1f} GB")
print(f"con LoRA:        {lora / GB:5.1f} GB")
print(f"archivo por tarea: {4 * Ne / 1e6:.1f} MB contra {4 * N / GB:.1f} GB")

entrenables: 1,597,440 (0.061%)
ajuste completo:  41.8 GB
con LoRA:         10.5 GB
archivo por tarea: 6.4 MB contra 10.5 GB

Ajuste completo contra LoRA

Lado a lado. El ajuste completo entrena los dos mil seiscientos millones, paga dieciséis bytes por cada uno, y deja una copia nueva del modelo por cada tarea. Con LoRA, la memoria es prácticamente la de usar el modelo, y cada tarea es un archivo de seis megas que se pone y se saca. La pregunta que queda es si alcanza: si un cambio tan chico puede hacer lo que hace el ajuste completo. Esa es la hipótesis del episodio dos.

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

1,2 TB → 350 GB · «On GPT-3 175B, we reduce the VRAM consumption during training from 1.2TB to 350GB.»

El paper que lo propuso, de Microsoft, dos mil veintiuno. Lo midieron con GPT tres, de ciento setenta y cinco mil millones de parámetros, entrenado con Adam. La VRAM, video RAM, la memoria de las placas, durante el entrenamiento bajó de uno coma dos terabytes a trescientos cincuenta gigas. La explicación del paper es la misma que contamos hoy: no hace falta guardar el estado del optimizador para los parámetros congelados.

Parte 4 · La apuesta

Ahorrar memoria es fácil: entrenás menos. Lo difícil es que siga funcionando.

Lo que la serie tiene que mostrar

• Que el cambio útil cabe en pocas direcciones
• Cómo se escribe: $W + \frac{\alpha}{r} B A$
• En qué matrices conviene ponerlo
• Que al final se funde con W y no cuesta nada

Congelar todo y entrenar uno coma seis millones solo sirve si esos pocos parámetros alcanzan. La hipótesis: el cambio que un fine-tuning le hace a cada matriz tiene rango bajo, cabe en pocas direcciones. Lo medimos en el episodio dos. La fórmula: W, la matriz original, más B por A escalado por alfa sobre r. En Álgebra lineal, episodio tres, vimos por qué B por A tiene rango bajo. En qué matrices conviene ponerlo, con los conteos de Gemma: episodio cuatro. Y que después del entrenamiento, B por A se suma a W y el modelo vuelve a tener el mismo tamaño y velocidad.

Resumiendo

• Gemma 2 2B: 2.614 M parámetros; el MLP es el 63 %
• Adam guarda $m$ y $v$: dos copias más del modelo
• Ajuste completo: 16 bytes por parámetro, 41,8 GB sin activaciones
• Congelar: sin gradiente ni estado; solo pagan los entrenables
• LoRA r = 8 en q y v: 1,6 M entrenables, 10,5 GB, 6,4 MB por tarea
• Próximo: por qué alcanza con rango bajo

Papers del episodio

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