Fusionar y servir

LoRA. Episodio siete: fusionar y servir. LoRA es low-rank adaptation, adaptación de rango bajo. Ya sabemos entrenar un adapter con peft. Hoy, la otra mitad: cómo se usa en producción. Fusionar el adapter con la matriz original para que no cueste nada extra, qué se pierde al deshacer esa fusión, y cómo servir varias tareas con un solo modelo base.

Dos formas de servir

• Sin fusionar: $W_0$ por un lado, $A$ y $B$ por otro
• Fusionado: una sola matriz $W_0 + \Delta W$
• Varios adapters sobre la misma base
• Un batch con un adapter distinto por fila

Hay dos formas de servir un adapter. Sin fusionar: la matriz original, W sub cero, queda como está, y en cada forward se calcula el desvío por A y por B, como en el episodio seis. Fusionado: sumás el cambio una sola vez, y te queda una matriz común, del mismo tamaño que la original. La primera forma permite tener varios adapters cargados sobre la misma base. Y hasta mezclar, en un mismo batch, pedidos que usan adapters distintos.

Parte 1 · Fusionar

El desvío es lineal. Entonces se puede sumar adentro de la matriz, una vez, antes de servir.

Fusionar es sacar factor común

sin fusionar: dos caminos que se suman

factor común x: una sola matriz

multiplicaciones por token: sin y con merge

Sin fusionar, la salida h es W sub cero por x, más alfa sobre r por B por A por x. En el dibujo, dos columnas de números que se suman. Las dos tienen a x a la derecha: sacamos factor común. Queda una sola matriz, W sub cero más delta W, por x. Y delta W, alfa sobre r por B por A, se calcula una vez: en el dibujo, B alta y flaca por A baja y ancha da una matriz de d por k. El costo. Sin fusionar, cada token paga d por k multiplicaciones, más r por d más k del desvío. Fusionado, solo d por k. En la q proj de Gemma, con rango ocho, la diferencia es cero coma siete cuatro por ciento. En el dibujo, ni se nota.

merge_and_unload

Con el bloque de juguete del episodio seis, y B al azar. Guardamos W sub cero y el delta W que calcula peft. merge and unload suma delta W adentro de cada capa y saca todo lo de LoRA. Comprobamos que la matriz nueva menos la vieja es delta W. Comparamos las salidas. Y contamos parámetros. q proj volvió a ser una Linear de torch, común. La resta da delta W. Las salidas difieren en dos por diez a la menos siete: redondeo de float32, porque sumar antes o después de multiplicar no da los mismos bits. Y quedan seiscientos ochenta parámetros, los del bloque original: A y B desaparecieron.

q = modelo.base_model.model.q_proj
W0 = q.base_layer.weight.detach().clone()
dW = q.get_delta_weight("default")             # (α/r)·B·A
y_antes = modelo(x)

fusionado = modelo.merge_and_unload()          # W ← W0 + ΔW
print("q_proj ahora:", type(fusionado.q_proj).__name__,
      "de", type(fusionado.q_proj).__module__.split(".")[0])
W = fusionado.q_proj.weight
print("W − W0 = ΔW:", torch.allclose(W - W0, dW, atol=1e-6))
print("máx |antes − después|:",
      f"{(y_antes - fusionado(x)).abs().max().item():.1e}")
print("parámetros:", sum(p.numel() for p in fusionado.parameters()))

q_proj ahora: Linear de torch
W − W0 = ΔW: True
máx |antes − después|: 2.1e-07
parámetros: 680

¿Cuánto cuesta no fusionar?

Medimos tiempos en la CPU de esta Mac. Cuatro capas de dos mil trescientos cuatro por dos mil trescientos cuatro, quinientos doce tokens, y la mediana de veinte forwards. LoRA de rango ocho en las cuatro. Medimos la base sola y con LoRA sin fusionar. Después, fusionado. Y la cuenta teórica. La base tarda alrededor de ocho milisegundos. Sin fusionar, entre un diez y un veinticinco por ciento más: el número exacto cambia con cada corrida y con lo que esté haciendo la máquina. Fusionado, lo mismo que la base. Pero las multiplicaciones extra eran cero coma sesenta y nueve por ciento. ¿Por qué tanto? Porque el desvío lanza operaciones chicas aparte, y cada una paga su costo fijo. Los números exactos dependen de la máquina; la diferencia entre contar y medir, no.

