Un bloque de transformer son seis tensores y un bus: léelo como un pipeline

La parte 4 se quedó en una sola etapa, pero un decodificador es una pila de bloques (doce en GPT-2 small) y la atención es una de las dos etapas de cada uno. Aquí leo un bloque entero como se lee un pipeline: qué entra, qué escribe cada etapa y adónde van a parar los parámetros y el coste. Ya te adelanto que no van al mismo sitio.

He leído Block en el model.py de nanoGPT con las formas anotadas al margen y le he cargado los pesos de GPT-2. Luego lo he roto a propósito, primero quitando los bloques de uno en uno y después haciendo que cada uno sobrescribiera su entrada sin sumarle nada.

El bloque son cuatro líneas

nanoGPT es un GPT-2 mínimo repartido en dos ficheros. Su README remite ahora a un sucesor, pero model.py sigue siendo la manera más corta de leer un bloque de GPT-2 sin trampas. Block.forward está en model.py:103-106.

def forward(self, x):                   # x: [B, T, 768], the stream
    x = x + self.attn(self.ln_1(x))     # read a normalised copy, add attention's output
    x = x + self.mlp(self.ln_2(x))      # read again, add the feed-forward's output
    return x                            # [B, T, 768]: the shape it came in with

Dentro de esas dos llamadas, en model.py:52-92, hay seis tensores que merece la pena nombrar. El código de abajo está condensado, sin escalado, máscara ni dropout, y las formas del margen salen de ejecutar un modelo con las dimensiones de GPT-2 small, d = 768 y 12 cabezas.

qkv = self.attn.c_attn(self.ln_1(x))         # [B, T, 2304]    q, k, v side by side
q, k, v = qkv.split(768, dim=2)              # three views, [B, T, 768] each
q = q.view(B, T, 12, 64).transpose(1, 2)     # [B, 12, T, 64]  12 heads of 64; k, v alike
att = (q @ k.transpose(-2, -1)).softmax(-1)  # [B, 12, T, T]   Part 4's square
y = (att @ v).transpose(1, 2).reshape(B, T, 768)
y = self.attn.c_proj(y)                      # [B, T, 768]     attention out
x = x + y                                    # [B, T, 768]     the stream, attention added
h = F.gelu(self.mlp.c_fc(self.ln_2(x)))      # [B, T, 3072]    the feed-forward's hidden layer
out = self.mlp.c_proj(h)                     # [B, T, 768]     feed-forward out

Contando el propio flujo, el bloque son x, las proyecciones, las puntuaciones, la salida de la atención, la capa oculta y la salida de la red feed-forward. Tres de esos tensores comparten la forma del flujo. Todo lo demás que aparece en el código es una vista, un paso elemento a elemento o la disposición por cabeza [B, 12, T, 64].

Conviene detenerse en cómo se reparten el trabajo las etapas. La atención es la única que mezcla posiciones (la fila t de su salida se construye con las filas 0..t del flujo), mientras que LayerNorm y la red feed-forward trabajan posición a posición.

LayerNorm resta a cada fila de 768 números su media, la divide por su desviación típica y después aplica una escala y un desplazamiento aprendidos. En model.py:27, todo eso cabe en una sola llamada a F.layer_norm. La red feed-forward ensancha cada fila a 3072, aplica GELU (una versión suave de max(0, x)) y la vuelve a estrechar. Si has escrito un map sobre filas, ya has escrito el flujo de control de la red feed-forward.

Elige la capa oculta de la red feed-forward y mira sus parámetros. Luego elige las puntuaciones, que no tienen ninguno. Ahora arrastra T y fíjate en el término de atención, las puntuaciones y la suma ponderada. Los parámetros no se mueven, pero la aritmética sí.

La arquitectura está en el signo más

Si vuelves a leer las dos líneas, verás que ninguna etapa sustituye a x. Cada una lee una copia normalizada y le suma al original lo que ha calculado, así que el flujo sale del bloque con dos retoques encima y el bloque siguiente lee el resultado.

Ese flujo, [B, T, 768] desde el embedding hasta el último bloque, es lo que la literatura llama flujo residual y lo que aquí llamo el bus. Todas las etapas lo leen y escriben en él, pero no es de ninguna. Si alguna vez has escrito total += delta en un bucle, ya conoces el contrato: cada etapa suma al total acumulado sin poder sobrescribirlo, y para deshacer una escritura anterior hace falta que una etapa posterior la sume con el signo cambiado.

Las cifras confirman que el contrato se cumple. He pasado 1024 tokens de una entrada anterior por GPT-2 small y he medido longitudes (la norma de cada vector), tomando la mediana de todas las posiciones. En los bloques 1 a 10, cada etapa suma un vector de entre el 16 % y el 40 % de la longitud del flujo que lee. El flujo sale apuntando casi adonde apuntaba al entrar: la similitud coseno (1 significa la misma dirección) va de 0,93 a 0,97. Por el camino la norma también crece, de 49 tras el bloque 0 a 480 al final. El bus en sí no se normaliza nunca. Cada etapa normaliza su propia copia y ln_f normaliza una vez más antes de que la capa de salida convierta la última fila en logits, las puntuaciones del token siguiente.

