QLoRA

LoRA. Episodio ocho: QLoRA, quantized LoRA. LoRA es low-rank adaptation, adaptación de rango bajo. Con LoRA ya no guardamos gradientes ni estados del optimizador para la base. Pero la base sigue ocupando memoria: diez gigabytes en float32 para Gemma 2 2B. QLoRA la guarda en cuatro bits por número. Hoy, a mano: la librería de verdad pide una placa NVIDIA.

La idea de QLoRA

• La base no se entrena: se puede guardar comprimida
• Cuatro bits por peso, con una escala por bloque
• En cada forward se descomprime a 16 bits
• A y B, en 16 bits, entrenan encima

La idea entra en cuatro líneas. La base está congelada: nunca la actualizamos, así que no hace falta guardarla con toda su precisión. Cada peso se guarda en cuatro bits, más una escala compartida por cada bloque de pesos. Eso es quantization: representar números con menos bits. En cada forward, la capa descomprime sus pesos a dieciséis bits y multiplica como siempre. Y A y B, las matrices de LoRA, viven en dieciséis bits y son lo único que entrena.

Parte 1 · Cuatro bits por número

Con cuatro bits hay dieciséis valores posibles. La pregunta es cuáles, y cómo se reparte un peso real entre ellos.

absmax: la escala del bloque

Un bloque de ocho pesos inventados. absmax es el mayor valor absoluto del bloque: cero coma cinco. Dividimos por absmax, así todo cae entre menos uno y uno; multiplicamos por siete y redondeamos. Quedan enteros de menos siete a siete: quince valores, que entran en cuatro bits. Para volver, el camino inverso. Y probamos con un valor atípico. Los códigos: dos, menos cuatro, uno, seis, y así. De vuelta, los pesos se parecen: el error máximo es cero coma cero veinticuatro. Pero con un cuatro en el bloque, casi todo da cero: el atípico se quedó con toda la escala.

w = torch.tensor([0.12, -0.31, 0.05, 0.44, -0.08, 0.27, -0.50, 0.19])
absmax = w.abs().max()                       # la constante del bloque
codigo = torch.round(w / absmax * 7)         # enteros de −7 a 7: 4 bits
w_hat = codigo / 7 * absmax                  # de vuelta a números reales
print("absmax:", absmax.item())
print("código:", codigo.int().tolist())
print("vuelta:", w_hat)
print(f"error máx: {(w - w_hat).abs().max():.3f}")

w[3] = 4.0                                   # un valor atípico en el bloque
absmax = w.abs().max()
codigo = torch.round(w / absmax * 7)
print("con atípico:", codigo.int().tolist())

absmax: 0.5
código: [2, -4, 1, 6, -1, 4, -7, 3]
vuelta: tensor([ 0.14, -0.29,  0.07,  0.43, -0.07,  0.29, -0.50,  0.21])
error máx: 0.024
con atípico: [0, -1, 0, 7, 0, 0, -1, 0]

absmax, en fórmulas

$w$: el bloque, 8 números · $c_i$: un entero de −7 a 7

volver: un paso vale absmax / 7

redondear erra medio paso, como mucho

Cada peso w sub i del bloque se divide por el máximo valor absoluto, se multiplica por siete y se redondea al entero más cercano: eso es c sub i. En el dibujo, los ocho pesos; las barras rojas son negativas, y la más alta es absmax, cero coma cinco. Para volver, c sub i por absmax sobre siete: cada paso entre códigos vale absmax sobre siete. En el dibujo, los códigos. Redondear se equivoca, como mucho, en medio paso: absmax sobre catorce. Acá, cero coma cero treinta y seis; el peor error real fue cero coma cero veinticuatro.

Un valor atípico arruina su bloque

Los dos casos lado a lado. Con absmax cero coma cinco, los ocho pesos usan casi todos los códigos. Con un atípico de cuatro, el paso pasa a valer cuatro sobre siete, cero coma cincuenta y siete: seis pesos se vuelven cero y dos quedan en menos uno. Por eso la escala no es una para toda la matriz, sino una por bloque, de sesenta y cuatro pesos en QLoRA: un atípico solo arruina a sus vecinos.

Parte 2 · NormalFloat

Los pesos de una red no se reparten parejo: se amontonan cerca de cero. Entonces los dieciséis niveles tampoco deberían ser parejos.

Dieciséis niveles por cuantiles