cfg = LoraConfig(r=8, lora_alpha=16, target_modules=["0", "1", "2", "3"])
con_lora = get_peft_model(copy.deepcopy(base), cfg).eval()
t_base, t_lora = medir(base), medir(con_lora)
t_merge = medir(con_lora.merge_and_unload())

print(f"base sola:      {t_base:5.2f} ms")
print(f"LoRA sin merge: {t_lora:5.2f} ms  (+{t_lora / t_base - 1:.0%})")
print(f"LoRA fusionado: {t_merge:5.2f} ms")
r = 8
print(f"multiplicaciones extra: {r * (d + d) / (d * d):.2%}")

base sola:       7.72 ms
LoRA sin merge:  8.76 ms  (+13%)
LoRA fusionado:  7.43 ms
multiplicaciones extra: 0.69%

Tiempo por forward, 512 tokens

Las tres barras. La base. Sin fusionar, la barra roja: de un diez a un veinticinco por ciento más, por un cero coma siete por ciento de cuentas. Fusionado, la verde: igual a la base, dentro del ruido de la medición. Si servís una sola tarea, fusionar es gratis y te ahorra ese sobrecosto.

• base: 12,52 ms
• sin merge: 21,76 ms
• fusionado: 10,50 ms

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

0 latencia extra · «we can explicitly compute and store W = W0 + BA and perform inference as usual»

Es la propiedad que el paper original destaca frente a otros adapters, que agregan capas y por eso agregan latencia. Calculás y guardás W sub cero más B por A, e inferís como siempre: sin latencia adicional. El mismo paper marca el precio: si fusionás, ya no es directo mezclar en un mismo batch pedidos de tareas con A y B distintas. A eso volvemos en la parte tres.

Parte 2 · Ida y vuelta

peft también puede fusionar sin descargar, y deshacer la fusión después. ¿Volvés exactamente al principio?

merge_adapter y unmerge_adapter

Una capa con la forma de la q proj de Gemma: dos mil trescientos cuatro entradas, dos mil cuarenta y ocho salidas. La probamos en dos tipos de número: float32 y bfloat16, el formato de dieciséis bits con que se suelen servir los modelos. merge adapter suma delta W adentro de W. unmerge adapter lo resta. Y buscamos el peso que más se movió. En float32, el error máximo es dos por diez a la menos nueve: nada. En bfloat16, uno coma dos por diez a la menos cuatro, y cambió el siete por ciento de los pesos. El ejemplo: menos cero coma cero uno siete dos uno dos pasó a menos cero coma cero uno siete cero nueve cero.

for tipo in [torch.float32, torch.bfloat16]:
    torch.manual_seed(0)
    cfg = LoraConfig(r=8, lora_alpha=16, target_modules=["q_proj"])
    m = get_peft_model(Capa().to(tipo), cfg)
    with torch.no_grad():
        m.base_model.model.q_proj.lora_B["default"].weight.normal_(std=0.01)
    W = m.base_model.model.q_proj.base_layer.weight
    W0 = W.detach().clone()

    m.merge_adapter()                          # W ← W0 + ΔW
    m.unmerge_adapter()                        # W ← W − ΔW
    error = (W.float() - W0.float()).abs()
    print(f"{str(tipo):15} máx error {error.max().item():.1e},",
          f"pesos cambiados {(error > 0).float().mean().item():.1%}")

i = error.argmax()                             # el peso que más se movió
print(f"ej. bf16: {W0.flatten()[i].item():.6f} → {W.flatten()[i].item():.6f}")

torch.float32   máx error 1.9e-09, pesos cambiados 5.4%
torch.bfloat16  máx error 1.2e-04, pesos cambiados 7.1%
ej. bf16: -0.017212 → -0.017090

Por qué bfloat16 no vuelve

