La caché de claves y valores: cuánto cuesta cada token

Sesenta y dos tokens tras La Nela, lo que le cabe al mini-GPT en su ventana, le cuestan unas dos mil filas de cálculo. Cada token sale de un sorteo, como estableció la lección sobre el muestreoBloque 1 · Lección 4Muestreo: temperatura, top-k y top-pUna distribución no es un texto: para escribir hay que quedarse con un token en cada posición. La lección deriva la familia entera de reglas a partir de un solo mando, la temperatura, cuyos dos extremos son la favorita de siempre y el sorteo uniforme; demuestra por qué la generación voraz acaba en un ciclo del que no sale y por qué sortear sin más acaba sacando tokens de la cola; y presenta top-k y top-p como dos maneras de cortar esa cola antes de sortear. Todo medido sobre el mini-GPT., y cada sorteo necesita una pasada del texto entero por la red, de la que sólo se aprovecha la última fila. De esas dos mil filas se usan 62; las demás ya se habían calculado antes, idénticas.

Esta lección pone números a ese derroche y a su remedio. ¿Cuánto más caro sale un token si todo se rehace? La respuesta es exacta y no depende del modelo: $t$ veces, en la posición $t$. ¿Y qué cuesta no rehacerlo? Memoria, que crece con cada token y que, en un modelo grande con un texto largo, llega a ocupar más que el propio modelo. Ese número decide cuánto contexto cabe en una máquina.

Cada posición de una capa de atención produce tres vectores, como en el curso anterior: una consulta, con la que busca; una clave, con la que la encuentran las posiciones que vienen detrás; y un valor, lo que les entrega cuando la eligen. La consulta sólo le sirve a su propia posición, y una vez. La clave y el valor los usan todas las posiciones posteriores, y siempre son los mismos: salen de la fila de su posición, y la máscara hace que esa fila no vea nada de lo que hay a su derecha. Un token más al final no les cambia una cifra.

Así que se guardan. En el explorable, el mini-GPT llena su ventana; debajo del texto está lo guardado, una columna por posición y una fila por matriz (claves y valores de sus dos capas), y debajo, las multiplicaciones de cada token. Haz tres cosas. Con caché, avanza token a token o pulsa Animar: cada token calcula su columna y lee las demás, y la barra verde casi no se mueve. Cambia a sin caché: cada token recalcula todas las columnas, y las barras grises forman una rampa. Y lleva $t$ hasta 64: sin caché, ese token cuesta 64 veces más; sumados, los 64 tokens cuestan 33 veces más.

Las claves y los valores de una posición no cambian

Fijemos una capa (uno de los $L = 2$ bloques del mini-GPT, que su código llama capas) y una de sus cabezas. Tras leer $t$ posiciones, la cabeza ha calculado las matrices $\mathbf{Q}$, $\mathbf{K}$ y $\mathbf{V}$ de la lección del curso anterior sobre el producto interno escaladoDeep Learning para NLP: del Perceptrón al TransformerBloque 5 · Lección 3La raíz que faltaba en la atenciónDe la fórmula de la atención queda un solo símbolo sin justificar, y no se justifica midiendo: se cuenta. Una puntuación es una suma de d_k términos, y lo que crece al añadir sumandos no es su tamaño sino cuánto se aleja del cero la suma, que va con la raíz de d_k. Lo que ese tamaño le hace al softmax es lo que el curso ya vio dos veces con otra ropa: la fila se cierra sobre una posición, las dos derivadas del softmax se apagan a la vez y las dos únicas matrices que hay antes de las puntuaciones dejan de recibir gradiente. Dividir entre la raíz devuelve las puntuaciones a una escala que no menciona la anchura del modelo, y con ellas el reparto y el gradiente., de forma $t \times d_k$, con filas $\mathbf{q}_i^{\top}$, $\mathbf{k}_i^{\top}$ y $\mathbf{v}_i^{\top}$. (Esta $\mathbf{q}_t$, con su posición, es la consulta de siempre, no la distribución $\mathbf{q}$ del muestreo.) La máscara pone $-\infty$ en las puntuaciones de la fila $t$ contra las columnas $j > t$, y no hay ninguna: la posición $t$ es la última. La fila $t$ de la salida es, entonces,

sin máscara: la posición nueva mira a todas las anteriores y a sí misma. Pide su consulta y las $t$ filas de $\mathbf{K}$ y de $\mathbf{V}$. La consulta y la última fila de cada matriz salen de la fila $t$ de la entrada de la capa. Las otras $t - 1$ son el asunto.