NF4 viene de NormalFloat de cuatro bits, el formato de QLoRA. El nuestro es una versión de juguete con la misma idea. La distribución normal estándar: media cero, desvío uno. icdf es la inversa de la función de distribución: le das una probabilidad p y te devuelve el valor que deja p de la masa a su izquierda. Ocho cuantiles para los negativos, el último en cero. Nueve para los positivos, el primero en cero. Juntamos, sin repetir el cero: dieciséis. Y normalizamos a menos uno, uno. Los niveles se amontonan cerca de cero: cero coma cero ocho, cero coma dieciséis, y en las puntas, saltos de cero coma veintiocho. int4 tiene quince niveles con paso fijo de cero coma uno cuatro tres.

N = torch.distributions.Normal(0, 1)
d = 0.03                                     # recorte de juguete en las colas
neg = N.icdf(torch.linspace(d, 0.5, 8))      # 8 cuantiles: −… hasta 0
pos = N.icdf(torch.linspace(0.5, 1 - d, 9))  # 9 cuantiles: 0 hasta +…
nf4 = torch.cat([neg, pos[1:]])              # 16 niveles, un solo cero
nf4 = nf4 / nf4.abs().max()                  # normalizados a [−1, 1]

int4 = torch.arange(-7, 8) / 7               # los 15 niveles uniformes
print("NF4 de juguete:", len(nf4), "niveles")
for i in range(0, 16, 4):
    print("  ", " ".join(f"{v:+.2f}" for v in nf4[i:i + 4]))
print("int4:", len(int4), "niveles, paso", round(1 / 7, 3))

NF4 de juguete: 16 niveles
   -1.00 -0.69 -0.52 -0.39
   -0.28 -0.18 -0.09 +0.00
   +0.08 +0.16 +0.24 +0.33
   +0.44 +0.56 +0.72 +1.00
int4: 15 niveles, paso 0.143

Los niveles, en orden

Los niveles ordenados: en el eje horizontal, el número de nivel; en el vertical, su valor. int4 es una recta: cada nivel está a la misma distancia del anterior. NF4 es una ese: en el medio, cerca de cero, la curva es chata y los niveles están juntos; en las puntas, empinada. Donde hay más pesos, hay más niveles, y el error promedio baja.

Cuantiles: un nivel por cada porción de la campana

$\Phi$: la función de distribución normal · $p_i$: probabilidades parejas

un cero exacto, compartido

$c$: un índice de 4 bits, de 0 a 15

Phi mayúscula es la función de distribución de la normal: qué fracción de la campana queda a la izquierda de un valor. Su inversa, Phi a la menos uno, hace lo contrario. Con probabilidades parejas, cada nivel cubre la misma porción de pesos. En el dibujo, los ocho positivos: cada salto es más grande que el anterior. Ocho negativos y nueve positivos comparten el cero: dieciséis niveles, dos a la cuarta, y el cero es exacto. Volver es buscar en una tabla: el código c es un índice, se lee el nivel y se multiplica por absmax. En el dibujo, cada código elige su fila.

Una matriz real de Gemma, en 4 bits

Ahora con pesos de verdad: la q proj de la capa doce de Gemma, leída del archivo. La cargamos. El error relativo: la norma de la diferencia sobre la norma de W. int4 con un solo absmax para toda la matriz. int4 por bloques de sesenta y cuatro. Y nuestro NF4, por bloques. El desvío de los pesos es cero coma cero once, y el máximo, cero coma ciento catorce: diez desvíos. Ese atípico arruina el absmax global: error cuarenta y uno por ciento. Por bloques, once por ciento. Con NF4, nueve coma tres: los niveles en la campana ayudan.

W = cargar("model.layers.12.self_attn.q_proj.weight").float()
print("W:", tuple(W.shape))
print(f"desvío {W.std():.4f}, máx {W.abs().max():.3f}")

def error(W_hat):                            # error relativo de Frobenius
    return ((W - W_hat).norm() / W.norm()).item()

print(f"int4, un absmax para todo: {error(cuantizar(W, INT4, W.numel())):.1%}")
print(f"int4, bloques de 64:       {error(cuantizar(W, INT4, 64)):.1%}")
print(f"NF4,  bloques de 64:       {error(cuantizar(W, NF4, 64)):.1%}")

W: (2048, 2304)
desvío 0.0113, máx 0.114
int4, un absmax para todo: 41.2%
int4, bloques de 64:       11.0%
NF4,  bloques de 64:       9.3%

Error relativo sobre q_proj, capa 12

En barras. Un absmax para todo: cuarenta y uno por ciento. Bloques de sesenta y cuatro: once. La mayor ganancia viene de los bloques. Y NF4 baja a nueve coma tres. Nueve por ciento por matriz es mucho. QLoRA no promete que la base cuantizada sea igual a la original: promete que LoRA, entrenado encima, recupera la calidad.

• int4, global: 41,2 %
• int4, bloques 64: 11,0 %
• NF4, bloques 64: 9,3 %

Parte 3 · Cuánta memoria