Normalizar fuera del bus fue una decisión. Llamemos pre-LN al orden de GPT-2, que normaliza dentro de la rama, x + attn(ln(x)), y deja el bus como está. El transformer original y GPT-1 eran post-LN y normalizaban después de la suma, x = ln(x + attn(x)), con lo que reescalaban el propio bus en cada etapa. El artículo de GPT-2 deja constancia del cambio en una sola frase, la misma que añade ln_f, y cita un precedente, pero no da ninguna razón propia.

Un experimento pequeño muestra una. He entrenado dos modelos de seis bloques con el texto de este blog, carácter a carácter, con la misma configuración y sin warmup (el warmup consiste en arrancar con un tamaño de paso pequeño e ir subiéndolo). La pérdida mide la sorpresa media del modelo ante cada token siguiente. Cuanto más baja, mejor, y 0 significaría acertar siempre y con total seguridad.

Tras 300 pasos, el modelo pre-LN llega a unos 2,4 y queda por debajo del 2,56 de un modelo que solo conoce la letra anterior. El post-LN se atasca en 3,2 (más o menos lo que dan por sí solas las frecuencias de las letras) y, en una semilla que he alargado hasta los 600 pasos, sigue ahí. Con un warmup de 100 pasos se desatasca entre los pasos 200 y 250. Es un experimento de juguete, con dos semillas y pérdidas leídas de los batches de entrenamiento, pero el resultado va en la línea del análisis publicado sobre por qué post-LN necesita warmup.

Vuelvo a GPT-2 small, esta vez para quitarle bloques. Sobre esos 1024 tokens, el modelo intacto da una pérdida de 3,59. Si quitas cualquiera de los bloques 1 a 10 y dejas que el flujo pase de largo, la pérdida queda entre 3,61 y 3,89; si quitas el bloque 0, salta a 9,15.

Sobrescribir es peor que quitar. Si conservas el bloque pero haces que su salida sustituya al flujo sin sumarse a él, la pérdida queda por encima de 9,6 en los bloques 1 a 8, y en los bloques 2 a 7 supera el 10,8 de una predicción uniforme sobre todo el vocabulario. El modelo acaba peor que si adivinara a ciegas.

Cambia a sobrescribir y verás cómo los bloques 2 a 7 cruzan la línea de puntos. Después elige el bloque 0.

El bloque 0 es la excepción, y la norma de los vectores explica por qué. El embedding entra con una norma de 4,6 y la atención del bloque 0 escribe un vector siete veces más largo. Tras el bloque 0, el flujo es sobre todo la salida de ese bloque, así que se comporta más como un segundo embedding que como un retoque. Sobrescribir ahí no cuesta nada (3,59), porque sumar ya equivalía a sustituir.

Dos tercios de los pesos

Los parámetros de un bloque se pueden contar a mano, con d = 768 y la anchura de la red feed-forward f = 4d. Las dos proyecciones de la atención tienen 4d² + 4d (2 362 368 números); las dos de la red feed-forward, 2·d·f + f + d, que da 8d² + 5d (4 722 432), y las dos LayerNorm, 3072. En todos los tamaños de GPT-2 la red feed-forward se lleva dos tercios de cada bloque, porque esa proporción la fija el 4 * config.n_embd de model.py:82, sea cual sea la anchura. Llama 3.1 8B cambia las dos etapas y su red feed-forward tiene el 81 % de un bloque; en la parte 7 veremos por qué.

Los doce bloques suman 85 millones de parámetros, y GPT-2 small tiene 124 439 808. Casi todo el resto es la tabla de tokens, con 38,6 millones de números, el 31 % del modelo. Se usa dos veces, abajo del todo como la consulta de la parte 2 y arriba del todo como la proyección a logits. En model.py:138 queda compartida con la capa de salida (el mismo array, sin copia). En GPT-2 XL esa misma tabla se queda en el 5,2 % del modelo. Los modelos pequeños gastan casi un tercio de sus pesos en el vocabulario, y los grandes apenas lo notan.

El coste está en otra parte

En un programa corriente, el tamaño de una función y su tiempo de ejecución salen del mismo código, pero aquí van por separado: los parámetros están sobre todo en la red feed-forward, y la aritmética crece con T.