La lección sobre el modelo de lenguaje causalBloque 1 · Lección 1Una sola columna: el modelo de lenguaje causalEl curso anterior eligió una columna del Transformer, le puso la máscara en todos los bloques y le pidió el token siguiente, posición a posición. Esta lección hace la pregunta que aquél no hizo, qué probabilidad le da ese modelo a un texto entero, y de la respuesta sale todo lo demás: la regla de la cadena convierte esa probabilidad, sin perder nada, en un producto de condicionales; la máscara causal es lo que permite calcular las T de una ventana en una sola pasada; y la pérdida es el logaritmo de ese producto, cambiado de signo y repartido por token. Sin código a propósito: la cuenta se hace a mano sobre una ventana de cuatro palabras, y sale un dieciseisavo por los dos caminos. citó del curso anterior que, con la máscara, la fila $i$ de lo que sale de una capa depende sólo de las filas $1$ a $i$ de lo que entró, y que eso sobrevive a apilar capas. En la entrada de cada capa, por tanto, la fila $i$ es función de $x_{\le i}$ y de nada más, y $\mathbf{k}_i$ y $\mathbf{v}_i$ son esa fila, normalizada y multiplicada por una matriz. Con el argumento diciendo qué texto se le dio a la red, como allí,

La clave y el valor de la posición $i$ se calcularon cuando $i$ era la última, y son los mismos, cifra a cifra, en cualquier texto que sólo crezca por la derecha.

Guardarlos es la caché de claves y valores (key–value cache, o KV cache): $\mathcal{K}^{(l)}$ y $\mathcal{V}^{(l)}$, las $h$ matrices $\mathbf{K}$ y $\mathbf{V}$ de la capa $l$, que tras $t$ posiciones tienen $t \times d_k$ números cada una, $t \times d_{\text{model}}$ entre las $h$. Con ella, el token de la posición $t$ son cuatro cosas:

1. la fila $t$ de la entrada: el embedding del token más el vector de su posición;
2. en cada capa, $\mathbf{q}_t$, $\mathbf{k}_t$ y $\mathbf{v}_t$ a partir de esa fila, y $\mathbf{k}_t^{\top}$, $\mathbf{v}_t^{\top}$ al final de la caché de cada cabeza;
3. la ecuación de arriba, con la caché entera como $\mathbf{K}$ y $\mathbf{V}$, y el resto de la capa (juntar las cabezas, el perceptrón por posiciones) sobre esa sola fila;
4. al salir de la última capa, los logits $\mathbf{z}_t$, con los que se sortea el token $t + 1$.

Lo que sale son los logits de adelante sobre el texto entero. No aproximadamente: los mismos.

El prompt no se sortea. Sus tokens ya existen, y por la misma máscara sus filas se calculan todas de una pasada, como en el entrenamiento, dejando la caché con una fila por token del prompt. Esa primera pasada se llama prefill; después viene un token cada vez.

La caché vale mientras nadie cambie de posición

La fila $i$ lleva dentro su posición. El mini-GPT suma a cada embedding un vector aprendido por posición, de una tabla con $T_{\text{ctx}} = 64$ filas: la opción que la lección sobre la codificación posicionalDeep Learning para NLP: del Perceptrón al TransformerBloque 5 · Lección 5Cada posición, un puñado de relojesLa capa de atención está terminada y sigue leyendo la frase como un montón de vectores sin orden, así que lo que hay que arreglar es la entrada: cada fila tiene que llegar sabiendo dónde estaba. Meter el número de la posición rompe la escala, dividirlo entre la longitud de la frase hace que la misma coordenada signifique cosas distintas en textos distintos, y aprender un vector por posición deja la 7 y la 8 sin relación ninguna. Los senos y cosenos del artículo resuelven las tres: acotados por construcción, con una longitud de onda por pareja de coordenadas, y con una propiedad que se demuestra en tres líneas —moverse k posiciones es girar un ángulo fijo dentro de cada pareja, con la misma matriz en toda la frase, y el producto escalar entre dos posiciones sólo depende de la distancia que las separa. mencionó y apartó, porque la tabla sólo sabe de las posiciones que se entrenaron. Mientras el texto crece por la derecha, cada token conserva la suya y todo lo guardado vale.