signo · exponente · mantisa (float32: 1 + 8 + 23)

$2^{-13} \approx 0{,}000122$

cada operación redondea al paso más cercano

Un número bfloat16 tiene dieciséis bits: uno de signo, ocho de exponente y siete de mantisa, los dígitos. float32 tiene veintitrés de mantisa. Con siete bits, entre dos potencias de dos hay solo ciento veintiocho valores posibles. Nuestro peso, cero coma cero diecisiete, cae entre dos a la menos seis y dos a la menos cinco, donde los valores están separados por dos a la menos trece: cero coma cero cero cero ciento veintidós. En el dibujo, el peso bajó exactamente un paso. Cada suma y cada resta redondean al valor representable más cercano, y no siempre se deshacen. En el dibujo, el error de float32 ni se ve. Moraleja: para cambiar de adapter en bfloat16, no restes; recargá la base.

Parte 3 · Varios adapters, una base

Una base congelada de diez gigabytes y muchos adapters de pocos megabytes. Ahora, cómo se eligen.

Una base, varios adapters

La idea, en un dibujo. En el centro, la base, una sola copia en memoria. Alrededor, los adapters: el de la tarea a. El de la tarea b. Y la base sola, sin ningún desvío. set adapter elige cuál se suma en el próximo forward: pasar de uno a otro es cambiar un puntero, no copiar diez gigabytes.

Cien tareas: cien modelos o una base

cien modelos fusionados (GPT-3, del paper)

una base y cien adapters

Gemma 2 2B, q y v, r = 8 (episodio seis)

¿Cuánto ahorra? El paper hace la cuenta en una nota al pie, con GPT-3. Cien tareas con cien modelos fusionados: cien por trescientos cincuenta gigabytes, treinta y cinco terabytes. En el dibujo, contra la otra opción, que casi no se ve. Una base y cien adapters de treinta y cinco megabytes: trescientos cincuenta y cuatro gigabytes. En el dibujo, los adapters juntos suman tres coma cinco. La misma cuenta con Gemma y los números del episodio seis: diez coma cinco gigabytes de base, más cien adapters de seis coma cuatro megabytes, once coma uno en total. Cien modelos fusionados serían más de un terabyte.

Dos adapters, dos tareas

Dos tareas sintéticas sobre una capa de dieciséis por dieciséis: cada una pide un cambio de rango dos distinto sobre W sub cero. El primer adapter se llama tarea a. add adapter agrega el segundo, con la misma receta. set adapter activa uno y lo deja como el único entrenable, y entrenamos cada uno en su tarea. disable adapter, como contexto, apaga todos. Cada adapter, en su tarea, error cero: rango dos alcanza para un cambio de rango dos. En la tarea del otro, seis coma cuatro cuatro. Y la base sola, uno coma nueve nueve y tres coma ocho tres: está más cerca de cada tarea que el adapter equivocado.

cfg = LoraConfig(r=2, lora_alpha=4, target_modules=["proj"])
modelo = get_peft_model(base, cfg, adapter_name="tarea_a")
modelo.add_adapter("tarea_b", cfg)               # segundo adapter

for nombre, Y in tareas.items():
    modelo.set_adapter(nombre)                 # activo y entrenable
    entrenar(modelo, Y)

filas = {}
for activo in ["tarea_a", "tarea_b", None]:
    modelo.set_adapter(activo or "tarea_a")
    if activo is None:                          # base sola
        with modelo.disable_adapter():
            filas["base"] = [error(modelo, Y) for Y in tareas.values()]
    else:
        filas[activo] = [error(modelo, Y) for Y in tareas.values()]
print("error en:  tarea_a tarea_b")
for k, v in filas.items():
    print(f"{k:9} {v[0]:7.3f} {v[1]:7.3f}")

error en:  tarea_a tarea_b
tarea_a     0.000   6.443
tarea_b     6.443   0.000
base        1.986   3.828

Error de cada adapter en cada tarea

La misma tabla, como mapa de calor. El adapter de la tarea a: cero en su tarea, el peor valor en la otra. El de la tarea b, al revés: la diagonal queda en cero. Y la base sola, en el medio. Lo que muestra: un adapter no es una mejora general. Es un cambio para una tarea, y aplicado donde no va, empeora.

