Dónde ponerlo

LoRA, low-rank adaptation. Episodio cuatro: dónde ponerlo. Ya tenemos la fórmula. Ahora la pregunta práctica: un modelo tiene cientos de matrices, ¿a cuáles les agregás B y A, y con qué rango? Vamos a contar las de Gemma 2 2B una por una, ver qué midió el paper, y probarlo en el juguete del episodio dos.

Parte 1 · Las siete matrices de una capa

Cada una de las veintiséis capas de Gemma tiene siete matrices de pesos grandes. Ahí es donde se puede poner LoRA.

Una capa de Gemma, por dentro

Una capa, de punta a punta, como la vimos en Alineación, episodio uno. Primero la atención: se normaliza la entrada, q proj, k proj y v proj la convierten en queries, keys y values, la atención mezcla los tokens, y o proj, la proyección de salida, devuelve el resultado al residual stream, la corriente que atraviesa todas las capas. Después el MLP, multilayer perceptron: gate proj y up proj agrandan el vector, se combinan con una GELU, gaussian error linear unit, y down proj lo vuelve a achicar. Las siete cajas con proj son matrices: candidatas a LoRA.

norma → q, k, v_proj → atención → o_proj → suma al residual → norma → gate y up_proj → GELU × producto → down_proj → suma al residual

La ficha de Gemma 2 2B

Las formas salen de la ficha del modelo, config punto json. La leemos. Miramos el ancho, las capas, y dos números de la atención: cuántas heads de query y cuántas de key y value. Una head es una de las atenciones que corren en paralelo. Las salidas de q proj: heads por dimensión de cada head. Y las de k y v. Ancho dos mil trescientos cuatro, veintiséis capas. Ocho heads de query y cuatro de key y value, de doscientos cincuenta y seis cada una. Así que q proj saca dos mil cuarenta y ocho números, y k y v, mil veinticuatro.

cfg = json.load(open(d + "config.json"))                # la ficha del modelo
for clave in ["hidden_size", "num_hidden_layers", "num_attention_heads",
              "num_key_value_heads", "head_dim", "intermediate_size"]:
    print(f"{clave:20} {cfg[clave]}")

q = cfg["num_attention_heads"] * cfg["head_dim"]         # salidas de q_proj
kv = cfg["num_key_value_heads"] * cfg["head_dim"]        # salidas de k y v
print("q_proj sale de", q, "| k_proj y v_proj, de", kv)

hidden_size          2304
num_hidden_layers    26
num_attention_heads  8
num_key_value_heads  4
head_dim             256
intermediate_size    9216
q_proj sale de 2048 | k_proj y v_proj, de 1024

Las formas, explicadas

8 heads de query de 256

GQA: cada key y value la comparten 2 heads de query

el MLP agranda 4 veces y vuelve

Cada W es una matriz de d filas, las salidas, por k columnas, las entradas. W sub q: ocho heads por doscientos cincuenta y seis, dos mil cuarenta y ocho filas, por dos mil trescientos cuatro columnas. En el dibujo, un token entra con dos mil trescientos cuatro números y sale con dos mil cuarenta y ocho. W sub k y W sub v tienen la mitad de filas: Gemma usa GQA, grouped-query attention, donde cada par de heads de query comparte una key y un value. Y el MLP: gate y up llevan de dos mil trescientos cuatro a nueve mil doscientos dieciséis, cuatro veces más ancho, y down vuelve.

Las siete matrices, y LoRA en cada una

Ahora contamos, leyendo solo las formas del archivo de pesos. Rango ocho. Para cada una de las siete matrices de la capa cero, d por k. LoRA agrega B, de d por r, y A, de r por k: r por d más k números. Y qué fracción de la matriz es eso. q y o: cuatro coma siete millones cada una; LoRA, treinta y cuatro mil ochocientos, el cero coma setenta y cuatro por ciento. k y v: dos coma cuatro millones; LoRA, el uno coma trece. Y las tres del MLP: veintiún millones cada una; LoRA, noventa y dos mil, apenas el cero coma cuarenta y tres por ciento.

Cada línea, su término:

• r = 8 → $r = 8$ — el rango, igual para todas
• dout, k = forma( → $W \in \mathbb{R}^{d \times k}$ — d salidas, k entradas
• lora = r * (dout + k) → $B: d \times r,\;\; A: r \times k \;\Rightarrow\; r(d+k)$ — lo que agrega LoRA
• lora / (dout * k) → $\frac{r(d+k)}{d\,k}$ — qué fracción de W es

r = 8
filas = []
for m in ["self_attn.q_proj", "self_attn.k_proj", "self_attn.v_proj",
          "self_attn.o_proj", "mlp.gate_proj", "mlp.up_proj", "mlp.down_proj"]:
    dout, k = forma(f"model.layers.0.{m}.weight")       # W es d × k
    lora = r * (dout + k)                                # B: d × r, A: r × k
    filas.append((m.split(".")[1], dout, k, dout * k, lora))
    print(f"{m.split('.')[1]:9} {dout:5}×{k:<5} {dout * k:>11,}"
          f"  LoRA {lora:>7,} ({lora / (dout * k):.2%})")

q_proj     2048×2304    4,718,592  LoRA  34,816 (0.74%)
k_proj     1024×2304    2,359,296  LoRA  26,624 (1.13%)
v_proj     1024×2304    2,359,296  LoRA  26,624 (1.13%)
o_proj     2304×2048    4,718,592  LoRA  34,816 (0.74%)
gate_proj  9216×2304   21,233,664  LoRA  92,160 (0.43%)
up_proj    9216×2304   21,233,664  LoRA  92,160 (0.43%)
down_proj  2304×9216   21,233,664  LoRA  92,160 (0.43%)

Millones de pesos por matriz, en cada capa

En barras, el tamaño de cada matriz de una capa. Azul, las cuatro de la atención: cuatro coma siete, dos coma cuatro, dos coma cuatro, cuatro coma siete millones. Violeta, las tres del MLP: veintiún millones cada una. Las tres del MLP son más del ochenta por ciento de la capa. Si LoRA va solo en la atención, el cambio pasa solo por la parte chica.

• q_proj: 4,7 M
• k_proj: 2,4 M
• v_proj: 2,4 M
• o_proj: 4,7 M
• gate_proj: 21,2 M
• up_proj: 21,2 M
• down_proj: 21,2 M

Parte 2 · Lo que cuesta cada lugar

El costo de LoRA en una matriz depende de su alto más su ancho, no de su área. Eso cambia las cuentas.

r·(d + k) contra d·k

crece con el perímetro, no con el área

dividimos arriba y abajo por dk

$L = 26$ capas · $m$: cada matriz elegida

LoRA en una matriz de d por k cuesta r por d más k; la matriz, d por k. En el dibujo, gate proj con rango ocho: noventa y dos mil contra veintiún millones. Dividiendo las dos cosas por d por k, la fracción es r por uno sobre k más uno sobre d. Cuanto más grande la matriz, más barata le sale LoRA, en proporción. En el dibujo, el porcentaje real de cada una: las del MLP son las más baratas. Y el total: L, las veintiséis capas, por r, por la suma de d más k de cada matriz elegida. En el dibujo, las siete juntas suman cuarenta y nueve mil novecientos veinte por capa.

Tres elecciones típicas, contadas

La cuenta de la fórmula, para cualquier lista de matrices. La elección del paper original: q y v. La otra punta: todas las lineales de cada capa. Con rango ocho y con sesenta y cuatro. q y v con rango ocho: un millón seiscientos mil, el cero coma cero seis por ciento, el número del episodio uno. Las cuatro de la atención: tres coma dos millones. Todas las lineales con rango ocho: diez coma cuatro millones, el cero coma cuatro por ciento. Y con rango sesenta y cuatro, ochenta y tres millones: el tres por ciento del modelo. En el episodio seis, peft va a dar exactamente estos números.

def entrenables(objetivos, r, capas=26):
    por_capa = sum(r * (filas[m][1] + filas[m][2]) for m in objetivos)
    return capas * por_capa

opciones = {"q, v": ["q_proj", "v_proj"],
            "q, k, v, o": ["q_proj", "k_proj", "v_proj", "o_proj"],
            "todas las lineales": list(filas)}
for nombre, objetivos in opciones.items():
    for r in [8, 64]:
        n = entrenables(objetivos, r)
        print(f"{nombre:18} r={r:2}: {n:>11,} ({n / N:.2%})")

q, v               r= 8:   1,597,440 (0.06%)
q, v               r=64:  12,779,520 (0.49%)
q, k, v, o         r= 8:   3,194,880 (0.12%)
q, k, v, o         r=64:  25,559,040 (0.98%)
todas las lineales r= 8:  10,383,360 (0.40%)
todas las lineales r=64:  83,066,880 (3.18%)

Parte 3 · Lo que midió el paper

En el paper de LoRA, con GPT-3: mismo presupuesto de parámetros, distintas matrices.

GPT-3, 18 M de parámetros entrenables (Tabla 5)

La Tabla cinco del paper. Fijaron dieciocho millones de entrenables en GPT-3 y probaron dónde ponerlos. WikiSQL es pasar preguntas en inglés a consultas de base de datos; MultiNLI, decidir si una frase se sigue de otra. Todo en W sub q, con rango ocho: setenta coma cuatro. En W sub k: setenta. En W sub v: setenta y tres. En W sub o: setenta y tres coma dos. Repartido en q y v, con rango cuatro: setenta y tres coma siete. Y en las cuatro, con rango dos: setenta y tres coma siete, y el mejor MultiNLI. Repartir gana.

precisión de validación (%); WikiSQL ±0,5 · MultiNLI ±0,1

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

más matrices > más rango · «it is preferable to adapt more weight matrices than adapting a single type of weights with a larger rank.»

La conclusión del paper, sección siete punto uno. Con el mismo presupuesto, conviene adaptar más matrices con rango chico que una sola con rango grande. Un detalle: el paper solo probó la atención; dejó el MLP congelado y lo anotó como trabajo futuro.

GPT-3, WikiSQL según el rango (Tabla 6)

Y la Tabla seis: ahora cambia el rango. Con LoRA solo en W sub q, el rango mueve algo: de sesenta y ocho coma ocho a setenta coma cinco. En q y v, de rango uno a sesenta y cuatro todo queda entre setenta y tres coma tres y setenta y tres coma ocho: dentro del margen de error. Y en las cuatro, igual: con rango uno, setenta y cuatro coma uno, el mejor número de la tabla. Sesenta y cuatro veces más parámetros no compran nada.

precisión de validación en WikiSQL (%), ±0,5

Dettmers et al., 2023 — QLoRA: Efficient Finetuning of Quantized LLMs · arXiv:2305.14314

todas las lineales · «LoRA on all linear transformer block layers are required to match full finetuning performance.»

Dos años después, el paper de QLoRA, quantized LoRA, que vemos en el episodio ocho, midió con LLaMA de siete mil millones. Con LoRA solo en q y v no lograron igualar al ajuste completo. Encontraron que lo más importante es cuántas matrices llevan LoRA, y que hace falta ponerlo en todas las lineales del bloque, MLP incluido. El rango, en cambio, casi no cambió los resultados.

Parte 4 · En el juguete

La red de dígitos del episodio dos, con LoRA en distintas capas y el mismo presupuesto.

¿Qué capa conviene adaptar?

La tarea es la misma: dígitos en espejo. La red base está congelada, y LoRA entrena con un learning rate propio, así que no comparamos contra el ajuste completo, solo entre ubicaciones. Cuatro pruebas con unos dos mil quinientos entrenables. Todo en la primera capa, la que lee los píxeles. Todo en la última. Repartido en las tres, con rango dos. con lora es una versión mínima de la capa que escribimos entera en el episodio cinco. Primera capa: noventa y seis por ciento. Del medio: noventa y dos coma nueve. Última: ochenta y seis coma nueve. Las tres con rango dos: noventa y cinco coma siete.

# capas 0, 2 y 4 de la red de juguete: 64→256, 256→256, 256→10
pruebas = [("primera, r=8", (0,), 8),                    # ~2.500 entrenables
           ("del medio, r=5", (2,), 5),
           ("última, r=10", (4,), 10),
           ("las tres, r=2", (0, 2, 4), 2)]
res = {}
for nombre, capas, r in pruebas:
    n, acc = con_lora(capas, r)                          # congela, agrega, entrena
    res[nombre] = acc
    print(f"{nombre:15} {n:>5,} entrenables → {acc:.1%}")

primera, r=8    2,560 entrenables → 96.0%
del medio, r=5  2,560 entrenables → 92.9%
última, r=10    2,660 entrenables → 86.9%
las tres, r=2   2,196 entrenables → 95.7%

Mismo presupuesto, distinta capa

La primera capa gana. La del medio, un poco atrás. La última, la que menos. Y repartir en las tres queda casi arriba, con menos parámetros. Tiene sentido: el espejo cambia el orden de los píxeles, y la primera capa es la que los lee. La lección no es la primera capa siempre, sino que dónde poner LoRA depende de qué cambia la tarea, y que repartir es la apuesta segura cuando no lo sabés.

• primera, r=8: 96,0 %
• del medio, r=5: 92,9 %
• última, r=10: 86,9 %
• las tres, r=2: 95,7 %

¿Y el rango?

Ahora fijamos el lugar, las tres capas, y cambiamos el rango. De uno a treinta y dos. Cada corrida congela, agrega LoRA y entrena trescientos pasos. Con rango uno, noventa y uno coma siete. Con dos, noventa y cinco coma siete, igual que con cuatro y con ocho. Dieciséis, noventa y seis coma cinco; treinta y dos, noventa y seis. Desde rango dos, multiplicar los parámetros por dieciséis casi no cambia nada: lo mismo que la Tabla seis del paper.

res = {}
for r in [1, 2, 4, 8, 16, 32]:
    n, acc = con_lora((0, 2, 4), r)                      # LoRA en las tres capas
    res[r] = acc
    print(f"r={r:2}: {n:>6,} entrenables → {acc:.1%}")

r= 1:  1,098 entrenables → 91.7%
r= 2:  2,196 entrenables → 95.7%
r= 4:  4,392 entrenables → 95.7%
r= 8:  8,784 entrenables → 95.7%
r=16: 17,568 entrenables → 96.5%
r=32: 35,136 entrenables → 96.0%

Precisión según el rango, LoRA en las tres capas

En curva. Sube de rango uno a dos, y después queda plana, con ruido de unas décimas. Coincide con el episodio dos: el cambio que hace falta tiene pocas direcciones, así que agregar rango más allá de eso no agrega nada.

Cómo elegir, con lo que vimos

• Repartir en más matrices antes que subir el rango
• Todas las lineales: 10,4 M en Gemma con r = 8 (0,4 %)
• Solo q y v: 1,6 M; anda en muchas tareas, no en todas
• r chico (4–16) suele alcanzar; probá antes de subirlo

Lo que nos llevamos. Con un presupuesto fijo, primero más matrices, después más rango: lo dicen el paper de LoRA, el de QLoRA y el juguete. Poner LoRA en todas las lineales de Gemma con rango ocho son diez coma cuatro millones: sigue siendo menos del medio por ciento. Solo q y v es lo mínimo, la elección original, y anda bien en muchas tareas, pero no en todas. Y el rango: chico suele alcanzar. Subirlo es lo último que conviene probar, no lo primero.

Resumiendo

• Por capa: q, k, v, o en la atención; gate, up, down en el MLP
• GQA: k y v miden la mitad que q (1024 × 2304)
• Costo de LoRA: $r(d+k)$; en proporción, más barato en matrices grandes
• Gemma, r = 8: q y v 1,6 M · todas las lineales 10,4 M
• Paper y juguete: repartir gana; el rango importa poco
• Próximo: LoRA desde cero, en PyTorch

Papers del episodio

• Hu et al., 2021 — LoRA: Low-Rank Adaptation of Large Language Models
• Dettmers et al., 2023 — QLoRA: Efficient Finetuning of Quantized LLMs