Cuatro bits por peso no es todo: cada bloque guarda su absmax. QLoRA comprime también eso.

Gemma 2 2B en cada formato

Bits por parámetro, para los dos mil seiscientos catorce millones de Gemma. Cuenta idealizada: en la práctica, los embeddings suelen quedar en dieciséis bits. float32. bfloat16. Cuatro bits, más un absmax de treinta y dos bits cada sesenta y cuatro pesos. Y la double quantization de QLoRA: los absmax se guardan en ocho bits, con su propia escala de treinta y dos bits cada doscientos cincuenta y seis bloques. Diez coma cuarenta y seis gigabytes, cinco coma veintitrés, uno coma cuarenta y siete, y uno coma treinta y cinco.

GB = 8e9                                      # bits en un gigabyte
bits = {
    "float32": 32,
    "bfloat16": 16,
    "4 + absmax": 4 + 32 / 64,
    "4 + doble": 4 + 8 / 64 + 32 / (64 * 256),
}
for nombre, b in bits.items():
    print(f"{nombre:11} {b:6.3f} bits {total * b / GB:6.2f} GB")

float32     32.000 bits  10.46 GB
bfloat16    16.000 bits   5.23 GB
4 + absmax   4.500 bits   1.47 GB
4 + doble    4.127 bits   1.35 GB

Bits por parámetro

cada bloque de 64 paga un absmax de 32 bits

double quantization

Gemma 2 2B, todo en 4 bits

Cuatro bits del código, más treinta y dos bits de absmax repartidos entre sesenta y cuatro pesos: medio bit cada uno. Cuatro coma cinco en total. Con double quantization, el absmax ocupa ocho bits, un octavo de bit por peso, más una escala de segundo nivel que se reparte entre sesenta y cuatro por doscientos cincuenta y seis pesos. En el dibujo, ese segundo nivel casi no se ve. Cuatro coma ciento veintisiete. Multiplicado por los parámetros y dividido por ocho, de bits a bytes: uno coma treinta y cinco gigabytes, contra diez coma cuarenta y seis en float32.

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

65B en una GPU de 48 GB · QLoRA «reduces memory usage enough to finetune a 65B parameter model on a single 48GB GPU».

Con NF4, double quantization y paged optimizers, que mueven estados del optimizador a la memoria de la CPU cuando la GPU, graphics processing unit, se llena, ajustan un modelo de sesenta y cinco mil millones de parámetros en una sola placa de cuarenta y ocho gigabytes. Y reportan que NF4 con double quantization recupera por completo el rendimiento de LoRA en dieciséis bits, en su benchmark de modelos de siete a sesenta y cinco mil millones.

Parte 4 · LoRA encima

Una capa que guarda la base en 4 bits, la descomprime en cada forward y suma el desvío de LoRA.

Una capa QLoRA, a mano

La capa entera. Partimos los pesos en bloques de sesenta y cuatro y calculamos un absmax por bloque. Para cada peso, el nivel de NF4 más cercano: argmin de la distancia. Guardamos los códigos como buffers, no como parámetros: viajan con el modelo, pero el optimizador no los ve. Usamos un byte por código; un formato real empaqueta dos. A y B, como siempre. En el forward, descomprimimos W y sumamos el desvío. Una capa de doscientos cincuenta y seis por doscientos cincuenta y seis: mil veinticuatro bloques. Entrenan solo A y B. Y la salida difiere de la lineal original, a lo sumo, en cero coma diecinueve.

class Lineal4bit(nn.Module):                  # base en 4 bits + LoRA
    def __init__(self, lin, r=8, alpha=16, bloque=64):
        super().__init__()
        w = lin.weight.detach().reshape(-1, bloque)
        absmax = w.abs().max(dim=1, keepdim=True).values
        codigo = ((w / absmax)[..., None] - NF4).abs().argmin(-1)
        self.register_buffer("codigo", codigo.to(torch.uint8))
        self.register_buffer("absmax", absmax)
        self.register_buffer("bias", lin.bias.detach().clone())
        self.forma = lin.weight.shape
        self.A = nn.Parameter(torch.randn(r, self.forma[1]) * 0.01)
        self.B = nn.Parameter(torch.zeros(self.forma[0], r))
        self.escala = alpha / r

    def forward(self, x):
        W = (NF4[self.codigo.long()] * self.absmax).reshape(self.forma)  # dequant
        lora = x @ self.A.T @ self.B.T            # el desvío, en float32
        return F.linear(x, W, self.bias) + self.escala * lora

El gradiente atraviesa la base de 4 bits

dq: descomprimir · Ŵ: la base de 4 bits

el gradiente sigue hacia las capas de abajo

los códigos no son parámetros