Por qué el cruce da lo mismo

la tarea a: la base más un cambio $\Delta_a$ de rango 2

el adapter b aprendió exactamente Δb

por eso 6,443 en las dos celdas

¿Por qué las dos celdas cruzadas dan seis coma cuatro cuatro tres? La tarea a es la base más un cambio delta sub a. En el dibujo, dos columnas que se suman. El adapter b aprendió exactamente delta sub b. Usado en la tarea a, su error es el promedio, sobre las entradas, del cuadrado de la diferencia: delta b menos delta a, por x. En el dibujo, un error por entrada, y su promedio. Y esa diferencia, al cuadrado, es la misma en los dos sentidos. Las dos celdas tienen que coincidir. Es un buen control: si en tu juguete no coinciden, algo no convergió.

Un adapter distinto por fila del batch

Ahora lo que el paper decía que no era directo. Primero, set adapter. Miramos qué entrena: solo el adapter activo. Cuatro pedidos, y para cada uno, qué adapter usar; doble guion bajo base doble guion bajo quiere decir ninguno. peft acepta adapter names en el forward, solo en modo evaluación. Comparamos cada fila contra su adapter usado solo. Entrenan solo lora A y lora B de tarea b. Y las cuatro filas dan diferencia cero. Por dentro, peft agrupa las filas por adapter y calcula cada desvío por separado. Funciona porque no fusionaste: W sub cero sigue limpia.

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

modelo.eval()
quien = ["tarea_a", "tarea_b", "__base__", "tarea_a"]  # uno por fila
mezcla = modelo(x, adapter_names=quien)

for i, nombre in enumerate(quien):
    if nombre == "__base__":
        with modelo.disable_adapter():
            solo = modelo(x)
    else:
        modelo.set_adapter(nombre)
        solo = modelo(x)
    print(f"fila {i} ({nombre}): diferencia",
          (mezcla[i] - solo[i]).abs().max().item())

entrena: proj.lora_A.tarea_b.weight
entrena: proj.lora_B.tarea_b.weight
fila 0 (tarea_a): diferencia 0.0
fila 1 (tarea_b): diferencia 0.0
fila 2 (__base__): diferencia 0.0
fila 3 (tarea_a): diferencia 0.0

Fusionar o no fusionar

La decisión. Fusionado: la misma latencia que la base, pero cada copia del modelo sirve un solo adapter, y para cambiar conviene recargar. Es lo que querés si servís una tarea. Sin fusionar: pagás un sobrecosto, en este juguete entre diez y veinticinco por ciento, a cambio de cargar muchos adapters sobre una base y hasta mezclarlos en un batch. Es lo que querés si servís muchas tareas, o muchos clientes, cada uno con su adapter.

Recetas para producción

• Una tarea: merge_and_unload y guardás el modelo entero
• Muchas tareas: una base y adapters sin fusionar
• Cambiar de adapter en bfloat16: recargá, no restes
• Medí en tu hardware: contar cuentas no alcanza

Cuatro recetas. Si servís una sola tarea, fusioná con merge and unload y guardá el modelo resultante: es un modelo común, cualquier servidor lo carga. Si servís muchas, dejá la base sin tocar y cargá los adapters aparte. Si fusionaste en bfloat16 y querés cambiar, recargá la base desde el archivo: restar deja errores. Y medí: en el juguete, cero coma siete por ciento de cuentas fueron entre diez y veinticinco por ciento de tiempo. En una GPU, graphics processing unit, la proporción puede ser otra.

Resumiendo

• Fusionar: $W = W_0 + \tfrac{\alpha}{r} BA$, factor común de $x$
• merge_and_unload: misma salida, misma latencia que la base
• Sin fusionar: +0,7 % de cuentas, +10 a 25 % de tiempo
• merge/unmerge en bfloat16: el redondeo no se deshace
• set_adapter, disable_adapter, adapter_names por fila
• Próximo: QLoRA, la base en 4 bits

Papers del episodio

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