Una pasada hacia delante cuesta unas dos operaciones por token, una multiplicación y una suma, por cada número de una matriz de pesos (la parte 4 las contaba como FLOPs). Si llamamos N a la cantidad de esos números, los parámetros cuestan 2N por token. A eso hay que sumarle el término de atención, 4·L·T·d (donde L es el número de bloques), que cubre las puntuaciones y la suma ponderada sobre el cuadrado completo. En GPT-2 small con 1024 tokens salen 247 millones de operaciones por token de los parámetros y 38 millones de la atención. El FlopCounterMode de torch coincide cuando la atención va escrita en sus cuatro líneas, por la ruta eager de la parte 4, y da 284,8 millones por token, justo la cifra de la fórmula (285,1 si N es el get_num_params() de nanoGPT, que cuenta también los sesgos y las LayerNorm).

Con sdpa, la opción por defecto, el mismo contador daba 247 millones tanto con 256 tokens como con 1024. En torch 2.14, la versión que he usado, el contador no tiene entrada para el kernel de atención fusionado de CPU, la única rutina a la que llama sdpa para sustituir esas cuatro líneas. Por eso el término de atención no llega nunca al total, y nada te avisa. La rama principal de PyTorch ya registra ese kernel (PR #195801), de modo que la próxima versión menor también debería contar el término de atención con sdpa.

Dentro de un bloque, el término de atención es el 18 % de la aritmética con 1024 tokens y adelanta a las capas lineales a los 4608 tokens, seis veces d. Los tiempos lo confirman. En la CPU de mi portátil, la etapa de atención de un bloque, proyecciones incluidas, se lleva el 38 % del tiempo con 256 tokens, el 47 % con 1024 y el 57 % con 4096, aunque tiene un tercio de los parámetros del bloque.

Para ser justos con el reparto, en las longitudes para las que se diseñó GPT-2 son las capas lineales las que hacen la mayor parte de la aritmética, y dentro de ellas la red feed-forward se lleva más de la mitad. El término de atención solo toma el relevo pasada una longitud que GPT-2 no admite. Aun así, al elegir un modelo eliges dos presupuestos. Los parámetros fijan la memoria y no cambian desde que lo cargas; el contexto, que llega con cada petición, marca el coste y una parte creciente de la memoria.

Qué vigila el ingeniero de ML

Son cuatro hábitos, y los dos primeros son de aritmética.

Cuenta los parámetros por componente, sin quedarte en el total. El primer gráfico hace esa cuenta para un bloque, con 4d² para la atención y 2·d·f para la red feed-forward. Para la tabla, suma una sola vez el producto del vocabulario por d. Si comparas un modelo de 124 millones de parámetros con uno de 1500 millones, pregúntate cuánto de cada uno es vocabulario.

Estima el coste en 2N por token y en 6N para entrenar. Entrenar es una pasada hacia delante y otra hacia atrás, y la de atrás, que calcula los gradientes, cuesta lo que unas dos hacia delante. Eso da 6N por token y 6·N·D para una ejecución de D tokens. La configuración de nanoGPT para GPT-2, config/train_gpt2.py, procesa 491 520 tokens por iteración durante 600 000 iteraciones, 295 000 millones de tokens en total. Con eso, 6·N·D da 2,2 × 10²⁰ operaciones, que son 195 horas de una A100 a su pico en bf16, o un día de un nodo de ocho GPU. El fichero dice que cuentes con unos cinco días y el README habla de cuatro, lo que deja la ejecución entre una quinta y una cuarta parte del pico. El estimador del propio nanoGPT añade el término de atención (model.py:296), que con 1024 tokens supone otro 15 %.

La profundidad es latencia. Reparte de cuatro maneras los mismos 85 millones de parámetros de los bloques, en 6 bloques de anchura 1088, 12 de 768, 24 de 544 o 48 de 384. Con batch 1, un token tarda 3,2 ms, 4,0 ms, 4,7 ms y 6,7 ms en atravesar la pila en mi CPU. Cada bloque añade un sobrecoste fijo y los bloques se ejecutan uno detrás de otro. Si sirves las peticiones de una en una, lo que esperas lo marca la profundidad, aunque la aritmética sea la misma.

Mide un bloque antes de quitarlo. La poda de capas (desplegar el modelo con menos bloques) y la salida temprana (pararse en un bloque intermedio cuando la respuesta parece decidida) se apoyan en el bus, y el bus no es uniforme. Quitar un bloque intermedio de GPT-2 small sube la pérdida entre un 0,5 % y un 8 %, pero quitar el bloque 0 la sube un 155 %, y quitar el último, un 41 %. Las cifras de arriba me han llevado un minuto en un portátil, con texto de este blog; las tuyas, sácalas de tu propio tráfico.

Esto tampoco es trabajo de modelado. Un bloque son cuatro líneas alrededor de una suma, y es la suma lo que permite que una pila de doce se comporte como un solo pipeline. Dónde están los pesos y adónde se va el tiempo son preguntas con respuestas distintas, y las dos se resuelven con cuentas que puedes hacer antes de cargar nada.