El forward: Ŵ, W con sombrero, es la base descomprimida desde los códigos c y los absmax. Se arma, se usa y se tira. En el dibujo, Ŵ por x, una columna. En el backward, la pérdida L tiene que llegar a las capas de abajo, porque ahí también hay A y B. El gradiente respecto de la entrada x es la matriz transpuesta por delta, el gradiente que llega desde arriba: la base de cuatro bits participa, aunque no se actualice. En el dibujo, Ŵ transpuesta por delta. Lo que no se calcula es el gradiente respecto de Ŵ: los códigos no son parámetros.

LoRA sobre una base de 4 bits

Dígitos escritos a mano de scikit-learn, imágenes de ocho por ocho. La base: una red de tres capas lineales, entrenada para reconocerlos. La tarea nueva: los mismos dígitos, en espejo. La base, en dígitos normales y en espejo. Después, LoRA de rango ocho en las tres capas, sobre la base en float32 y sobre la base en NF4. Medimos cuánto perdió la base al comprimirla. Y entrenamos solo A y B en la tarea del espejo. La base acierta noventa y siete coma siete por ciento, pero en espejo solo cuarenta y cinco coma nueve. En NF4 pierde poco: noventa y siete coma tres. Con LoRA en espejo: noventa y cinco coma cinco sobre float32, noventa y cinco coma ocho sobre NF4. Ocho mil setecientos ochenta y cuatro entrenables.

print("base      dígitos → espejo")
print(f"sin LoRA:  {acierto(base, X):.1%}  →  {acierto(base, espejo):.1%}")

for nombre, Capa in [("float32", LoRALinear), ("NF4", Lineal4bit)]:
    m = con_lora(Capa)                       # A al azar, B = 0
    antes = acierto(m, X)                    # ¿cuánto perdió la base?
    entrenar(m, espejo)                      # solo A y B
    print(f"{nombre + ':':9}  {antes:.1%}  →  {acierto(m, espejo):.1%}")
n = sum(p.numel() for p in m.parameters() if p.requires_grad)
total = sum(p.numel() for p in base.parameters())
print(f"entrenables {n:,} (base: {total:,})")

base      dígitos → espejo
sin LoRA:  97.7%  →  45.9%
float32:   97.7%  →  95.5%
NF4:       97.3%  →  95.8%
entrenables 8,784 (base: 85,002)

Acierto en dígitos en espejo

El resultado. La base sola, en espejo: cuarenta y seis por ciento. LoRA sobre float32: noventa y cinco y medio. LoRA sobre NF4: noventa y cinco coma ocho. La diferencia entre las dos es ruido de este experimento; lo que importa es que la base comprimida no estorbó. Es un juguete: muestra el mecanismo, no reproduce el paper.

• base, sin LoRA: 45,9 %
• LoRA, base float32: 95,5 %
• LoRA, base NF4: 95,8 %

La versión de verdad, con bitsandbytes

Así se hace en una GPU NVIDIA, con la librería bitsandbytes. Este código no corre acá, por eso dice ilustrativo. Cargar en cuatro bits. Con el tipo NF4, el de verdad. Con double quantization. Las cuentas, en bfloat16: es el ŵ descomprimido de nuestra capa. prepare model for kbit training, de peft, prepara el modelo para entrenar encima: por ejemplo, deja en float32 las capas de normalización. Y después, LoRA como en el episodio seis.

cuatro_bits = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",               # NormalFloat de 4 bits
    bnb_4bit_use_double_quant=True,          # cuantiza también los absmax
    bnb_4bit_compute_dtype=torch.bfloat16,   # las cuentas, en 16 bits
)
modelo = AutoModelForCausalLM.from_pretrained(
    "google/gemma-2-2b", quantization_config=cuatro_bits)
modelo = prepare_model_for_kbit_training(modelo)

cfg = LoraConfig(r=8, lora_alpha=16, target_modules="all-linear")
modelo = get_peft_model(modelo, cfg)         # A y B en 16 bits, encima

ImportError: Using `bitsandbytes` 4-bit quantization requires bitsandbytes: `pip install -U bitsandbytes>=0.46.1`

Resumiendo

• absmax: $c = \operatorname{round}(7\, w / \max\lvert w\rvert)$, una escala por bloque
• Bloques de 64: un atípico solo arruina a sus vecinos
• NF4: 16 niveles por cuantiles de la normal, con cero exacto
• Gemma: 10,5 GB en float32 → 1,35 GB en 4 bits
• La base se descomprime en cada forward; A y B entrenan
• Próximo: variantes (DoRA, rsLoRA, LoRA+)

Papers del episodio

• Dettmers et al., 2023 — QLoRA: Efficient Finetuning of Quantized LLMs