En el token 64 se llena la ventana de contexto, lo más que la red puede leer de una vez. (No es la ventana de entrenamiento de la lección sobre el entrenamientoBloque 1 · Lección 3Entrenar un mini-GPT en NumPyEl checkpoint del bloque es el resultado de 2 000 pasos de un bucle que esta lección escribe entero. Cada paso sigue la pérdida media de unas pocas ventanas del corpus sacadas al azar, y la lección deriva por qué eso apunta, de media, adonde apuntaría el corpus entero, y cuánto ruido cuesta ahorrarse el resto. El paso es AdamW, con una tasa de aprendizaje que se calienta y luego decae. Con todo ello el mini-GPT sigue entrenando en el navegador, sobre las últimas páginas de Marianela, que no había visto., el trozo de 65 tokens del corpus del que la red aprendía.) Para escribir el 65, generar suelta el primer token y cada uno baja una posición. Su fila cambia, con ella cambian sus claves y sus valores en todas las capas, y la caché entera deja de valer. Con la caché, $T$, las posiciones que el texto tiene, no pasa de $T_{\text{ctx}}$. Qué hacer con un texto que no cabe es un problema que el curso retoma más adelante.

Lo que cuesta un token, con caché y sin ella

Contemos multiplicaciones, como el curso anterior contó el precio de la atención en su lección sobre el adiós a la recurrenciaDeep Learning para NLP: del Perceptrón al TransformerBloque 5 · Lección 1Adiós a la recurrenciaLa atención ya no necesita a la recurrencia, pero la máquina que la rodea sigue siendo recurrente, y eso fija dos cosas que ningún ordenador arregla: una frase de T tokens obliga a T pasos que no se pueden solapar, y lo que la posición 1 le aporta a la posición T pasa por T menos 1 productos por la misma matriz. Quitar la recurrencia deja las dos en un número que no crece con la frase, a cambio de una rejilla de T por T puntuaciones que sale más barata que la recurrencia mientras la frase mida menos que el modelo. Y a cambio de otra cosa que no se paga en multiplicaciones: el orden.: multiplicar una matriz por un vector cuesta una multiplicación por casilla de la matriz. Lo demás (normalizar, el softmax, la ReLU) cuesta poco por número y queda fuera.

Una fila que atraviesa la red pasa, en cada capa, por las tres proyecciones de consulta, clave y valor, por la matriz que junta las cabezas y por las dos del perceptrón por posiciones; al final, por la capa de salida. Llamemos $c_{\text{fila}}$ a lo que cuesta:

Es una multiplicación por cada peso de una matriz. En el mini-GPT, $2\left(4 \cdot 64^{2} + 2 \cdot 64 \cdot 256\right) + 64 \cdot 512 = 131\,072$: todos sus $136\,448$ parámetros menos los que nunca multiplican una fila (la tabla de posiciones, que se consulta, los sesgos y los de las normalizaciones).

A eso se suma la atención. La consulta de una posición contra $t$ claves son $t$ productos escalares de $d_k$ términos, y mezclar $t$ valores, otras $t \cdot d_k$ multiplicaciones; entre las $h$ cabezas, $2t \cdot d_{\text{model}}$ por capa. Con la caché, el token de la posición $t$ es una fila por la red más esa atención en cada capa:

Sin caché es adelante sobre las $t$ filas: $t$ filas por las matrices y, en cada cabeza, la rejilla entera de $t \times t$ puntuaciones y su mezcla, que el código calcula completa y enmascara después,

exactamente $t$ veces más, en cada posición y para cualquier modelo. Las dos partes se multiplican por $t$: las matrices, porque se rehacen $t$ filas; la atención, porque se rehacen $t$ filas de puntuaciones.

Sumemos sobre los $T$ tokens de un texto, con $\sum_t t = T(T+1)/2$ y $\sum_t t^{2} = T(T+1)(2T+1)/6$:

La atención crece como $T^{2}$ con la caché y como $T^{3}$ sin ella, y las matrices como $T$ y como $T^{2}$: lo que los artículos escriben $O(T^{2}d)$ frente a $O(T^{3}d)$. Con la ventana del mini-GPT llena, $T = 64$, son $8\,921\,088$ multiplicaciones frente a $295\,526\,400$, 33 veces menos.

Qué parte pesa más depende del modelo. La atención de un token iguala a sus matrices cuando $2Lt \cdot d_{\text{model}} = c_{\text{fila}}$, que en el mini-GPT es $t = 512$, ocho veces su $T_{\text{ctx}}$. Por eso, con la caché, su coste por token es casi plano ($c(64)$ es un 12 % más que $c(1)$). En un modelo grande, con textos de decenas de miles de tokens, la atención sí manda.

Lo que se paga a cambio: memoria

La caché guarda, por cada posición y cada capa, una clave y un valor de $d_{\text{model}}$ números entre todas las cabezas. Un texto de $T$ tokens ocupa

El mini-GPT con la ventana llena guarda $2 \cdot 2 \cdot 64 \cdot 64 = 16\,384$, 128 KiB a 8 bytes por número. Un modelo abierto de unos 7 000 millones de parámetros, con $L = 32$ y $d_{\text{model}} = 4\,096$, guarda $262\,144$ números por token: con 4 096 tokens a 2 bytes por número, 2 GiB, y con 32 768, 16 GiB, más que sus propios pesos, unos 13 GiB. Y eso, por cada conversación abierta a la vez. Muchos modelos actuales la reducen haciendo que varias cabezas compartan claves y valores, una variante que queda fuera del curso.

La caché, escrita sobre el mini-GPT

Las dos celdas corren sobre el mini-GPT, el checkpoint del bloque, sin entrenar nada. La primera escribe dos funciones. llenar(ids) es el prefill: pasa el prompt por adelante una vez y se queda con las claves y los valores que la ida ya calcula y guarda para la vuelta del entrenamiento. avanzar(x, cache) es la lista de cuatro cosas de arriba para un token. Con las dos llena la ventana tras La Nela, sorteando con top-p, y después compara los logits de ocho posiciones con los que da adelante sobre el texto entero.

El texto es el del explorable, y los logits de las dos maneras difieren en $4 \times 10^{-15}$: son las mismas sumas hechas en otro orden, y ningún sorteo cambia por ese redondeo. La caché de la primera capa acaba con forma (4, 64, 16), $h \times T \times d_k$, y entre las dos capas guarda los $16\,384$ números de la cuenta de arriba.

Ahora cambia el 64 del while por 65. avanzar falla en la posición 65 con un IndexError: la tabla de posiciones, P, tiene 64 filas y no hay una más. Es la sección sobre las posiciones, en forma de error.

La segunda pone la cuenta al lado del reloj en cinco posiciones: las multiplicaciones de $c(t)$ y de $t\,c(t)$, y lo que tarda el token de las dos maneras. La última línea mira la posición 64 fila a fila.

Tus milisegundos serán otros; la forma, no. Sin caché, el tiempo sube en rampa, de unos 12 ms en la posición 8 a unos 65 en la 64, como $t\,c(t)$. Con ella se queda entre 5 y 9 ms en cualquier posición, como $c(t)$, que en ese tramo sólo crece un 12 %. Pero en la posición 64 la cuenta dice 64 veces y el reloj, unas diez.

Lo que falta es un coste que la cuenta no ve. Cada pasada por la red lanza decenas de operaciones de NumPy, y cada una paga una parte fija, sea del tamaño que sea. Una fila sola la paga entera; 64 filas en una pasada la reparten, y por eso salen a 1 ms por fila y la fila sola a unos 6. No es cosa del navegador. En una tarjeta gráfica (graphics processing unit, GPU), la parte fija es sobre todo traer de la memoria todos los pesos del modelo, que sirven igual para una fila que para cien. Por eso un modelo lee el prompt, cientos de filas por pasada, mucho más deprisa de lo que escribe la respuesta, que va de fila en fila: cada token necesita el anterior ya sorteado.

Comprueba tu intuición

Cinco preguntas y un reto: la atención de un token contra la caché, fuera de la red.

El mini-GPT escribe ahora a una fila por token, y esta lección dice cuánto cuesta cada una. Lo que ninguna de sus cuentas dice es si lo que escribe vale algo. En la primera celda, estrelas, que no existe, y personas, que sí, salen igual de baratos, y la caché hace que los dos lleguen antes.

Hasta aquí, los veredictos sobre el mini-GPT han sido dos: una pérdida de $3.2$ por token en páginas de Marianela que no vio, y un recuento de palabras que existen en la novela. ¿Es mucho $3.2$? ¿Frente a qué, si el mismo modelo con otro tokenizador daría otro número? Convertir esa pérdida en una medida que se pueda comparar, y ponerle delante el modelo más pobre posible como vara de medir, es la lección siguiente, sobre la perplejidad.